# An Improved Empirical Fisher Approximation for Natural Gradient Descent

Xiaodong Wu1∗ Wenyi Yu2∗ Chao Zhang2 Philip Woodland1
1University of Cambridge 2Tsinghua University
<EMAIL_ADDRESS>{ywy22<EMAIL_ADDRESS>
###### Abstract

Approximate Natural Gradient Descent (NGD) methods are an important family of
optimisers for deep learning models, which use approximate Fisher information
matrices to pre-condition gradients during training. The empirical Fisher (EF)
method approximates the Fisher information matrix empirically by reusing the
per-sample gradients collected during back-propagation. Despite its ease of
implementation, the EF approximation has its theoretical and practical
limitations. This paper first investigates the inversely-scaled projection
issue of EF, which is shown to be a major cause of the poor empirical
approximation quality. An improved empirical Fisher (iEF) method, motivated as
a generalised NGD method from a loss reduction perspective, is proposed to
address this issue, meanwhile retaining the practical convenience of EF. The
exact iEF and EF methods are experimentally evaluated using practical deep
learning setups, including widely-used setups for parameter-efficient fine-
tuning of pre-trained models (T5-base with LoRA and Prompt-Tuning on GLUE
tasks, and ViT with LoRA for CIFAR100). Optimisation experiments show that
applying exact iEF as an optimiser provides strong convergence and
generalisation. It achieves the best test performance and the lowest training
loss for majority of the tasks, even when compared with well-tuned
AdamW/Adafactor baselines. Additionally, under a novel empirical evaluation
framework, the proposed iEF method shows consistently better approximation
quality to the exact Natural Gradient updates than both EF and the more
expensive sampled Fisher (SF). Further investigation also shows that the
superior approximation quality of iEF is robust to damping across tasks and
training stages. Improving existing approximate NGD optimisers with iEF is
expected to lead to better convergence ability and stronger robustness to
choice of damping.

##  1 Introduction

Parameter optimisation is a crucial research area in the field of deep
learning, where stochastic optimisers are commonly used which update the model
parameters iteratively to minimise a target loss function. Approximate Natural
Gradient Descent (NGD) [1] methods are an important family of approximate
second-order optimisers, which pre-condition the gradient with the
(approximate) Fisher information matrix (also called the Fisher matrix) to
accelerate training or improve generalisation.

Although there are many successful optimisers based on approximate NGD, many
of them in fact use the empirical Fisher (EF) as a pre-conditioner. These
methods are referred to as approximate empirical NGD methods [47]. The EF
method constructs an approximation to the exact Fisher matrix directly from
the gradients of training samples, which are usually readily computed during
the training process [18]. In contrast, the exact Fisher matrix needs to be
either sampled from the model output distribution [24], or requires repeated
evaluation of the matrix-vector product with the Fisher matrix [23], which are
both expensive operations. As a result, due to the ease of implementation
brought by EF, empirical NGD is used in many approximate NGD optimisers as the
default choice [29, 49, 39, 10, 11, 47, 46].

Despite the prevalence of empirical NGD optimisers, it is known that EF is in
general a questionable approximation to the exact Fisher matrix [24, 18, 42].
The poor approximation quality of EF-based updates has been experimentally
verified for small-scale experimental setups by [18, 42]. However, traditional
evaluation methods are used in [18, 42] where the exact NG update and Fisher
matrix need to be explicitly computed, making their findings impossible to
verify for large deep learning setups for practical tasks. Hence, a more
generally applicable evaluation framework is needed. There is also a need for
an improved approximation to the exact Fisher matrix (as a pre-conditioner)
than EF, while being as efficient to implement. This paper aims to fill these
gaps.

Our Contributions:  In this paper, an improved EF (iEF) approximation for NGD
is proposed, which provides a better approximation to the exact Natural
Gradient (NG) updates, while maintaining the practical convenience of EF. This
method allows for a straightforward upgrade for all existing approximate
empirical NGD optimisers. To achieve this, a theoretical investigation into
the behaviour of EF update is first carried out, where the impact of the EF
update on each involved sample is analysed (the “involved samples” refers to
training samples whose gradient is used to construct the EF update). This
leads to finding the inversely-scaled projection issue of EF (see Sec. 4).
Accordingly, the iEF method is proposed to overcome this issue by introducing
a diagonal scaling matrix to the standard formulation of the EF pre-
conditioner. It is motivated as an approximate Gauss-Newton algorithm from a
loss reduction perspective, and is shown to guarantee global convergence under
mild assumptions (see Sec. 5). A novel empirical evaluation framework for
approximate NGD methods is then proposed to enable accurate comparison of
approximate Fisher pre-conditioners (e.g. EF and iEF) in large-scale
optimisation setups (see Sec. 6). We then conducted experiments that compare
the exact EF and iEF methods in a range of practical deep learning setups
including computer vision and fine-tuning large language models. Under our
evaluation framework, iEF demonstrates better approximation quality to exact
NG updates than EF and even sampled Fisher (SF, see Appendix E), and its
approximation quality is significantly more robust to damping than EF and SF
across tasks and training stages. Direct application of iEF as optimiser also
shows consistently strong generalisation and convergence, even when compared
to well-tuned AdamW/Adafactor baselines (see Sec. 7).

##  2 Related Work

Approximate (Empirical) NGD: There are many existing approximate (empirical)
NGD methods, most of which use EF despite its theoretical limitations. Some
prior work, e.g. [39, 36], uses the Woodbury identities [32] to exactly
compute EF updates. Recent block-diagonal methods (based on K-FAC [25]) has
gained popularity due to their efficiency, which includes work that modify the
K-FAC approximation [49, 37, 10, 4, 47, 3] or distributively apply K-FAC as
optimisers [30, 44]. Sometimes Adagrad-based methods [9, 12, 15] are also
regarded as empirical NGD methods. However, the connection is questionable
[18] as these methods use the square-root of the EF matrix, instead of the EF
matrix itself, as a pre-conditioner.

Limitations of EF Approximation: The limitations of EF as an approximate
Fisher matrix have been discussed and demonstrated in several papers [24, 18,
42], among which [18] provided a thorough review and analysis. However, as far
as we are aware, there has been no prior work that analysed the exact EF
method in larger deep-learning setups, and these observations are limited to
small-scale problems for theoretical machine-learning studies. It is known,
however, that practical EF-based optimisers usually require a large damping
factor or a sophisticated damping scheme to work well [29, 27]. It has even
been suggested that an infinitely large damping should be used with the
gradient covariance term [4, 31]. These observations can be tied to the
theoretical limitations of EF.

Empirical Evaluation of Approximate NGD Quality: An accurate evaluation of the
approximation quality to exact NG updates is of great importance for
approximate NGD methods. Usually, the performance of the method of interest is
evaluated on machine learning benchmarks [49, 25, 11], which provide crucial
information from the optimisation perspective. However, limited information
about the approximation quality to exact NGD can be drawn from these
experiments. Therefore, additional small-scale experiments are usually
performed to compare against the exact Fisher matrices [42, 25, 11], or the
exact NG updates [18, 37, 3], which are extremely difficult to do for
commonplace large-scale models. This limits our understanding of these methods
in the context of large-scale tasks.

##  3 Preliminaries

##### Supervised Learning for Classification Model with Softmax Activation:

This paper considers supervised learning of categorical classification, where
a probabilistic model is trained to predict outputs $y\in\\{c|c=1,2,\dots
C\\}$ of $C$ categories from inputs $\bm{x}\in\mathbb{X}$. The target model
$\bm{z}=f_{\bm{\theta}}(\bm{x})$ has $\bm{\theta}\in\mathbb{R}^{P}$ as the
model parameters, which outputs the logits $\bm{z}\in\mathbb{R}^{C}$. Assume a
softmax activation is used on the logits, the model can be expressed as a
conditional probability of $p_{\bm{\theta}}(y|\bm{x})$. Given $N$ i.i.d.
training samples $(\bm{x}_{n},y_{n})_{n=1}^{N}$ (assuming $N\ll P$), the
following accumulated loss is minimised

$\mathcal{L}(\bm{\theta})=\sum\nolimits_{n}-\log
p_{\bm{\theta}}(y=y_{n}|\bm{x}_{n})=\sum\nolimits_{n}l_{n},$ (1)

where $l_{n}=-\log p_{\bm{\theta}}(y=y_{n}|\bm{x}_{n})$ is the categorical
cross-entropy loss for the $n$-th training sample. For brevity, we denote
$p_{\bm{\theta}}(y=c|\bm{x}_{n})=p_{n}(c)$.

A vectorised representation of loss $\bm{l}\in\mathbb{R}^{N}$ is used where
$\bm{l}=\begin{bmatrix}l_{1},l_{2},\cdots,l_{N}\end{bmatrix}^{\top}$. The
accumulated loss then becomes
$\mathcal{L}(\bm{\theta})=\sum\nolimits_{n}l_{n}=\bm{l}^{\top}\bm{1}$ where
$\bm{1}$ is an all 1 vector of matching dimension, and the accumulated
gradient can be re-written as
$\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})=\nabla_{\bm{\theta}}\bm{l}^{\top}\bm{1}$
where $\nabla_{\bm{\theta}}\bm{l}\in\mathbb{R}^{N\times P}$ is the Jaccobian
of per-sample losses w.r.t. model parameters.

##### NGD and Empirical NGD

In a first-order optimisation method, say SGD [38], the updating direction on
the model parameter is an estimate of the accumulated gradient
$\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})$. In an NGD method [24], the
gradient is pre-conditioned by the Fisher matrix $\mathbf{F}$ (i.e.
$\mathbf{F}^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})$) to accelerate
convergence. The exact Fisher matrix can be computed from the model output
distribution using available training samples as follows

$\mathbf{F}\mathrel{\mathop{\ordinarycolon}}=\sum\nolimits_{n}\sum\nolimits_{c}p_{n}(c)\left[\nabla_{\bm{\theta}}\log
p_{n}(c)\nabla_{\bm{\theta}}\log p_{n}(c)^{\top}\right],$ (2)

which can be estimated with Monte-Carlo sampling [25] (see Appendix E), which
is termed SF in this paper. Alternatively, when the model is well trained and
$p_{n}(y_{n})\to 1$ for all $N$ samples, it is possible to approximate the
exact Fisher with EF using the empirical gradient as follows

$\tilde{\mathbf{F}}\mathrel{\mathop{\ordinarycolon}}=\sum\nolimits_{n}\left[\nabla_{\bm{\theta}}\log
p_{n}(y_{n})\nabla_{\bm{\theta}}\log
p_{n}(y_{n})^{\top}\right]=\nabla_{\bm{\theta}}\bm{l}^{\top}\nabla_{\bm{\theta}}\bm{l}.$
(3)

Pre-conditioning the gradient with the EF matrix (i.e.
$\tilde{\mathbf{F}}^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})$) yields
the empirical NGD method. Although NGD is prevalent due to the efficiency of
computing the EF matrix in practice, the approximation quality of EF to the
exact Fisher matrix is still worth questioning [18, 42].

##  4 Inversely-Scaled Projection Issue of Empirical Fisher

Despite the practical prevalence of the EF method, it is generally believed to
be a poor approximation of the exact NGD method [18]. To better understand the
cause of the limited approximation quality of the EF method, an analysis of
the impact of the EF update on each of the involved samples is presented
below. This leads to finding the “inversely-scaled projection issue” of the EF
method, which provides a focus for the improvement of the EF method.

###  4.1 Formal Definition

Recall the definition of EF in Eqn. (3). The empirical NG update (or just the
EF update) can be defined as follows

$\Delta\bm{\theta}_{\text{EF}}=-\eta\,\tilde{\mathbf{F}}^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})=-\eta\,(\nabla_{\bm{\theta}}\bm{l}^{\top}\nabla_{\bm{\theta}}\bm{l}+\lambda\mathbf{I})^{-1}(\nabla_{\bm{\theta}}\bm{l}^{\top}\bm{1})$
(4)

where $\lambda\in\mathbb{R}^{+}$ is a small damping factor to facilitate
inversion (gradient covariance matrix
$\nabla_{\bm{\theta}}\bm{l}^{\top}\nabla_{\bm{\theta}}\bm{l}\in\mathbb{R}^{P\times
P}$ cannot be directly inverted for over-parameterised models). Using the
Woodbury identity [32], the EF update can be re-expressed as follows:

$\Delta\bm{\theta}_{\text{EF}}=-\eta\,\nabla_{\bm{\theta}}\bm{l}^{\top}(\nabla_{\bm{\theta}}\bm{l}\nabla_{\bm{\theta}}\bm{l}^{\top}+\lambda\mathbf{I})^{-1}\bm{1}.$
(5)

The loss induced on each sample (denoted as $\Delta\bm{l}_{\text{EF}}$) when
applying the EF update to the model can be estimated using the Jacobian
$\nabla_{\bm{\theta}}\bm{l}$ as follows:

$\Delta\bm{l}_{\text{EF}}=-\nabla_{\bm{\theta}}\bm{l}\Delta\bm{\theta}_{\text{EF}}=-\eta\,\nabla_{\bm{\theta}}\bm{l}\nabla_{\bm{\theta}}\bm{l}^{\top}(\nabla_{\bm{\theta}}\bm{l}\nabla_{\bm{\theta}}\bm{l}^{\top}+\lambda\mathbf{I})^{-1}\bm{1}\approx-\eta\bm{1},$

which is an exact estimate as long as the Jacobian
$\nabla_{\bm{\theta}}\bm{l}$ has full row rank, and both $\lambda$ and $\eta$
are infinitesimally small. This result means that EF updates have the property
of inducing an equal loss reduction on every involved sample. For the $n$-th
sample, the projection of the EF update onto gradient direction
$\nabla_{\bm{\theta}}l_{n}$ (denoted as $(\kappa_{n})_{\text{EF}}$) can be
computed as follows

$(\kappa_{n})_{\text{EF}}=\Delta\bm{\theta}_{\text{EF}}^{\top}\frac{\nabla_{\bm{\theta}}l_{n}}{\|\nabla_{\bm{\theta}}l_{n}\|_{2}}=-\frac{\eta}{\|\nabla_{\bm{\theta}}l_{n}\|_{2}},$
(6)

where $\|\nabla_{\bm{\theta}}l_{n}\|_{2}$ denotes the $l_{2}$ norm of the
$n$-th per-sample gradient. This means that the EF update projection onto
every sample gradient is inversely proportional to each sample gradient norm.
Note that a smaller $\|\nabla_{\bm{\theta}}l_{n}\|_{2}$ generally indicates
the sample is better trained (or more converged, or closer to its minimum).
The EF update is therefore easily biased towards well-trained samples, and
tends to have a larger norm as training progresses
($\|\nabla_{\bm{\theta}}l_{n}\|_{2}$ decreases) [18]. We term this the
inversely-scaled projection issue of the EF update, which is further
illustrated in the following section.

###  4.2 Visual Illustration

The detrimental impact of the inversely-scaled projection issue is illustrated
in Fig. 1 (third plot), where EF updates are “attracted” to the minimum of
each training sample (the dashed lines), leading to distorted update vector
field and inefficient training trajectories. Meanwhile, it can also be seen
that EF updates have a much larger norm when either training sample is nearly
converged, suggesting the necessity of a complicated step-size scheduler.
Please refer to Appendix B for a detailed description and discussion. Note
that the effect of inversely-scaled projection issue is also validated by
experiments (E1) and (E2) of large-scale deep learning setups in Sec. 7.

Figure 1: A visual comparison of Fisher, iEF and EF as pre-conditioners for a
2D linear least-squares regression problem (refer to Appendix B for details).
All three plots are loss landscapes with the $x$-axis and $y$-axis
representing $\theta_{0}$ and $\theta_{1}$ respectively. The first plot shows
the gradient vector field of the loss function and 5 sample training
trajectories for SGD updates. Similarly, the second plot is for NGD/iEF
updates and the third plot is for EF updates (with a zoomed view). The global
minimum (0, 0) is marked by a star where visible. The two dashed lines on all
plots represent the optimal parameter sets for each training sample. It can be
seen that the EF method has a highly distorted update vector field while the
iEF and NGD methods adapt to the curvature of the problem successfully.

##  5 Improved Empirical Fisher

The EF method is a widely used approximate NGD method, since it can be
implemented conveniently by constructing the EF matrix with the per-sample
gradients that are readily computed during backpropagation. In this section,
we propose the improved EF (iEF) method which preserves the implementational
convenience of the EF method, meanwhile resolving the inversely-scaled
projection issue. The iEF method is justified as an approximate (generalised)
GDN method from a loss reduction perspective. A continuous-time convergence
analysis shows that the iEF method guarantees sub-linear/linear convergence to
the global minimum under mild assumptions.

###  5.1 Update Formulation

The proposed iEF method resolves the EF issues by rescaling the projection on
each involved sample direction with the square of the model output logits-
level gradient norm $\|\nabla_{\bm{z}_{n}}l_{n}\|^{2}_{2}$. The projection
norm for iEF is constructed as

$(\kappa_{n})_{\text{iEF}}=-\frac{\eta\,\|\nabla_{\bm{z}_{n}}l_{n}\|^{2}_{2}}{\|\nabla_{\bm{\theta}}l_{n}\|_{2}},$
(7)

where the ratio
${\|\nabla_{\bm{z}_{n}}l_{n}\|^{2}_{2}}/{\|\nabla_{\bm{\theta}}l_{n}\|_{2}}$
is now aligned with the convergence level of each sample (closeness to
minimum), thus resolving the inversely-scaled projection issue.

A viable formulation of the iEF update $\Delta\bm{\theta}_{\text{iEF}}$ which
relies only on the per-sample gradients is proposed as

$\Delta\bm{\theta}_{\text{iEF}}=-\eta\,\nabla_{\bm{\theta}}\bm{l}^{\top}(\nabla_{\bm{\theta}}\bm{l}\nabla_{\bm{\theta}}\bm{l}^{\top}+\lambda\mathbf{I})^{-1}\bm{s}_{\text{iEF}}.$
(8)

where $\bm{s}_{\text{iEF}}\in\mathbb{R}^{N}$ is a scaling vector defined as

$\bm{s}_{\text{iEF}}=\begin{bmatrix}\|\nabla_{\bm{z}_{1}}l_{1}\|_{2}^{2}&\|\nabla_{\bm{z}_{2}}l_{2}\|_{2}^{2}&\cdots&\|\nabla_{\bm{z}_{N}}l_{N}\|_{2}^{2}\end{bmatrix},$

and this scaling vector can be obtained along with back-propagation in Pytorch
[28] with negligible overhead. For the toy example in Fig. 1, by switching
from EF (the third plot) to iEF (the second plot), the distortion of the EF
update vector field is completely removed, indicating the effectiveness of
iEF. Further experimental results in Sec. 7 validate that iEF achieves
consistently better approximation quality to NG updates than EF (and SF)
method in practical deep learning setups (experiment (E1)), meanwhile being
robust to the choice of damping $\lambda$ across tasks and training stages
(experiment (E3)).

###  5.2 Theoretical Connection to Generalised NGD

In the iEF method, the choice of scaling vector $\bm{s}_{\text{iEF}}$ is
motivated by the Gauss-Newton (GN) algorithm, which is a type of generalised
NGD method [26]. The update for the GN algorithm is defined as

$\Delta\bm{\theta}_{\text{GN}}=-\eta\,\hat{\mathbf{G}}^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}),$
(9)

where
$\hat{\mathbf{G}}=\sum\nolimits_{n}\nabla_{\bm{\theta}}\bm{z}_{n}^{\top}\nabla_{\bm{\theta}}\bm{z}_{n}$
is the GN matrix. The GN algorithm can be viewed as gradient descent on the
model output logits space ($\bm{z}_{n}$-space) and the loss reduction induced
for the $n$-th sample by the GN update is approximately

$\Delta(l_{n})_{\text{GN}}\approx-\eta\,\|\nabla_{\bm{z}_{n}}l_{n}\|_{2}^{2}.$
(10)

Meanwhile, the iEF update given in Eqn. (8) induces the same loss reduction on
each involved sample

$\Delta(l_{n})_{\text{iEF}}=\nabla_{\bm{\theta}}l_{n}^{\top}\Delta\bm{\theta}_{\text{iEF}}\approx-\eta\,\|\nabla_{\bm{z}_{n}}l_{n}\|^{2}_{2}.$
(11)

Therefore, the iEF method can be regarded as an efficient approximation to the
GN algorithm in terms of its loss-reduction behaviour. In particular, the iEF
method is equivalent to the GN algorithm and/or the exact NGD method for some
important machine learning setups (see Appendix A).

###  5.3 Convergence Analysis

In this section, two continuous time convergence analyses are provided for the
non-stochastic version of the iEF method, which shows its sub-linear or linear
global convergence guarantee for different types of objective functions (see
Appendix C for proofs). The analysis can be considered an extension to proofs
provided in [48] for non-regression models and cross-entropy objective
function. The two base assumptions used by the two convergence analysis are as
follows:

######  Assumption 5.1.

At time $t$, the full-batch, un-damped iEF update to model parameters
$\bm{\theta}(t)$ is

$\dfrac{\operatorname{d\\!}{}\bm{\theta}(t)}{\operatorname{d\\!}{t}}=-\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}[\nabla_{\bm{\theta}(t)}\bm{l}(t)\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}]^{-1}\bm{s}_{\text{iEF}}(t),$
(12)

######  Assumption 5.2.

$\forall t>0$, the gradient covariance matrix (or Gram matrix)
$[\nabla_{\bm{\theta}(t)}\bm{l}(t)][\nabla_{\bm{\theta}(t)}\bm{l}(t)]^{\top}$
is always full rank.

The two main conclusions of the analysis are described below.

##### Sub-linear Global Convergence for Softmax + Cross-Entropy Objective

When the target model uses softmax output and cross-entropy loss (as described
in Sec. 3), the Theorem 5.3 can be proved.

######  Theorem 5.3.

Suppose Assumption 5.2 holds, $\forall n\in\\{1,\dots,N\\}$, the target
probability
$\hat{p}_{n}(t)\mathrel{\mathop{\ordinarycolon}}=p_{\bm{\theta}(t)}(y=y_{n}|\bm{x}_{n})$
for the $n$-th training sample is bounded as follows

$\hat{p}_{n}(t)>1-\frac{2}{t+C_{0}+1},$ (13)

where
$C_{0}=\frac{1}{1-\hat{p}_{n}(0)}+\log\frac{\hat{p}_{n}(0)}{1-\hat{p}_{n}(0)}$
and $t>max\\{-1-C_{0},0\\}$.

##### Linear Global Convergence for Strongly Convex Objective

When the target model uses an $m$-strongly convex objective function [2] (see
Assumption C.2), the Theorem 5.4 can be proved.

######  Theorem 5.4.

Suppose Assumption 5.2 and C.2 holds, $\forall n\in\\{1,\dots,N\\}$, the per-
sample loss $l_{n}(t)$ for the $n$-th training sample is bounded as follows

$l_{n}(t)-l_{n}^{\star}\leq e^{-2mt}(l_{n}(0)-l_{n}^{\star}),$ (14)

where $l_{n}^{\star}$ is the minimum loss for the $n$-th sample.

Remark: The global convergence guarantee of iEF can be intuitively understood:
the iEF update would become zero if and only if all samples have reached their
respective minimum. This means that iEF method does not get stuck in local
minima, unlike the gradient descent method.

###  5.4 Applications of IEF

There are two main applications of iEF. First, the exact iEF method can be
used directly as an optimiser (see Algorithm 1) for models with a small number
of trainable parameters. Its performance is evaluated in experiment (E2) in
Sec. 7, which demonstrates competitive convergence and generalisation against
well-tuned baselines. Refer to Appendix. D.1 for discussions on the
implementation and complexity.

More importantly, the iEF method provides a straightforward guide to improve
existing approximate empirical NGD optimisers. The evaluation in Sec. 7 shows
that iEF achieves consistently better approximation quality to NGD than both
EF and SF, meanwhile obviating the need for damping tuning. Additionally, the
corresponding iEF matrix $\tilde{\mathbf{F}}^{\star}\in\mathbb{R}^{P\times P}$
is

$\tilde{\mathbf{F}}^{\star}=\nabla_{\bm{\theta}}\bm{l}^{\top}\text{diag}(\bm{s}_{\text{iEF}})^{-1}\nabla_{\bm{\theta}}\bm{l},$
(15)

which has a highly similar form to the EF matrix as proved in Eqn. (41). This
makes integrating iEF into existing EF-based methods particularly
straightforward. Refer to Appendix. D.2 for an example of iEF-based
improvement to the popular empirical K-FAC optimiser.

##  6 Empirical Evaluation Framework for Approximate NGD Methods

Traditional evaluation methods for quality of approximate NGD methods have
high memory and time complexity, which is infeasible for large setups (see
discussion in Sec. 2). In order to accurately evaluate the quality of
approximate NGD methods (EF, iEF, SF etc.) in practical deep-learning setups,
we introduce an efficient empirical evaluation framework which enables a
quantitative comparison of different approximate NGD methods under large-scale
setups. For a given approximate NGD method that generates an update
$\Delta\bm{\theta}$, our proposed evaluation framework satisfies the following
requirements: 1) provides a quantitative evaluator $\gamma(\Delta\bm{\theta})$
that measures the (direction-wise) approximation quality to the exact NG
update; 2) the evaluation process is efficient in modern auto-grad frameworks,
and it poses no constraints on the size or structure of the target model. The
implementation and theoretical motivations of this empirical evaluation
framework are discussed in the following sections.

###  6.1 Efficient Indicator of Approximation Quality

The proposed evaluation framework revolves around the indicator
$\gamma(\Delta\bm{\theta})$ which is designed to accurately reflect the
quality of an approximate NG update, while being efficient to compute. For an
update $\Delta\bm{\theta}$ of interest, the proposed indicator
$\gamma(\Delta\bm{\theta})\in\mathbb{R}^{+}$ is defined as

$\gamma(\Delta\bm{\theta})=\frac{(\Delta\bm{\theta}^{\top}\mathbf{F}\Delta\bm{\theta})^{\frac{1}{2}}}{\Delta\bm{\theta}^{\top}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})},$
(16)

and the smaller the value of $\gamma(\Delta\bm{\theta})$, the better the
approximation quality of $\Delta\bm{\theta}$ to the exact NG update. Note that
this indicator requires only a matrix-vector product with the exact Fisher
matrix (i.e. $\textbf{F}\Delta\bm{\theta}$), which can be efficiently computed
in modern auto-grad frameworks [28]. This allows for the application of this
framework to large-scale models. Refer to Appendix F.1 for implementation
details, algorithm complexity and a comparison with traditional methods.

###  6.2 Theoretical Motivation

In this section, the proposed indicator $\gamma(\cdot)$ is justified as a
theoretically appropriate evaluator of the quality of an approximate NG
update. An alternative definition for the NGD is first proposed, which
formulates the NG update direction with an unconstrained optimisation problem
as

$\zeta^{{}^{\prime}}\mathbf{F}^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})=\operatorname*{arg\,min}_{\Delta\bm{\theta}}\gamma(\Delta\bm{\theta})^{2},$
(17)

where $\zeta^{{}^{\prime}}\in\mathbb{R}$ is an arbitrary non-zero scalar. By
this definition, any approximate NG update with a smaller
$\gamma(\Delta\bm{\theta})$ is a “strictly better approximation” to the exact
NG update (which achieves the minimum for $\gamma(\Delta\bm{\theta})$). This
alternative definition for NGD is implicitly used in the Hessian-free method
[23]. It is shown that the linear conjugate gradient (CG) algorithm used in
Hessian-free to solve for the exact NG update is a locally optimal minimiser
for $\gamma(\Delta\bm{\theta})^{2}$ (see Appendix F.2).

Furthermore, $\gamma(\Delta\bm{\theta})$ can also be justified from a second-
order optimisation perspective. $\frac{1}{2\gamma(\Delta\bm{\theta})^{2}}$ is
shown to quantify the maximum achievable loss reduction for a given update
direction under a local quadratic approximation of the loss function (see
Appendix. F.3). Consequently, the proposed indicator can be used to
(accurately) predict the convergence ability of a target update generation
method (this is demonstrated in experiment (E2) in Sec. 7).

##  7 Experiments

Experimental results are presented in this section. The main goal of the
experiments is to verify that the behaviour of exact EF and iEF methods align
with our theories in practical deep learning setups. Mainly three
approximation methods are compared: EF, iEF and SF (a strong yet more
expensive approximation method, see Appendix E). The exact updates of each
method are generated based on Eqn. (5), (8), (48) respectively. Fifteen
different setups are used to evaluate the optimisation performance and the
approximation quality of these methods, including widely used parameter-
efficient fine-tuning (PEFT) for pre-trained models. These include T5-base
with LoRA and Prompt-Tuning on GLUE tasks [7], and ViT with LoRA for CIFAR100
[13]. PEFT of pre-trained models is investigated because it involves large-
scale practical models, while having a small trainable parameter size (the
implementation of exact EF, iEF and SF methods are memory intensive, see
Appendix D.1). Please refer to Appendix H.1 for detailed experimental setups.
The following three findings are demonstrated with our experiments.

(E1) The approximation quality (to exact NG updates) of EF, iEF, SF and SGD
has been evaluated and compared using the proposed evaluation framework in
Sec. 6 on all setups. It is shown that iEF consistently improves on SGD
updates and is superior to both EF and SF methods for the majority of the
training stages for all setups.

(E2) The optimisation performance of EF, iEF, SF and SGD was evaluated on all
setups. For each task, an appropriate baseline (Adafactor/AdamW) was also
compared. It is shown that iEF achieves comparable performance to the best
baseline method, while iEF and SF suffer from unstable training to different
extents.

(E3) The impact of damping on the approximation quality of EF, iEF and SF is
analysed under the proposed evaluation framework. It is shown that careful
damping tuning is crucial to the quality of traditional EF and SF methods,
while iEF works well with near-zero damping.

##### E1: Approximation Quality to NG Updates

The behaviour of updates generated with EF, iEF, SF and SGD methods were
compared using the proposed empirical evaluation framework in terms of their
approximation quality to exact NG updates. The updates for EF, iEF and SF were
generated according to Eqns. (5), (8), and (48) respectively, and the
evaluation framework follows Algorithm 4. The “un-damped” behaviour of these
methods is analysed and a near-zero damping factor is used for update
generation. The checkpoints at the end of each epoch generated by the baseline
optimisation methods (AdamW/Adafactor) for each task were used for evaluation.
In each evaluation. For each checkpoint $\bm{\theta}(t)$, indicators are
computed from 100 batches of randomly picked training samples of the target
task of batch size $M=160$. The averaged indicator for each update were
evaluated: $\bar{\gamma}(\Delta\bm{\theta}_{\text{SGD}}(t))$,
$\bar{\gamma}(\Delta\bm{\theta}_{\text{EF}}(t))$,
$\bar{\gamma}(\Delta\bm{\theta}_{\text{iEF}}(t))$,
$\bar{\gamma}(\Delta\bm{\theta}_{\text{SF}}(t))$ (which are denoted as
$\gamma_{\text{SGD}},\gamma_{\text{iEF}},\gamma_{\text{EF}},\gamma_{\text{SF}}$
for simplicity). The relationship among these indicators across epochs and
tasks is shown in Fig. 2. Note that results are presented for only 3
representative setups due to space limit (indicator plots for all tasks are
shown in Appendix H.5.1). Three findings can be concluded from these figures.
1)  EF achieves poorer approximation quality even than SGD updates for most
training stages and tasks. This is aligned with the finding in prior work that
EF is a questionable approximation to NGD. 2)  The fourth plot shows that the
gradient norm imbalance gets larger as training progresses. This correlates
well with both the EF and SF curves, while impacting iEF less. This means that
the inversely-scaled projection issue indeed plays a significant role in
reducing the approximation quality of the EF (and SF) approximation. 3)
Comparing the first three plots, it can be seen that, for the majority of the
training stages, the approximation quality follows iEF $>$ SF $>$ EF. IEF
gives a consistently better approximation, and EF and SF are only able to beat
iEF at the start of training (where a good approximation to the NG update has
less impact).

Figure 2: Four (log-scaled) ratios computed for checkpoints at various stages
of training (sampled at the interval of one epoch) for 3 of the all 15 tasks.
For the $x$-axes, $0$ indicates the initialised model and $1$ indicates
checkpoints at the end of the last epoch. Each data point is averaged across
100 evaluations, and the error bars represent the standard deviation
(1-sigma). The first plot shows $\gamma_{\text{EF}}/\gamma_{\text{SGD}}$,
which denotes the relative approximation quality improvement of EF updates
w.r.t. SGD updates (the lower the better). The second plot shows
$\gamma_{\text{iEF}}/\gamma_{\text{SGD}}$, and the third plot shows
$\gamma_{\text{SF}}/\gamma_{\text{EF}}$. The last plot depicts the imbalance
of gradient norms, which is the average ratio between the maximum and minimum
gradient norm for each evaluated batch (a larger value indicates a more
imbalanced gradient norm). Overall, the approximation quality follows iEF $>$
SF $>$ EF.

##### E2: Optimisation Performance

The exact iEF, EF and SF methods were implemented as stochastic optimisers
(following Algorithms 1, 2, 3 respectively). The same near-zero damping factor
was used as in (E1). The averaged test metrics for GLUE and CIFAR100 for each
optimiser are shown in Table 1 (see full test results in Table 7, validation
result in Table 6, final training loss in Table 5 and training curves in Fig.
10 and 11. The following three observations can be made. 1)  From the final
training loss reported in Table 5, the ranking of final training loss
generally follows iEF $<$ AdamW/Adafactor $<$ SGD $<$ SF $<$ EF (the lower the
better). This ranking of training convergence follows the ranking of
indicators in (E1) closely, demonstrating the effectiveness of the empirical
evaluation framework in predicting the training behaviour of optimisers. 2)
Particularly, for most of the tasks, EF always suffer from unstable training
(see training curves in Fig. 10 and 11), while iEF consistently reaches the
lowest training loss at the end of training (even when compared with well-
tuned Adafactor/AdamW baselines). This further confirms the inversely-scaled
projection issue of EF, and demonstrates the strong convergence ability of the
proposed iEF method. 3)  From test results in Table 1, it can be seen that iEF
achieves the best generalisation for Prompt Tuning tasks (outperformed
Adafactor in 6 out of 7 tasks). For LoRA tasks, iEF remains competitive to
AdamW with each of them outperformed the other in 4 out of 8 tasks. This is
likely because LoRA setups (which on average have 50 times more trainable
parameters than Prompt Tuning) have a stronger reliance on regularisation and
momentum, which have not been properly extended to use together with the iEF
optimiser yet. Overall, iEF achieves the best generalisation for the majority
of tasks (10 out of 15), which indicates its potential as a strong optimiser
for PEFT for pre-trained models.

Table 1: Average test performance of different optimisers for GLUE and CIFAR100. For GLUE tasks, the metric results for the 7 tasks are averaged on to produce the final test score. For tasks with two metrics, these metrics are averaged first [45]. For all tasks, the test result is computed for the best validation accuracy checkpoint. Refer to Table 7 for a more complete test performance report and detailed explanations on metrics. | AdamW | Adafactor | SGD | EF | SF | iEF
---|---|---|---|---|---|---
GLUE + T5 + Prompt Tuning | - | $77.1$ | $67.4$ | $48.1$ | 69.7 | $\bm{79.3}$
GLUE + T5 + LoRA | $\bm{80.1}$ | - | $77.3$ | $63.1$ | $76.5$ | $79.3$
CIFAR100 + ViT + LoRA | $93.9$ | - | $91.3$ | $31.0$ | $92.8$ | $\bm{94.3}$

##### E3: Impact of Damping

As is discussed in Sec. 2, practical approximate NGD optimisers rely heavily
on a good damping schedule, which is typically chosen based on empirical
experience [25]. Using the proposed evaluation framework, it is
straightforward to observe the impact of damping on the approximation quality
of EF, SF and iEF. For the CoLA+T5+LoRA task, the indicator $\gamma$ w.r.t.
damping $\lambda$ curve is computed at the start, mid-way and end of the
training, as is shown in Fig. 3 (graphs for other tasks are shown in Appendix
H.5.2). Two findings can be concluded from these graphs. 1)  A well-chosen
damping factor significantly improves the approximation quality of EF and SF,
which is aligned with observations in prior work on approximate NGD optimisers
[25, 29, 30]. However, the optimal damping factor changes greatly for
different tasks and training stages, and this makes the damping schedule for
SF or EF based optimisers hard-to-design yet necessary in practice. 2)  Across
all tasks and training stages, iEF robustly achieves great approximation
quality with a near-zero damping factor. More importantly, its approximation
quality is consistently better than EF method, and is comparable to the
optimally-damped SF method (which is much more expensive, particularly when
the cost of damping tuning is considered). This makes iEF a great potential
improvement method for both EF and SF based approximate NGD optimisers.

Figure 3: Approximation quality (relative to SGD) of EF, SF and iEF methods
w.r.t. damping factor $\lambda$ at different training stages of task
CoLA+T5+LoRA, $x$-axes show the value of the damping factor, $y$-axes depict
the relative approximation quality improvement of the target update method
w.r.t. SGD (the lower the better). Each data point is averaged across 100
evaluations, and the error-bars represent the standard deviation (1-sigma).
The first plot is for checkpoint saved at the end of the first training epoch,
the second plot at the mid-way epoch and the third plot at the end of
training. It can be observed that iEF achieves the best approximation quality
robustly for a near-zero $\lambda$. In contrast, $\lambda$ has a non-linear
impact on both SF and EF. When optimally tuned, an EF update can achieve
better approximation quality than SGD, and an SF update can achieve comparable
approximation quality to iEF. Note that the optimal damping factor for EF and
SF changes greatly with training stages.

##  8 Conclusions

This paper presents the iEF method, which resolves the inversely-scaled
projection issue of the EF approximation for NGD, while maintaining the
implementational convenience. Furthermore, a novel empirical evaluation
framework for the quality of general approximate NGD update is proposed, which
enables quantified comparison of approximate NGD methods in large deep
learning setups111Codebase for both will be publicly released.. Based on the
experiments with practical PEFT of pre-trained models for NLP and CV
classification tasks, the exact iEF optimiser shows superior convergence and
generalisation for majority of the tasks, supporting the applicability of iEF
directly as an optimiser. Further evaluation on approximation quality
concludes that iEF achieves consistently better approximation quality than
both EF and SF. It also demonstrates the superior property of being robust to
choice of damping factor across different tasks and training stages. It is
expected that, by integrating iEF into existing approximate NGD optimisers
(which is straightforward to do), the improved optimiser is likely to obtain
stronger convergence ability and significantly reduced reliance on
sophisticated damping scheduling.

## Acknowledgements

Xiaodong Wu is in part funded by the Cambridge Trust, a donation from Meta
Systems and Christ’s College, Cambridge. This work has in part been performed
using resources provided by the Cambridge Tier-2 system operated by the
University of Cambridge Research Computing Service (www.hpc.cam.ac.uk) funded
by EPSRC Tier 2 capital grant EP/T022159/1.

## References

  * [1] Shun-ichi Amari “Natural Gradient Works Efficiently in Learning” In _Neural Computation_ 10.2, 1998, pp. 251–276 DOI: 10.1162/089976698300017746
  * [2] Andersen Ang “Convergence of gradient flow” Dept. Combinatorics & Optimization, University of Waterloo, Canada, 2022 URL: https://angms.science/doc/CVX/GradientFlowConvAna.pdf
  * [3] Achraf Bahamou, Donald Goldfarb and Yi Ren “A Mini-Block Fisher Method for Deep Neural Networks” In _International Conference on Artificial Intelligence and Statistics_ , 2022 URL: https://api.semanticscholar.org/CorpusID:253158046
  * [4] Frederik Benzing “Gradient Descent on Neurons and its Link to Approximate Second-order Optimization” In _Proceedings of the 39th International Conference on Machine Learning_ 162, Proceedings of Machine Learning Research PMLR, 2022, pp. 1817–1853 URL: https://proceedings.mlr.press/v162/benzing22a.html
  * [5] “Machine Learning and Knowledge Discovery in Databases - European Conference, ECML PKDD 2019, Würzburg, Germany, September 16-20, 2019, Proceedings, Part III” In _ECML/PKDD (3)_ 11908, Lecture Notes in Computer Science Springer, 2020 URL: http://dblp.uni-trier.de/db/conf/pkdd/pkdd2019-3.html
  * [6] Tom B Brown et al. “Language models are few-shot learners” In _arXiv preprint arXiv:2005.14165_ , 2020
  * [7] Ning Ding et al. “Parameter-efficient fine-tuning of large-scale pre-trained language models” In _Nature Machine Intelligence_ 5, 2023, pp. 1–16 DOI: 10.1038/s42256-023-00626-4
  * [8] Alexey Dosovitskiy et al. “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale” In _International Conference on Learning Representations_ , 2021 URL: https://openreview.net/forum?id=YicbFdNTTy
  * [9] John Duchi, Elad Hazan and Yoram Singer “Adaptive Subgradient Methods for Online Learning and Stochastic Optimization” In _Journal of Machine Learning Research_ 12.61, 2011, pp. 2121–2159 URL: http://jmlr.org/papers/v12/duchi11a.html
  * [10] Kaixin Gao et al. “Eigenvalue-Corrected Natural Gradient Based on a New Approximation” In _Asia-Pacific Journal of Operational Research (APJOR)_ 40.01, 2023, pp. 1–18 DOI: 10.1142/S0217595923400055
  * [11] Thomas George et al. “Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis” In _Neural Information Processing Systems_ , 2018 URL: https://api.semanticscholar.org/CorpusID:47018147
  * [12] Vineet Gupta, Tomer Koren and Yoram Singer “Shampoo: Preconditioned Stochastic Tensor Optimization”, 2018 URL: http://proceedings.mlr.press/v80/gupta18a/gupta18a.pdf
  * [13] Xuehai He et al. “Parameter-efficient model adaptation for vision transformers” In _Proceedings of the Thirty-Seventh AAAI Conference on Artificial Intelligence and Thirty-Fifth Conference on Innovative Applications of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelligence_ , AAAI’23/IAAI’23/EAAI’23 AAAI Press, 2023 DOI: 10.1609/aaai.v37i1.25160
  * [14] Edward J Hu et al. “LoRA: Low-Rank Adaptation of Large Language Models” In _International Conference on Learning Representations_ , 2022 URL: https://openreview.net/forum?id=nZeVKeeFYf9
  * [15] Diederik P. Kingma and Jimmy Ba “Adam: A Method for Stochastic Optimization” In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015 URL: http://arxiv.org/abs/1412.6980
  * [16] Andrew V. Knyazev “Toward the Optimal Preconditioned Eigensolver: Locally Optimal Block Preconditioned Conjugate Gradient Method” In _SIAM Journal on Scientific Computing_ 23.2, 2001, pp. 517–541 DOI: 10.1137/S1064827500366124
  * [17] A. Krizhevsky and G. Hinton “Learning multiple layers of features from tiny images” In _Master’s thesis, Department of Computer Science, University of Toronto_ Citeseer, 2009
  * [18] Frederik Kunstner, Lukas Balles and Philipp Hennig “Limitations of the Empirical Fisher Approximation for Natural Gradient Descent” In _Proceedings of the 33rd International Conference on Neural Information Processing Systems_ Red Hook, NY, USA: Curran Associates Inc., 2019
  * [19] Yann LeCun and Corinna Cortes “MNIST handwritten digit database”, http://yann.lecun.com/exdb/mnist/, 2010 URL: http://yann.lecun.com/exdb/mnist/
  * [20] Brian Lester, Rami Al-Rfou and Noah Constant “The Power of Scale for Parameter-Efficient Prompt Tuning” In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ OnlinePunta Cana, Dominican Republic: Association for Computational Linguistics, 2021, pp. 3045–3059 DOI: 10.18653/v1/2021.emnlp-main.243
  * [21] Brian Lester, Rami Al-Rfou and Noah Constant “The Power of Scale for Parameter-Efficient Prompt Tuning” In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ OnlinePunta Cana, Dominican Republic: Association for Computational Linguistics, 2021, pp. 3045–3059 DOI: 10.18653/v1/2021.emnlp-main.243
  * [22] Ilya Loshchilov and Frank Hutter “Decoupled Weight Decay Regularization” In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_ OpenReview.net, 2019 URL: https://openreview.net/forum?id=Bkg6RiCqY7
  * [23] James Martens “Deep learning via Hessian-free optimization.” In _ICML_ Omnipress, 2010, pp. 735–742 URL: http://dblp.uni-trier.de/db/conf/icml/icml2010.html#Martens10
  * [24] James Martens “New Insights and Perspectives on the Natural Gradient Method” In _Journal of Machine Learning Research_ 21.146, 2020, pp. 1–76 URL: http://jmlr.org/papers/v21/17-678.html
  * [25] James Martens and Roger Grosse “Optimizing Neural Networks with Kronecker-Factored Approximate Curvature” In _Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37_ , ICML’15 Lille, France: JMLR.org, 2015, pp. 2408–2417
  * [26] Levon Nurbekyan, Wanzhou Lei and Yunan Yang “Efficient Natural Gradient Descent Methods for Large-Scale PDE-Based Optimization Problems” In _SIAM Journal on Scientific Computing_ 45.4, 2023, pp. A1621–A1655 DOI: 10.1137/22M1477805
  * [27] Kazuki Osawa et al. “Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks” In _2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2018, pp. 12351–12359 URL: https://api.semanticscholar.org/CorpusID:67856631
  * [28] Adam Paszke et al. “PyTorch: An Imperative Style, High-Performance Deep Learning Library” In _Advances in Neural Information Processing Systems 32_ Curran Associates, Inc., 2019, pp. 8024–8035
  * [29] J. Pauloski et al. “Convolutional neural network training with distributed K-FAC” In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ , SC ’20 Atlanta, Georgia: IEEE Press, 2020
  * [30] J. Pauloski et al. “Deep Neural Network Training With Distributed K-FAC” In _IEEE Transactions on Parallel and Distributed Systems_ 33.12, 2022, pp. 3616–3627 DOI: 10.1109/TPDS.2022.3161187
  * [31] Felix Petersen et al. “ISAAC Newton: Input-based Approximate Curvature for Newton’s Method” In _The Eleventh International Conference on Learning Representations_ , 2023 URL: https://openreview.net/forum?id=0paCJSFW7j
  * [32] K.. Petersen and M.. Pedersen “The Matrix Cookbook” Version 20081110 Technical University of Denmark, 2008 URL: http://www2.imm.dtu.dk/pubdb/p.php?3274
  * [33] Daniel Povey, Xiaohui Zhang and Sanjeev Khudanpur “Parallel training of DNNs with Natural Gradient and Parameter Averaging” In _International Conference on Learning Representations 2014_ , 2014
  * [34] Alec Radford et al. “Robust Speech Recognition via Large-Scale Weak Supervision” In _Proceedings of the 40th International Conference on Machine Learning_ 202, Proceedings of Machine Learning Research PMLR, 2023, pp. 28492–28518 URL: https://proceedings.mlr.press/v202/radford23a.html
  * [35] Colin Raffel et al. “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer” In _Journal of Machine Learning Research_ 21.140, 2020, pp. 1–67 URL: http://jmlr.org/papers/v21/20-074.html
  * [36] Yi Ren and Donald Goldfarb “Efficient Subsampled Gauss-Newton and Natural Gradient Methods for Training Neural Networks” In _CoRR_ abs/1906.02353, 2019 arXiv: http://arxiv.org/abs/1906.02353
  * [37] Yi Ren and Donald Goldfarb “Tensor Normal Training for Deep Learning Models” In _Advances in Neural Information Processing Systems_ 34 Curran Associates, Inc., 2021, pp. 26040–26052 URL: https://proceedings.neurips.cc/paper_files/paper/2021/file/dae3312c4c6c7000a37ecfb7b0aeb0e4-Paper.pdf
  * [38] Herbert Robbins and Sutton Monro “A Stochastic Approximation Method” In _The Annals of Mathematical Statistics_ 22.3 Institute of Mathematical Statistics, 1951, pp. 400–407 DOI: 10.1214/aoms/1177729586
  * [39] Nicolas Roux, Pierre-antoine Manzagol and Yoshua Bengio “Topmoumoute Online Natural Gradient Algorithm” In _Advances in Neural Information Processing Systems_ 20 Curran Associates, Inc., 2007 URL: https://proceedings.neurips.cc/paper_files/paper/2007/file/9f61408e3afb633e50cdf1b20de6f466-Paper.pdf
  * [40] John Schulman et al. “Trust Region Policy Optimization” In _Proceedings of the 32nd International Conference on Machine Learning_ 37, Proceedings of Machine Learning Research Lille, France: PMLR, 2015, pp. 1889–1897 URL: https://proceedings.mlr.press/v37/schulman15.html
  * [41] Jonathan Richard Shewchuk “An Introduction to the Conjugate Gradient Method Without the Agonizing Pain”, 1994 URL: http://www.cs.cmu.edu/~quake-papers/painless-conjugate-gradient.pdf
  * [42] Valentin Thomas et al. “On the interplay between noise and curvature and its effect on optimization and generalization” In _International Conference on Artificial Intelligence and Statistics_ , 2019
  * [43] Hugo Touvron et al. “LLaMA: Open and Efficient Foundation Language Models”, 2023 arXiv:2302.13971 [cs.CL]
  * [44] Yuichiro Ueno et al. “Rich Information is Affordable: A Systematic Performance Analysis of Second-order Optimization Using K-FAC” In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, KDD ’20 Virtual Event, CA, USA: Association for Computing Machinery, 2020, pp. 2145–2153 DOI: 10.1145/3394486.3403265
  * [45] Alex Wang et al. “GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding” In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ Brussels, Belgium: Association for Computational Linguistics, 2018, pp. 353–355 DOI: 10.18653/v1/W18-5446
  * [46] Minghan Yang et al. “An Efficient Fisher Matrix Approximation Method for Large-Scale Neural Network Optimization” In _IEEE Transactions on Pattern Analysis and Machine Intelligence_ 45.5, 2023, pp. 5391–5403 DOI: 10.1109/TPAMI.2022.3213654
  * [47] Minghan Yang et al. “Sketch-Based Empirical Natural Gradient Methods for Deep Learning” In _J. Sci. Comput._ 92.3 USA: Plenum Press, 2022 DOI: 10.1007/s10915-022-01911-x
  * [48] Guodong Zhang, James Martens and Roger B Grosse “Fast Convergence of Natural Gradient Descent for Over-Parameterized Neural Networks” In _Advances in Neural Information Processing Systems_ 32 Curran Associates, Inc., 2019 URL: https://proceedings.neurips.cc/paper_files/paper/2019/file/1da546f25222c1ee710cf7e2f7a3ff0c-Paper.pdf
  * [49] Lin Zhang, Shaohuai Shi and Bo Li “Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation” In _The Eleventh International Conference on Learning Representations_ , 2023 URL: https://openreview.net/forum?id=_Mic8V96Voy
  * [50] Yingqiu Zhu et al. “Automatic, dynamic, and nearly optimal learning rate specification via local quadratic approximation” In _Neural Networks_ 141, 2021, pp. 11–29 DOI: https://doi.org/10.1016/j.neunet.2021.03.025

##  Appendix A Relation between iEF, GN and NGD in Different Machine Learning
Setups

The connection between iEF, GN and NGD methods is discussed for different
common machine-learning scenarios. It is explained that for a scalar output
model, the iEF method and the GN algorithm is equivalent. Furthermore, for a
least-squares problem, the iEF, GN and NGD methods are all equivalent.

###  A.1 Scalar Output Model Problem

In this section, a supervised learning problem for a target model with only a
scalar output is discussed. The definition of the setup is as follows.

Consider a target scalar output model
$f_{\bm{\theta}}(\cdot)\in\mathbb{R}^{P}\to\mathbb{R}$, where
$\bm{\theta}\in\mathbb{R}^{P}$ is the trainable parameters of size $P$. Given
$N$ i.i.d. training samples $(\bm{x}_{n},y_{n})_{n=1}^{N}$ (where $\bm{x}_{n}$
is the input feature vector and $y_{n}\in\mathbb{R}$ is the scalar label), the
output of the model for the $n$-th sample is
$z_{n}=f_{\bm{\theta}}(\bm{x}_{n})$ where $z_{n}\in\mathbb{R}$ is a scalar.
For the $n$-th sample, the per-sample loss $l_{n}$ is defined as

$l_{n}=\mathcal{F}^{\text{obj}}(z_{n},y_{n}),$ (18)

where $\mathcal{F}^{\text{obj}}(\cdot)\in\mathbb{R}\to\mathbb{R}$ is the per-
sample objective function, and the accumulated loss $\sum\nolimits_{n}l_{n}$
is to be minimised. Some examples of this problem setup include least-squares
regression problems where a mean-square error is used as the loss function,
and binary classification problems where the Sigmoid-activation plus binary
cross-entropy is used as the loss function.

Recall the definition of the iEF matrix in Eqn. (15). In this problem setup,
the iEF matrix can be computed as follows

$\tilde{\textbf{F}}^{\star}=\sum\nolimits_{n}\|\dfrac{\operatorname{d\\!}{}l_{n}}{\operatorname{d\\!}{z_{n}}}\|_{2}^{-2}(\dfrac{\operatorname{d\\!}{}l_{n}}{\operatorname{d\\!}{z_{n}}}\nabla_{\bm{\theta}}z_{n})(\dfrac{\operatorname{d\\!}{}l_{n}}{\operatorname{d\\!}{z_{n}}}\nabla_{\bm{\theta}}z_{n})^{\top}=\sum\nolimits_{n}\nabla_{\bm{\theta}}z_{n}\nabla_{\bm{\theta}}z_{n}^{\top}=\hat{\mathbf{G}},$
(19)

which takes the same form as the GN matrix. Therefore, the iEF and GN methods
are equivalent to supervised learning problems of scalar output models.

###  A.2 Linear Least-Squares Problem

The linear least-squares problem is a representative example of supervised
learning problems of scalar output models, which includes the toy example
shown in Fig. 1. The definition of this type of problem follows most of the
definitions in Appendix A.1, apart from the per-sample objective function,
which is elaborated as follows. For the linear least-squares problem, the
objective function is defined as follows

$\mathcal{F}^{\text{obj}}(z_{n},y_{n})=\frac{1}{2}(z_{n}-y_{n})^{2}.$ (20)

The Fisher matrix for the linear least-squares problem is defined as follows
[42]

$\mathbf{F}=\sum\nolimits_{n}\nabla_{\bm{\theta}}z_{n}\nabla_{\bm{\theta}}z_{n}^{\top}=\tilde{\textbf{F}}^{\star}=\hat{\mathbf{G}},$
(21)

which coincides with the iEF and GN matrix in this problem setup.
Consequently, the iEF, GN and NGD methods are equivalent for linear least-
squares problems. This is verified in the second plot of Fig. 1 where NGD and
iEF share the same update vector field.

##  Appendix B Visualisation for Linear Least-Squares Problem with Two
Training Samples

##### Setup Description

In this section, the Fig. 1 referenced in Sec. 4.2 is explained in detail. The
vector field graph is based on a simple linear least-squares regression
problem, with 2 training samples and 2 trainable parameters.

The linear least-squares problem is chosen for visualisation because it is not
only a highly representative machine learning setup [42, 18], but also there
has been a precedent of using this problem to visualise the distortion in EF
(see Figure 1 in [18]). The trainable parameter size of $P=2$ is chosen to
facilitate the 2D visualisation, and the training sample size of $N=2$ (i.e.
$N\leq P$) is chosen to better match the practical deep learning scenarios
where the model is over-parameterised.

The target linear model is formulated as
$f_{\bm{\theta}}(x_{n})=\bm{\theta}^{\top}[x_{n},1]^{\top}=\theta_{0}+\theta_{1}x_{n}$
with $x_{n}\in\mathbb{R}$ and $\bm{\theta}\in\mathbb{R}^{2}$. The two training
samples are $(x_{n},y_{n})_{n=1}^{N}=\\{(0,0),(1,0)\\}$ (with $N=2$) and the
target loss to be minimised is
$\mathcal{L}(\bm{\theta})=\sum_{n=1}^{2}\frac{1}{2}[y_{n}-f_{\bm{\theta}}(x_{n})]^{2}$.
It is obvious that the optimal parameter (at global minimum) is
$\bm{\theta}^{\star}=[0,0]^{\top}$.

The update vector fields are generated using the following definition:

$\begin{cases}\Delta\bm{\theta}_{\text{SGD}}&=\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})\\\
\Delta\bm{\theta}_{\text{NGD}}&=(\mathbf{F}+\lambda_{\text{NGD}}\mathbf{I})^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})\\\
\Delta\bm{\theta}_{\text{EF}}&=(\tilde{\mathbf{F}}+\lambda_{\text{EF}}\mathbf{I})^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})\\\
\Delta\bm{\theta}_{\text{iEF}}&=(\tilde{\mathbf{F}}^{\star}+\lambda_{\text{iEF}}\mathbf{I})^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}).\end{cases}$
(22)

The EF matrix $\tilde{\mathbf{F}}$ follows the definition in Eqn. (3). The iEF
matrix $\tilde{\mathbf{F}}^{\star}$ and Fisher matrix $\mathbf{F}$ shares the
same definition as shown in Eqn. (21) (see Appendix A.2). For this toy
problem, the damping factors $\lambda_{\text{iEF}}$ and $\lambda_{\text{NGD}}$
are set to zero. The damping factor $\lambda_{\text{EF}}$ is set to zero
everywhere, apart from when one of the per-sample gradient gives a norm of 0,
a damping factor of $1\times 10^{-4}\max(\text{diag}(\tilde{\mathbf{F}}))$ is
used to facilitate inversion. Finally, in the third plot of Fig. 1, the EF
update vectors are normalised for better visualisation, because EF updates
have hugely different scales across the contour plots. In the zoomed view next
to the third plot, the EF update vectors are not normalised, which should give
a better demonstration of the inverse scaling issue of EF updates.

The same set of 5 initial parameters (shown as black solid octagons) are used
to generate 5 training trajectories (coloured curves) on each plot in Fig. 1.
Each follows the corresponding update vector fields (each step follows the
direction of the update vector field, and is normalised to
$1\text{\times}{10}^{-2}$). To demonstrate how each sample affects the EF
updates, the optimal parameter set for each training sample (i.e. when
$\frac{1}{2}[y_{n}-f_{\bm{\theta}}(x_{n})]^{2}=0$) is added on all plots in
Fig. 1 as dashed lines. For training sample 1: (0, 0), the optimal parameter
set is $\theta_{0}=0$, which is shown as horizontal dashed lines. For training
sample 2: (1, 0), the optimal parameter set is $\theta_{0}+\theta_{1}=0$,
which is shown as diagonal dashed lines.

##### Observations

By observing the behaviour of the vector fields for different updates on the
loss landscape, along with the 5 sampled trajectories in Fig. 1, it can be
seen that all methods successfully reached the global minimum, with the
NGD/iEF updates having the most efficient training trajectory. However, the EF
method has highly distorted update vector field and training trajectories due
to the inversely-scaled projection issue discussed in Sec. 4.1. The following
conclusions can be drawn for EF updates:

  1. 1.

Directional bias towards converged samples:  The EF update vector field is
easily “attracted” to the dashed lines because the EF update is strongly
affected by the better-converged sample (due to inverse scaling). This leads
to a distorted update vector field and ineffective training trajectories. In
particular, when per-sample gradients have a highly imbalanced norm near the
dashed lines, the EF update vector field becomes almost parallel to the loss
contour. This causes them to deviate significantly from both the gradient
descent direction and the optimal direction (NG update direction), and
oscillations can be observed along the dashed line in the zoomed view.

  2. 2.

Inversely-scaled update norm:  EF updates have larger norms when either of the
training samples is nearly converged (see the zoomed plot). In fact, the EF
training trajectories managed to converge because we normalised every update
to a small norm of $1\text{\times}{10}^{-2}$, which is not needed for SGD, NGD
and iEF methods. Consequently, it is likely that a sophisticated learning rate
scheduler is necessary when training with the EF method in practical deep
learning setups (as is mentioned in [18]).

##  Appendix C Convergence Analysis

The proofs for Theorem. 5.3 and Theorem. 5.3 are provided in this section. The
proofs follow and extend the continuous-time analysis framework in [48] and
[2].

##### Justification of Full-rank Gram Matrix Assumption

Recall Assumption 5.2 which assumes
$[\nabla_{\bm{\theta}(t)}\bm{l}(t)][\nabla_{\bm{\theta}(t)}\bm{l}(t)]^{\top}$
to be full-rank throughout the training process. This is equivalent to
assuming $\nabla_{\bm{\theta}(t)}\bm{l}(t)\in\mathbb{R}^{N\times P}$ always
have full row rank. This is considered a mild assumption for practical deep
learning because:

  1. 1.

It is generally true that for each individual per-sample gradient
$\nabla_{\bm{\theta}(t)}l_{n}(t)$, the norm of the gradient will be zero if
and only if $\nabla_{\bm{z}_{n}(t)}l_{n}(t)=\bm{0}$ (i.e. at the minimum of
that sample). Therefore, it is unlikely that there exists zero norm per-sample
gradients in $\nabla_{\bm{\theta}(t)}\bm{l}(t)$ during training.

  2. 2.

The target model is usually highly over-parameterised with $P\gg N$ and the
model is highly complex. It is in general unlikely that
$\nabla_{\bm{\theta}(t)}\bm{l}(t)$ is rank-deficient, as long as there are no
duplicate training samples.

##### Ideal Per-sample Loss Change for IEF

Given Assumption 5.1 and Assumption 5.2, the following Lemma can be derived

######  Lemma C.1.

Suppose Assumption 5.2 holds, $\forall n\in\\{1,\dots,N\\}$, the loss
reduction induced on the $n$-th training sample by the iEF update follows

$\dfrac{\operatorname{d\\!}{}l_{n}(t)}{\operatorname{d\\!}{t}}=-\|\nabla_{\bm{z}_{n}(t)}l_{n}(t)\|^{2}_{2}.$
(23)

Proof of Lemma. C.1:
$\dfrac{\operatorname{d\\!}{}\bm{l}(t)}{\operatorname{d\\!}{t}}$ can be
computed as follows

$\dfrac{\operatorname{d\\!}{}\bm{l}(t)}{\operatorname{d\\!}{t}}=\nabla_{\bm{\theta}(t)}{\bm{l}(t)}\dfrac{\operatorname{d\\!}{}\bm{\theta}(t)}{\operatorname{d\\!}{t}}=-\nabla_{\bm{\theta}(t)}{\bm{l}(t)}\nabla_{\bm{\theta}(t)}{\bm{l}(t)}^{\top}[\nabla_{\bm{\theta}(t)}\bm{l}(t)\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}]^{-1}\bm{s}_{\text{iEF}}(t)=-\bm{s}_{\text{iEF}}(t).$
(24)

Therefore,
$\dfrac{\operatorname{d\\!}{}l_{n}(t)}{\operatorname{d\\!}{t}}=[\dfrac{\operatorname{d\\!}{}\bm{l}(t)}{\operatorname{d\\!}{t}}]_{n}=-\|\nabla_{\bm{z}_{n}(t)}l_{n}(t)\|^{2}_{2}$.

##### Global Sub-linear Convergence for Softmax + CE Objective Function

The proof for Theorem 5.3 is provided here. In this analysis, the supervised
learning setup with softmax + CE objective function is considered (as
described in Sec. 3). This is a setup commonly used in practical deep
learning, making this convergence analysis practically relevant.

As stated in Sec. 3, the target model is assumed to use a softmax activation
function. For the $n$-th sample at time step $t$,
$\bm{z}_{n}(t)=f_{\bm{\theta}(t)}(\bm{x}_{n})$, where
$\bm{z}_{n}(t)\in\mathbb{R}^{C}$ is the output logits. The output probability
for class $c$ can be computed from the logits using the softmax activation:
$[p_{n}(c)](t)=[\bm{\sigma}_{\text{SM}}(\bm{z}_{n}(t))]_{c}$ where
$\bm{\sigma}_{\text{SM}}(\cdot)\mathrel{\mathop{\ordinarycolon}}\mathbb{R}^{C}\to\mathbb{R}^{C}$
is the softmax activation function.

The per-sample loss in this setup is therefore defined as

$l_{n}(t)\mathrel{\mathop{\ordinarycolon}}=-\log[p_{n}(y_{n})](t)\mathrel{\mathop{\ordinarycolon}}=-\log\hat{p}_{n}(t)=-\log\left[\bm{\sigma}_{\text{SM}}(\bm{z}_{n}(t))\right]_{y_{n}},$
(25)

where $\left[\bm{\sigma}_{\text{SM}}(\bm{z}_{n})\right]_{y_{n}}$ represents
the output probability for the target class $y_{n}$. It can be seen that the
lowest loss for sample $n$ tends to 0, which is achieved when
$\hat{p}_{n}(t)\to 1$.

The gradient of $l_{n}(t)$ w.r.t. $\bm{z}_{n}(t)$ can be computed as

$\left[\nabla_{\bm{z}_{n}(t)}l_{n}(t)\right]_{c}=\begin{cases}-[p_{n}(c)](t)&c\neq
y_{n},\\\ 1-\hat{p}_{n}(t)&c=y_{n}.\end{cases}$ (26)

The norm of the gradient satisfies the following inequality

$\|\nabla_{\bm{z}_{n}(t)}l_{n}(t)\|_{2}^{2}\geq[1-\hat{p}_{n}(t)]^{2}.$ (27)

Combining Lemma. C.1 with the definition in Eqn. (25), the following equation
can be obtained

$\dfrac{\operatorname{d\\!}{}l_{n}(t)}{\operatorname{d\\!}{t}}=-\dfrac{\operatorname{d\\!}{}\log\hat{p}_{n}(t)}{\operatorname{d\\!}{t}}=-\frac{1}{\hat{p}_{n}(t)}\dfrac{\operatorname{d\\!}{}\hat{p}_{n}(t)}{\operatorname{d\\!}{t}}=-\|\nabla_{\bm{z}_{n}(t)}l_{n}(t)\|_{2}^{2}.$
(28)

Combining Eqn. (27)(28), the following inequality can be obtained

$\frac{1}{[\hat{p}_{n}(t)][1-\hat{p}_{n}(t)]^{2}}\text{d}\hat{p}_{n}(t)\geq\text{d}t.$
(29)

Integrating on both sides gives

$\frac{1}{1-\hat{p}_{n}(t)}+\log\frac{\hat{p}_{n}(t)}{1-\hat{p}_{n}(t)}-C_{0}\geq
t.$ (30)

where
$C_{0}=\frac{1}{1-\hat{p}_{n}(0)}+\log\frac{\hat{p}_{n}(0)}{1-\hat{p}_{n}(0)}$.
It is known that $x>\log x$ for $x>0$, therefore the inequality can be further
relaxed to

$\frac{1}{1-\hat{p}_{n}(t)}+\frac{\hat{p}_{n}(t)}{1-\hat{p}_{n}(t)}>t+C_{0},$
(31)

which is equivalent to

$\frac{2}{1-\hat{p}_{n}(t)}>t+C_{0}+1.$ (32)

Now consider a large $t$, s.t. $t+C_{0}+1>0$, a bound can be provided for
$\hat{p}_{n}(t)$ as follows

$\hat{p}_{n}(t)>1-\frac{2}{t+C_{0}+1}.$ (33)

This shows that iEF guarantees global sub-linear convergence for target
probability $\hat{p}_{n}\to 1$ for every training sample. It equivalently
implies that iEF guarantees sub-linear convergence to the global minimum of
the accumulated cross-entropy loss (achieved when $\forall n=\\{1,\dots N\\}$,
$\hat{p}_{n}=1$), given the per-sample objective function is the softmax +
cross-entropy combination.

###  C.1 Global Linear Convergence for Strongly Convex Objective Functions

The proof for Theorem 5.4 is provided here. Unlike for Theorem 5.3, where the
model is assumed to use a softmax + cross-entropy loss function, this theorem
is applicable to a group of common loss functions: $m$-strongly convex
objective functions.

The details of the setup used here mostly follow that in Sec. 3. The
differences are described as follows. Given $N$ i.i.d. training samples
$(\bm{x}_{n},y_{n})_{n=1}^{N}$ (where $\bm{x}_{n}$ is the input feature vector
and $y_{n}$ is a label of arbitrary form), the output of the model for the
$n$-th sample is $\bm{z}_{n}=f_{\bm{\theta}}(\bm{x}_{n})$ where
$z_{n}\in\mathbb{R}^{C}$ is a general model output vector (no longer logits
vector). A $m$-strongly convex objective loss function
$\mathcal{F}_{\text{obj}}(\cdot)$ is used to compute the per-sample loss for
the $n$-th sample $l_{n}\in\mathbb{R}$ as follows

$l_{n}=\mathcal{F}_{\text{obj}}(\bm{z}_{n},y_{n})\coloneq\mathcal{F}_{\text{obj}}(\bm{z}_{n}),$
(34)

where the label $y_{n}$ is omitted for simplicity. Finally, the following
accumulated loss is to be minimised

$\mathcal{L}(\bm{\theta})=\sum\nolimits_{n}l_{n}.$ (35)

In addition to Assumption 5.1 and 5.2, an additional assumption is made for
the objective loss function as follows

######  Assumption C.2.

For the $n$-th sample, the objective loss function
$\mathcal{F}_{\text{obj}}(\bm{z}_{n})$ is assumed to be $m$-strongly convex on
the model output space $\bm{z}_{n}$, which then satisfies the following
Polyak-Lojasiewicz inequality [5]

$\mathcal{F}_{\text{obj}}(\bm{z}_{n})-\mathcal{F}_{\text{obj}}(\bm{z}_{n}^{\star})\leq\frac{1}{2m}\|\nabla_{\bm{z}_{n}}\mathcal{F}_{\text{obj}}(\bm{z}_{n})\|^{2},$
(36)

where $\mathcal{F}_{\text{obj}}(\bm{z}_{n}^{\star})$ is the global minimum of
the loss for the $n$-th sample. For simplicity, the notation $l_{n}^{\star}$
is used in place of $\mathcal{F}_{\text{obj}}(\bm{z}_{n}^{\star})$. The
inequality is therefore rewritten as follows

$l_{n}-l_{n}^{\star}\leq\frac{1}{2m}\|\nabla_{\bm{z}_{n}}l_{n}\|^{2}.$ (37)

Assumption. C.2 is quoted from [2], which covers a wide range of loss
functions including mean-square-error. Note that under this assumption, both
per-sample losses and accumulated loss can still have an arbitrarily non-
convex landscape on the parameter space.

Based on Lemma C.1, when training with iEF updates, for the $n$-th sample, the
loss change is

$\dfrac{\operatorname{d\\!}{}l_{n}(t)}{\operatorname{d\\!}{t}}=-\|\nabla_{\bm{z}_{n}(t)}l_{n}(t)\|_{2}^{2}$
(38)

Using the Polyak-Lojasiewicz inequality in Assumption. C.2, the following
inequality can be obtained

$\dfrac{\operatorname{d\\!}{}(l_{n}(t)-l_{n}^{\star})}{\operatorname{d\\!}{t}}\leq-2m(l_{n}(t)-l_{n}^{\star})\Rightarrow\frac{1}{(l_{n}(t)-l_{n}^{\star})}\dfrac{\operatorname{d\\!}{}(l_{n}(t)-l_{n}^{\star})}{\operatorname{d\\!}{t}}\leq-2m.$
(39)

Integrating on both sides gives the following

$l_{n}(t)-l_{n}^{\star}\leq e^{-2mt}(l_{n}(0)-l_{n}^{\star}),$ (40)

which shows that iEF pre-conditioned gradient flow linearly converges to the
global minimum for every sample. This then implies that iEF has a global
linear convergence guarantee for an accumulated loss
$\mathcal{L}(\bm{\theta})$.

##  Appendix D Discussion on Practical Applications of IEF

###  D.1 Stochastic IEF/EF Optimiser

As is mentioned in Sec. 5.4, the stochastic version of the exact iEF method
can be directly used as an optimiser. It is described in Algorithm 1. In this
section, the implementation of this exact iEF optimiser is discussed, along
with its computation and memory complexity. Note that the stochastic exact EF
optimiser can be constructed in a similar form in Algorithm 2. The following
discussions should also apply to the exact EF optimiser.

Algorithm 1 Stochastic Optimisation with Exact IEF

Require: All $N$ training samples, initial model parameters $\bm{\theta}(0)$

for $t=0$ to $T-1$ do

Sample Mini-batch $\mathcal{M}(t)$ with size $M$

Perform forward pass on batch $\mathcal{M}(t)$ to obtain the per-sample loss
vector $\bm{l}(t)$

Perform back-propagation to obtain Jacobian $\nabla_{\bm{\theta}(t)}\bm{l}(t)$
and logits gradient norm vector $\bm{s}_{\text{iEF}}(t)$

Compute iEF update
$\Delta\bm{\theta}_{\text{iEF}}(t)=\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}(\nabla_{\bm{\theta}(t)}\bm{l}(t)\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}+\lambda\mathbf{I})^{-1}\bm{s}_{\text{iEF}}(t)$:
Eqn. (8)

Update model
$\bm{\theta}(t+1)=\bm{\theta}(t)-\eta\Delta\bm{\theta}_{\text{iEF}}(t)$

end for

Algorithm 2 Stochastic Optimisation with Exact EF

Require: All $N$ training samples, initial model parameters $\bm{\theta}(0)$

for $t=0$ to $T-1$ do

Sample Mini-batch $\mathcal{M}(t)$ with size $M$

Perform forward pass on batch $\mathcal{M}(t)$ to obtain the per-sample loss
vector $\bm{l}(t)$

Perform back-propagation to obtain Jacobian $\nabla_{\bm{\theta}(t)}\bm{l}(t)$

Compute EF update
$\Delta\bm{\theta}_{\text{EF}}(t)=\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}(\nabla_{\bm{\theta}(t)}\bm{l}(t)\nabla_{\bm{\theta}(t)}\bm{l}(t)^{\top}+\lambda\mathbf{I})^{-1}\bm{1}$:
Eqn. (5)

Update model
$\bm{\theta}(t+1)=\bm{\theta}(t)-\eta\Delta\bm{\theta}_{\text{EF}}(t)$

end for

##### Implementation Details

There are two main aspects of the implementation of this optimiser that are
non-trivial. These are discussed respectively as follows:

  1. 1.

Jacobian matrix $\nabla_{\bm{\theta}}\bm{l}$:

The computation of $\nabla_{\bm{\theta}}\bm{l}$ is effectively collecting the
per-sample gradients for a batch during the back-propagation process. In
Pytorch [28], the per-sample gradients are readily computed during back-
propagation, but are usually accumulated along the batch dimension to compute
the total gradient, and are not available for collection directly. Therefore,
additional backward hooks need to be attached to trainable modules to store
these per-sample gradients during the back-propagation process. This is a
standard procedure used in most approximate NGD optimisers [25, 4, 39]. Our
implementation is partially based on this K-FAC implementation. Note that this
additional procedure of computing per-sample gradients is negligible w.r.t.
forward and backward of large pre-trained models in PEFT setups, making iEF/EF
have comparable speed as standard AdamW/Adafactor/SGD optimisers.

  2. 2.

Logits gradient norm vector $\bm{s}_{\text{iEF}}$:

During back-propagation in Pytorch [28], the gradient of logits
$\nabla_{\bm{z}_{n}}l_{n}$ is already computed, but is not stored because
logits vector $\bm{z}_{n}$ is not a leaf-node. This can be easily changed by
calling the “.retain_grad()” method on the logits vector. Although this
operation is non-standard to common approximate NGD optimisers, it adds
negligible cost to standard back-propagation (in general deep learning setup,
not limited to PEFT) and its effect on speed can be ignored.

##### Time and Memory Complexity

Given a batch size of $M$, and trainable parameters of size $P$, assume the
per-sample gradients have already been computed through backpropagation, the
time complexity of computing each update is $O(M^{3}+M^{2}P)$, and the memory
complexity is $O(M^{2}+MP)$. Due to model over-parameterisation, we have $P\gg
M$. Then, the time complexity becomes $O(M^{2}P)$ and memory complexity
becomes $O(MP)$. In practical deep learning frameworks such as Pytorch [28],
the limiting factor for the applicability of such an exact EF-like method is
the memory complexity $O(MP)$, which is essentially the storage of the $M$
per-sample gradients involved in the computation of exact iEF updates. It is
therefore only possible to apply exact EF-like methods to models with small
trainable parameter sizes (either full tuning of small models or parameter-
efficient fine-tuning of large models) or small batch size $M$. This is the
reason why exact EF is never directly used to optimise modern deep learning
models, and additional approximation is always necessary (e.g. K-FAC [25] or
SVD-based pruning [39]). Nevertheless, given the rise of large pre-trained
models [43, 6, 35, 34], parameter-efficient fine-tuning has gained traction
[13, 14, 7] and direct application of exact iEF may still be beneficial (as
the trainable parameter size is usually $<1\%$ of the pre-trained model size).

###  D.2 Improving Existing EF-based Optimisers with IEF

It is mentioned in Appendix D.1 that exact iEF as an optimiser is limited by
its memory complexity. Given the descent training performance of exact iEF
optimiser, it would be interesting to observe the performance of iEF on more
general setups. As is mentioned in Sec. 5.4, it can be achieved by
incorporating iEF into existing EF-based optimisers such as empirical K-FAC
[25]. In this section, an improvement to the EF-based K-FAC optimiser with iEF
is proposed, which can act as a starting point for future work in improving
other approximate empirical NGD methods.

##### Expression of IEF Matrix

The iEF matrix can be computed according to Eqn. (15):
$\mathbf{F}^{\star}=\nabla_{\bm{\theta}}\bm{l}^{\top}\text{diag}(\bm{s}_{\text{iEF}})^{-1}\nabla_{\bm{\theta}}\bm{l}$.
This can be derived from Eqn. (8) using the identity
$(\mathbf{U}\mathbf{V}+\lambda\mathbf{I})^{-1}\mathbf{U}=\mathbf{U}(\mathbf{V}\mathbf{U}+\lambda\mathbf{I})^{-1}$.
For simplicity of derivation, we denote
$\mathbf{J}=\nabla_{\bm{\theta}}\bm{l}$ and assume $\mathbf{J}$ is a square
matrix and is full-rank. Also, we denote
$\mathbf{S}=\text{diag}(\bm{s}_{\text{iEF}})^{\frac{1}{2}}$ and assume the
diagonal matrix $\mathbf{S}$ is full-rank. Therefore, the proposed iEF matrix
can be expressed as
$\mathbf{F}^{\star}=\mathbf{J}^{\top}\mathbf{S}^{-2}\mathbf{J}$. The following
derivation can be made

$\displaystyle(\mathbf{F}^{\star})^{-1}(\nabla_{\bm{\theta}}\bm{l}^{\top}\bm{1})$
$\displaystyle=(\mathbf{J}^{\top}\mathbf{S}^{-2}\mathbf{J})^{-1}\mathbf{J}^{\top}\bm{1}=[(\mathbf{J}^{\top}\mathbf{S}^{-1})(\mathbf{S}^{-1}\mathbf{J})]^{-1}(\mathbf{J}^{\top}\mathbf{S}^{-1})(\mathbf{S}\bm{1})$
(41)
$\displaystyle=(\mathbf{J}^{\top}\mathbf{S}^{-1})[(\mathbf{S}^{-1}\mathbf{J})(\mathbf{J}^{\top}\mathbf{S}^{-1})]^{-1}\mathbf{S}\bm{1}=\mathbf{J}^{\top}\mathbf{S}^{-1}\mathbf{S}[\mathbf{J}\mathbf{J}^{\top}]^{-1}\mathbf{S}\mathbf{S}\bm{1}$
$\displaystyle=\mathbf{J}^{\top}[\mathbf{J}\mathbf{J}^{\top}]^{-1}\mathbf{S}^{2}\bm{1}=\mathbf{J}^{\top}[\mathbf{J}\mathbf{J}^{\top}]^{-1}\bm{s}_{\text{iEF}}.$

This shows that the update pre-conditioned by $\mathbf{F}^{\star}$ indeed is
the same as the iEF update described in Eqn. (8).

The expression for the iEF matrix in Eqn. (15) can be alternatively rewritten
as follows

$\mathbf{F}^{\star}=\nabla_{\bm{\theta}}\bm{l}^{\top}\text{diag}(\bm{s}_{\text{iEF}})^{-1}\nabla_{\bm{\theta}}\bm{l}=\sum\nolimits_{n}(\|\nabla_{\bm{z}_{n}}l_{n}\|^{-1}\nabla_{\bm{\theta}}l_{n})(\|\nabla_{\bm{z}_{n}}l_{n}\|^{-1}\nabla_{\bm{\theta}}l_{n})^{\top},$
(42)

which differs from the EF matrix in Eqn. (3) in that each per-sample gradient
is re-scaled with $\|\nabla_{\bm{z}_{n}}l_{n}\|^{-1}$. Such simple scaling is
easy to implement in most approximate empirical NGD optimisers.

##### Improving Empirical K-FAC with IEF

The empirical K-FAC [33] is a widely used version of the K-FAC [29, 30, 44,
27]. Its formulation can be described as follows.

Assume there are $n$ samples in the target batch. For one fully connected
layer in a target model, the weight matrix is denoted by matrix
$\mathbf{W}\in\mathbb{R}^{m\times k}$, the (batched) input is denoted by
$\bm{a}\in\mathbb{R}^{n\times m}$ and the (batched) output is denoted by row
vector $\bm{c}\in\mathbb{R}^{n\times k}$. They satisfy
$\bm{c}=\bm{a}\mathbf{W}$. The gradient
$\nabla_{\bm{c}}\bm{l}\in\mathbb{R}^{n\times k}$ is denoted by $\bm{g}$. The
block-diagonal portion of the EF matrix corresponding to this layer (denoted
as $\tilde{\mathbf{F}}_{\mathbf{W}}$) can be estimated using K-FAC
approximation as follows [25]

$\tilde{\mathbf{F}}_{\mathbf{W}}=\mathbb{E}[(\bm{g}^{\top}\bm{g})\otimes(\bm{a}^{\top}\bm{a})]\approx\mathbb{E}[\bm{g}^{\top}\bm{g}]\otimes\mathbb{E}[\bm{a}^{\top}\bm{a}].$
(43)

This is based on the gradient expression for weight matrix $\mathbf{W}$:
$\nabla_{\mathbf{W}}l=\bm{g}^{\top}\bm{a}$. By rescaling this gradient with
vector $\bm{s}_{\text{iEF}}$, the expression for
$\tilde{\mathbf{F}}_{\mathbf{W}}$ becomes:

$\tilde{\mathbf{F}}^{\star}_{\mathbf{W}}=\mathbb{E}[\bm{g}^{\top}\text{diag}(\bm{s}_{\text{iEF}})^{-1}\bm{g}]\otimes\mathbb{E}[\bm{a}^{\top}\bm{a}],$
(44)

where the $\bm{s}_{\text{iEF}}$ vector is easy to compute in Pytorch [28] with
negligible extra cost (as shown in Appendix D.1), and such diagonal rescaling
is straightforward to implement. In conclusion, the idea of iEF can be easily
integrated into existing approximate empirical NGD optimisers, and it is
interesting to observe the improvements due to such integration.

##  Appendix E Monte-Carlo Sampled Fisher Matrix

A commonly used method to estimate the exact Fisher matrix is through Monte-
Carlo sampling, which is used in the original K-FAC method [25]. This method
is referred to as sampled Fisher in this paper. Particularly, when only one
Monte-Carlo sample is generated for each training sample, the corresponding
method is termed SF for brevity. Recall the definition of the Fisher matrix in
Eqn. (2), it can be rewritten using expectations as follows

$\textbf{F}=\sum\nolimits_{n}\mathbb{E}_{c\sim
p_{\bm{\theta}}(y|\bm{x}_{n})}[\nabla_{\bm{\theta}}\log
p_{n}(c)\nabla_{\bm{\theta}}\log p_{n}(c)^{\top}].$ (45)

This means, by generating enough labels (with Monte-Carlo sampling) from the
output distribution $p_{\bm{\theta}}(y|x_{n})$, and computing the gradient
w.r.t. these labels $\nabla_{\bm{\theta}}\log p_{n}(c)$, the Fisher matrix can
be estimated without bias.

Assume $K$ labels are generated for each training sample, the exact expression
for sampled Fisher with $K$ samples (denoted as $\hat{\mathbf{F}}(K)$) is as
follows

$\hat{\mathbf{F}}(K)=\frac{1}{K}\sum_{n=1}^{N}\sum_{k=1}^{K}[\nabla_{\bm{\theta}}\log
p_{n}(\hat{y}_{n}^{(k)})\nabla_{\bm{\theta}}\log
p_{n}(\hat{y}_{n}^{(k)})^{\top}]=\frac{1}{K}\hat{\mathbf{J}}(K)^{\top}\hat{\mathbf{J}}(K),$
(46)

where $\hat{y}_{n}^{(k)}\sim p_{\bm{\theta}}(y|\bm{x}_{n})$ is the $k$-th
generated label for the $n$-th sample, Jacobian
$\hat{\mathbf{J}}(K)\in\mathbb{R}^{(NK)\times P}$ denotes the stacked sampled
gradients $[\nabla_{\bm{\theta}}\log
p_{1}(\hat{y}_{1}^{(1)}),\dots,\nabla_{\bm{\theta}}\log
p_{n}(\hat{y}_{n}^{(k)}),\dots,\nabla_{\bm{\theta}}\log
p_{N}(\hat{y}_{N}^{(K)})]^{\top}$.

$\hat{\mathbf{F}}(1)$, i.e. sampled Fisher with 1 sampling for each training
sample, is used as a baseline approximate NGD method in Sec. 7 and is compared
against EF and iEF. This has two reasons:

  1. 1.

For $\hat{\mathbf{F}}(K)$, with $K=1$ is more commonly chosen in practice [12,
25] than $K>1$. This is mainly because $\hat{\mathbf{F}}(K)$ requires $K$
additional back-propagations through the target batch, which becomes very
expensive for a large $K$.

  2. 2.

$\hat{\mathbf{F}}(1)$ has a maximum rank of $N$, which is the same rank as the
EF matrix ($\tilde{\mathbf{F}}$) and the iEF matrix
($\tilde{\mathbf{F}}^{\star}$).

Note that even for $\hat{\mathbf{F}}(1)$, as compared to EF and iEF methods,
requires an additional back-propagation through target batches. This makes
$\hat{\mathbf{F}}(1)$ hard to implement in practice, and becomes nearly twice
as expensive as EF/iEF. It is the leading reason that EF is commonly used in
favour of $\hat{\mathbf{F}}(1)$ in practice [29, 30, 44, 27].

##### Exact Pre-conditioning with $\hat{\mathbf{F}}(1)$

To properly evaluate the pre-conditioner $\hat{\mathbf{F}}(1)$ either through
optimisation or our evaluation framework (see Sec. 6), it is necessary to
exactly compute the pre-conditioned gradient. The Sherman-Morrison-Woodbury
(SMW) identity [36] can be used to achieve this, which states that
$(\mathbf{U}^{\top}\mathbf{U}+\lambda\mathbf{I})^{-1}=\frac{1}{\lambda}[\mathbf{I}-\mathbf{U}^{\top}(\mathbf{U}\mathbf{U}^{\top}+\lambda\mathbf{I})^{-1}\mathbf{U}]$.

The update pre-conditioned by $\hat{\mathbf{F}}(1)$ (denoted as
$\Delta\bm{\theta}_{\text{SF}}$) takes the following form

$\Delta\bm{\theta}_{\text{SF}}=(\hat{\mathbf{F}}(1)+\lambda\mathbf{I})^{-1}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}).$
(47)

By plugging in
$\hat{\mathbf{F}}(1)=\hat{\mathbf{J}}(1)^{\top}\hat{\mathbf{J}}(1)$ and using
the SMW identity, the update can then be computed as follows

$\Delta\bm{\theta}_{\text{SF}}=\frac{1}{\lambda}[\mathbf{I}-\hat{\mathbf{J}}(1)^{\top}(\hat{\mathbf{J}}(1)\hat{\mathbf{J}}(1)^{\top}+\lambda\mathbf{I})^{-1}\hat{\mathbf{J}}(1)]\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}),$
(48)

which is easy to compute once the Jacobian $(\hat{\mathbf{J}}(1)$ for sampled
gradients is collected (using the per-sample gradient collection method
described in Appendix D.1). Similar to EF and iEF, an exact stochastic
optimiser can also be constructed for SF, as is described in Algorithm 3. Note
that as compared to Algorithm 2,1, SF requires an additional back-propagation.

Algorithm 3 Stochastic Optimisation with Exact SF

Require: All $N$ training samples, initial model parameters $\bm{\theta}(0)$

for $t=0$ to $T-1$ do

Sample Mini-batch $\mathcal{M}(t)$ with size $M$

Perform forward pass on batch $\mathcal{M}(t)$ to obtain output distribution
$\bm{p}(t)$

Perform back-propagation to obtain accumulated loss
$\nabla_{\bm{\theta}(t)}\mathcal{L}(\bm{\theta}(t))$

Sample output labels $\hat{\bm{y}}(t)\sim\bm{p}(t)$

Compute cross-entropy loss between $\bm{p}(t)$ and sampled label
$\hat{\bm{y}}(t)$ to obtain the pseudo per-sample loss vector
$\hat{\bm{l}}(t)$

Perform additional back-propagation to obtain Jacobian
$\mathbf{A}(t)=\nabla_{\bm{\theta}(t)}\hat{\bm{l}}(t)$

Compute SF update
$\Delta\bm{\theta}_{\text{SF}}(t)=\frac{1}{\lambda}[\mathbf{I}-\mathbf{A}(t)^{\top}(\mathbf{A}(t)\mathbf{A}(t)^{\top}+\lambda\mathbf{I})^{-1}\mathbf{A}(t)]\nabla_{\bm{\theta}(t)}\mathcal{L}(\bm{\theta}(t))$:
Eqn. (48)

Update model
$\bm{\theta}(t+1)=\bm{\theta}(t)-\eta\Delta\bm{\theta}_{\text{SF}}(t)$

end for

##### Impact of Inversely-Scaled Projection on SF Update

Although the same analysis in Sec. 4 cannot be applied to SF due to the
different formulation of the update, the inversely-scaled projection issue is
still expected to impact SF. The sampled Fisher matrix $\hat{\mathbf{F}}(1)$
of SF is likely to include an un-scaled empirical gradient, and pre-
conditioning the gradient with the inverse of this matrix is likely to cause
the inverse-scaling issue to some extent (although to a smaller extent than
EF). Consequently, a poorer approximation quality of SF to the exact NG update
is still expected.

##  Appendix F Empirical Evaluation Framework for Approximate NGD Methods

###  F.1 Implementation Details

As introduced in Sec. 6, the proposed empirical evaluation framework is built
around the proposed indicator $\gamma(\cdot)$. Consider a given model
checkpoint $\bm{\theta}(t)$ and $N$ pairs of training samples
$(\bm{x}_{n},y_{n})_{n=1}^{N}$ from a target dataset, $K$ approximate NGD
methods $\bm{g}^{(k)}(\cdot)$ are to be compared. Each method can generate an
update $\Delta\bm{\theta}(t)^{(k)}$ for the provided samples and model as
$\Delta\bm{\theta}(t)^{(k)}=\bm{g}^{(k)}(\bm{\theta}(t),(\bm{x}_{n},y_{n})_{n=1}^{N})$.
The framework is implemented such that $K$ indicators
$\gamma(\Delta\bm{\theta}(t)^{(k)})$ are computed, which allows the ranking of
these update generation methods w.r.t. their approximation quality to the
exact NGD method. This can be repeated for model checkpoints at different
training stages $\bm{\theta}_{t^{{}^{\prime}}}$ and different batches of
training samples to provide a more accurate comparison. The detailed
implementation of the proposed evaluation framework for a given checkpoint
$\bm{\theta}(t)$, a given batch of $N$ samples and $M$ update generation
methods is described in Algorithm 4. Typically, multiple batches are used to
compute a series of $\gamma(\Delta\bm{\theta}^{(k)})$, and the averaged
indicator is used to provide a more accurate value.

In particular, assuming the target approximate NGD method generates a pre-
conditioned gradient, then it is important to compute the ratio
$\gamma(\Delta\bm{\theta}(t)^{(k)})/\gamma(\nabla_{\bm{\theta}(t)}\mathcal{L}(\bm{\theta}(t)))$.
This ratio should be smaller than 1 for a good approximate NGD method. If this
ratio is consistently larger than 1, then it means that the target approximate
NGD method $\bm{g}^{(k)}(\cdot)$ has a poorer approximation quality than the
SGD method and is likely to be unsafe to use in practice.

The four update methods investigated in the experiments (see Sec. 7) are SGD,
EF, iEF and SF (sampled Fisher with 1 Monte-Carlo sample for each training
datum). The SGD update is directly the gradient
$\nabla_{\bm{\theta}(t)}\mathcal{L}(\bm{\theta}(t))$; the iEF update is
computed according to Eqn. (8); the EF update is computed according to Eqn.
(5); the sampled Fisher update is computed according to Eqn. (48).

Finally, the computationally most expensive operation when computing indicator
is finding the matrix-vector product between the exact Fisher and a given
update direction, i.e. $\mathbf{F}\Delta\bm{\theta}$. A detailed discussion is
provided on its implementation and algorithm complexity in the following
section.

Algorithm 4 Empirical evaluation framework for approximate NGD Methods

Input: A batch of $M$ training data pairs $(\bm{x}_{n},y_{n})_{n=1}^{M}$,
model checkpoint $\bm{\theta}(t)$, loss function
$\mathcal{L}(\bm{\theta}(t))$, $K$ update generation methods
$\bm{g}^{(k)}(\bm{\theta}(t),(\bm{x}_{n},y_{n})_{n=1}^{M}),k\in\\{1,\cdots,K\\}$.

Execute:

Compute accumulated gradient $\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}(t))$
on $M$ training samples

for $k=1$ to $K$ do

Compute update
$\Delta\bm{\theta}(t)^{(k)}=\bm{g}^{(k)}(\bm{\theta}(t),(\bm{x}_{n},y_{n})_{n=1}^{M})$

Compute indicator
$\gamma(\Delta\bm{\theta}(t)^{(k)})=\frac{[(\Delta\bm{\theta}(t)^{(k)})^{\top}\mathbf{F}\Delta\bm{\theta}(t)^{(k)}]^{\frac{1}{2}}}{(\Delta\bm{\theta}(t)^{(k)})^{\top}\nabla_{\bm{\theta}(t)}\mathcal{L}(\bm{\theta}(t))}$

end for

Output: Return $M$ indicators
$\left[\gamma(\Delta\bm{\theta}(t)^{(k)})\right]_{k=1}^{K}$. Updates with a
smaller indicator are considered a better approximate NG update.

####  F.1.1 Matrix-Vector Product and Algorithm Complexity

The most important operation in Algorithm 4 when computing the proposed
indicator is the computation of the matrix-vector product $\mathbf{F}(\cdot)$.
It is the most expensive operation that dominates both the time and space
complexity of this algorithm. Fortunately, this can be efficiently computed in
a modern auto-grad framework using the double differentiation trick [40]. The
key to computing the matrix-vector product with the exact Fisher matrix is
identifying the equivalence of the Fisher matrix and the Generalised Gauss-
Newton matrix (GGN matrix, denoted as $\mathbf{G}$). It can be shown that for
the target model which uses softmax activation and cross-entropy loss
function, the following equality holds

$\mathbf{F}=\sum\nolimits_{n}\nabla_{\bm{\theta}}\bm{z}_{n}^{\top}(\nabla^{2}_{\bm{z}_{n}}l_{n})\nabla_{\bm{\theta}}\bm{z}_{n}=\mathbf{G}.$
(49)

Therefore, the matrix-vector product can be broken down into the computation
of three separate matrix-vector products. The double differentiation trick is
useful when computing the first two matrix-vector products:
$\nabla_{\bm{\theta}}\bm{z}_{n}(\cdot)$,
$\nabla^{2}_{\bm{z}_{n}}l_{n}(\cdot)$. The final matrix-vector product
$\nabla_{\bm{\theta}}\bm{z}_{n}^{\top}(\cdot)$ can be computed using the
standard back-propagation pipeline. Please refer to [40] for implementation
details.

Overall, with the double differentiation trick, the time and memory complexity
of the matrix-vector product with the Fisher matrix is comparable to standard
training with batch size $M$. In our implementation in Pytorch, the matrix-
vector product requires 2 forward passes and 2 backward passes through the
model on the $M$ samples. This makes the proposed empirical evaluation
framework particularly manageable.

Finally, for the reader’s information, by rewriting the code according to
[23], it is possible to further reduce the cost of the matrix-vector product
down to only 1 additional forward pass and backward pass. However, this is not
a preferable choice because it requires rewriting the code individually for
each target model.

####  F.1.2 Comparison to Traditional Evaluation Methods

As is summarised in Sec. 2, traditional evaluation of the approximation
quality to the NGD method requires the explicit storage and inversion of the
exact Fisher matrix. This is usually done using the definition in Eqn. (2).
For a general model, computing the Fisher matrix in Eqn. (2) requires $C$
separate backwards and forward passes through the model ($C$ being the output
category number). This can range from 10 (e.g. MNIST[19]) to >10,000 (for our
LLM finetune setup [7]). From a time complexity perspective, the traditional
method can be arbitrarily more expensive than our evaluation framework
depending on the output category size $C$. From a memory complexity
perspective, the storage of the Fisher matrix of $\mathbb{R}^{P\times P}$ is
infeasible for a large-scale model. Both of these concerns limit the
traditional evaluation of approximate NGD methods in large practical setups.
On the contrary, our proposed evaluation framework resolves these limitations
of the traditional evaluation method, meanwhile having a strong theoretical
motivation. This greatly facilitates the evaluation of approximate NGD methods
in large-scale deep-learning setups for methods not limited to EF and iEF.

###  F.2 Hessian-free Implicitly Minimises $\gamma(\cdot)$ Indicator

In Sec. 6.2, it is stated that the indicator $\gamma(\cdot)$ is implicitly
minimised in the linear CG algorithm to iteratively approximate the exact NG
update in Hessian-free method [23]. In this section, this statement is
formally introduced and justified.

The Hessian-free method is an important approximate NGD method. Unlike most
other approximate NGD methods, which approximate the Fisher matrix
$\mathbf{F}$ directly, the Hessian-free method does not explicitly express the
Fisher matrix. Instead, it uses an iterative method to directly solve for
$\Delta\bm{\theta}_{\text{HF}}$ in the following equation

$\mathbf{F}\Delta\bm{\theta}_{\text{HF}}=\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}).$
(50)

The iterative method used in Hessian-free is the linear CG algorithm [41],
which is a classical method that iteratively solves for
$\Delta\bm{\theta}_{\text{HF}}$ in Eqn. (50) with only the access to the
matrix-vector product function $\mathbf{F}(\cdot)$ (introduced in Sec. 6).
Since the Fisher is not required to be explicitly stored, the method enjoys a
memory complexity of $O(P)$ where $P$ is the number of trainable parameters in
a target model.

The linear CG algorithm (without pre-conditioning) used in HF is shown in
Algorithm 5.

Input: The maximum CG execution iterations $M_{\text{CG}}$

The gradient vector $\bm{b}=\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})$

The matrix-vector product function for the Fisher $\mathbf{F}$

Execute:

Set $\bm{r}_{0}\leftarrow\bm{b}$

Set $\bm{v}_{0}\leftarrow\bm{r}_{0}$

Set $m\leftarrow 0$

while $m<M_{\text{CG}}$ do

Compute $\|\bm{r}_{m}\|^{2}=\bm{r}_{m}^{\top}\bm{r}_{m}$

Set
$\alpha_{m}\leftarrow\|\bm{r}_{m}\|^{2}/(\bm{v}_{m}^{\top}\mathbf{F}\bm{v}_{m})$

Update $\bm{x}_{m+1}\leftarrow\bm{x}_{m}+\alpha_{m}\bm{v}_{m}$

Update $\bm{r}_{m+1}\leftarrow\bm{r}_{m}-\alpha_{m}\mathbf{F}\bm{v}_{m}$

Compute $\|\bm{r}_{m+1}\|^{2}=\bm{r}_{m+1}^{\top}\bm{r}_{m+1}$

Set $\beta_{m+1}\leftarrow\|\bm{r}_{m+1}\|^{2}/\|\bm{r}_{m}\|^{2}$

Update $\bm{v}_{m+1}\leftarrow\bm{r}_{m+1}+\beta_{m+1}\bm{v}_{m}$

Update $m\leftarrow m+1$

end while

Output: Return $\bm{x}_{M_{\text{CG}}}$ as an approximate solution for
$\Delta\bm{\theta}_{\text{EF}}$

Algorithm 5 The Linear CG Algorithm in Hessian-free Method

It is known that the CG algorithm is an efficient solver for Eqn. (50) and is
usually regarded as a locally optimal minimiser (each step uses an exact line
search) for the equivalent pseudo loss function [23, 41]

$L^{{}^{\prime}}_{\text{CG}}(\Delta\bm{\theta}_{\text{HF}})=-\Delta\bm{\theta}_{\text{HF}}^{\top}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})+\frac{1}{2}\Delta\bm{\theta}_{\text{HF}}^{\top}\mathbf{F}\Delta\bm{\theta}_{\text{HF}}.$
(51)

However, it is also possible to interpret this method as a locally optimal
minimiser [16] for the generalised Rayleigh Quotient of positive semi-definite
matrices $\mathbf{F}$ and
$\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})^{\top}$

$\gamma(\Delta\bm{\theta}_{\text{HF}})^{2}=\frac{(\Delta\bm{\theta}_{\text{HF}}^{\top}\mathbf{F}\Delta\bm{\theta}_{\text{HF}})}{\Delta\bm{\theta}_{\text{HF}}^{\top}[\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})^{\top}]\Delta\bm{\theta}_{\text{HF}}}$
(52)

where $\gamma(\Delta\bm{\theta}_{\text{HF}})$ happens to be the proposed
indicator of this paper. The proof of this statement is provided in the
following section.

####  F.2.1 Proof that Linear CG is a Locally Optimal Minimiser for Indicator
$\gamma(\cdot)$

Note that in every iteration in Algorithm 5, an update $\alpha_{m}\bm{v}_{m}$
is accumulated to the final solution $\bm{x}_{M_{\text{CG}}}$. This update
$\alpha_{m}\bm{v}_{m}$ can be shown to achieve the local maximum reduction in
$\gamma(\bm{x})^{2}$ in every iteration. Formally, it is to be shown that, for
the current partial solution $\bm{x}_{m}$ and the given search direction
$\bm{v}_{m}$, the scaling factor $\alpha_{m}$ achieves the maximum reduction
in $\gamma(\bm{x}_{M_{\text{CG}}})^{2}$:

$\alpha_{m}=\operatorname*{arg\,min}_{\alpha_{m}}\gamma(\bm{x}_{M_{\text{CG}}})^{2}.$
(53)

Now, assume the actual minimiser $\alpha^{{}^{\prime}}_{m}$ is different from
the $\alpha_{m}$ in the CG iterations. The true minimiser
$\alpha^{{}^{\prime}}_{m}$ can be acquired through a Ritz-Rayleigh analysis
[16]. By setting
$\mathbf{B}^{{}^{\prime}}=\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})^{\top}$,
and plugging in the true minimiser $\alpha^{{}^{\prime}}_{m}$, Eqn. (53)
becomes:

$\alpha^{{}^{\prime}}_{m}=\operatorname*{arg\,min}_{\alpha^{{}^{\prime}}_{m}}\frac{[(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})^{\top}\mathbf{F}(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})]}{[(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})^{\top}\mathbf{B}^{{}^{\prime}}(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})]}.$
(54)

The Ritz-Rayleigh analysis then gives the optimal scaling:

$\alpha^{{}^{\prime}}_{m}=\frac{\bm{b}^{\top}\bm{v}_{m}}{\bm{v}_{m}^{\top}\mathbf{F}\bm{v}_{m}}.$
(55)

Recall that
$\alpha_{m}=\frac{\bm{r}_{m}^{\top}\bm{r}_{m}}{\bm{v}_{m}^{\top}\mathbf{F}\bm{v}_{m}}$,
in order to prove $\alpha_{m}=\alpha^{{}^{\prime}}_{m}$, the following need be
proved $\bm{b}^{\top}\bm{v}_{m}=\bm{r}_{m}^{\top}\bm{r}_{m}$. This can be done
through recursion.

For $m=0$, it is obvious that
$\bm{b}^{\top}\bm{v}_{0}=\bm{b}^{\top}\bm{b}=\bm{r}_{0}^{\top}\bm{r}_{0}$.

For $m>1$, it is known that
$\bm{b}^{\top}\bm{v}_{m-1}=\bm{r}_{m-1}^{\top}\bm{r}_{m-1}$ is true, then:

$\begin{aligned}
\bm{b}^{\top}\bm{v}_{m}&=\bm{b}^{\top}(\bm{r}_{m}+\beta_{m}\bm{v}_{m-1})\\\
&=\bm{b}^{\top}\bm{r}_{m}+\frac{\bm{r}_{m}^{\top}\bm{r}_{m}^{\top}}{\bm{r}_{m-1}^{\top}\bm{r}_{m-1}}\bm{b}^{\top}\bm{v}_{m-1}\\\
&=\bm{r}_{0}^{\top}\bm{r}_{m}+\frac{\bm{r}_{m}^{\top}\bm{r}_{m}^{\top}}{\bm{r}_{m-1}^{\top}\bm{r}_{m-1}}\bm{r}_{m-1}^{\top}\bm{r}_{m-1}\\\
&=0+\bm{r}_{m}^{\top}\bm{r}_{m}\ \ \ (\because\forall i\neq
j,\bm{r}_{i}^{\top}\bm{r}_{j}=0)\\\
&=\bm{r}_{m}^{\top}\bm{r}_{m}\end{aligned}.$ (56)

In conclusion, $\alpha_{m}$ is indeed the true minimiser for
$\frac{[(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})^{\top}\mathbf{F}(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})]}{[(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})^{\top}\mathbf{B}^{{}^{\prime}}(\bm{x}_{m}+\alpha^{{}^{\prime}}_{m}\bm{v})]}$
in every CG iterations, making CG a locally optimal optimiser for
$\gamma(\bm{x})^{2}$.

It is known that Hessian-free iteratively approaches the exact NG update [23],
and a larger iteration number ${M_{\text{CG}}}$ leads to a better
approximation. Now that it has been shown that the indicator $\gamma(\bm{x})$
decreases (optimally) for every iteration of Hessian-free, it verifies that
indicator $\gamma(\bm{x})$ directly describes the approximation level of the
intermediate solution $\bm{x}_{m}$ to the exact NG update.

###  F.3 $\gamma(\cdot)$ Quantifies Loss Reduction Effectiveness under the
Local Quadratic Approximation

The indicator $\gamma(\Delta\bm{\theta})$ can also describe the loss-reduction
effectiveness of a target update direction under the Local Quadratic
Approximation. The loss change induced by a target update can be estimated
using a Taylor expansion up to the second order as follows:

$\mathcal{L}(\bm{\theta}({t+1}))\approx\mathcal{L}(\bm{\theta}(t))+\Delta\bm{\theta}(t)^{\top}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta}(t))+\frac{1}{2}\Delta\bm{\theta}(t)^{\top}\mathbf{B}_{t}\Delta\bm{\theta}(t)$
(57)

where $\mathbf{B}\in\mathbb{R}^{P\times P}$ is the choice of curvature matrix.
Such an approximation to the loss function is regarded as a Local Quadratic
Approximation [50].

Consider a fixed-direction target update $\Delta\bm{\theta}$ with an unknown
step size $\eta$, the maximum achievable loss reduction based on Local
Quadratic Approximation can be obtained as follows

$-\frac{1}{2\gamma(\Delta\bm{\theta})^{2}}=\min_{\eta}\eta\Delta\bm{\theta}^{\top}\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})+\frac{\eta^{2}}{2}\Delta\bm{\theta}^{\top}\mathbf{F}\Delta\bm{\theta},$
(58)

where $\mathbf{F}$ is used in place of the curvature matrix $\mathbf{B}$. This
implies that the maximum loss reduction along update direction
$\Delta\bm{\theta}$ is $-{1}/{2\gamma(\Delta\bm{\theta})^{2}}$, which is
inversely proportional to the square of the proposed indicator. For an update
with a smaller $\gamma(\Delta\bm{\theta})$, it is expected to induce a larger
loss reduction, with the exact NG update achieving the maximum loss reduction.
Consequently, the proposed indicator has a strong correlation to practical
convergence ability of the target update generation method.

##  Appendix G Limitations

##### Focus on Theoretical Approximation for NGD Method

The focus of the paper is on the theoretical approximation methods for NGD,
including the exact EF method, the proposed iEF method, and also the SF
method. Practical approximate (empirical) NGD optimisers are not the main
focus of this paper (e.g. K-FAC [25], EK-FAC [11], TNT [37] etc.), and no
experiments are conducted for them. Despite the importance of these practical
optimisers in the scene of NGD-based optimisation, they can all be considered
a further (structure-aware) approximation to the EF or SF methods. However, it
is an important future work to apply iEF to improve existing approximate NGD
optimisers (an example of using iEF to improve empirical K-FAC is provided in
Appendix D.2).

##### Limitations of Exact Update Generation

The updates generated by EF, iEF and SF methods (defined in Eqn. (5), (8),
(48) respectively) investigated in this paper are “exact”, in order to provide
a more theoretically rigorous comparison and analysis. In our experiments,
they are generated based only on the current provided batch $\mathcal{M}(t)$
as is described in Algorithms 2, 1, 3 respectively. The implementation of
exact EF, iEF and SF updates requires storage of the $M$ per-sample gradients,
which is memory-demanding in practice. This limits the scope of application of
the exact iEF (and EF/SF) optimisers to setups where trainable parameter size
is small, such as PEFT for pre-trained models. However, given the rise of PEFT
setups, and consider the competitive optimisation performance of the exact iEF
optimiser, such limitations may be out-weighted by the gain in practice.
Additionally, the exact update formulation of iEF means momentum and weight
decay cannot be directly applied with the resultant optimiser. This affects
the optimisation performance on certain setups, and further work is required
to integrate these key techniques with exact iEF optimiser. However, none of
these limitations would affect the future work where iEF is incorporated into
practical approximate NGD optimisers, where memory constraints and
momentum/weight decay integration is already resolved in these practical NGD
optimisers.

##  Appendix H Detailed Experiment Information

###  H.1 Overall Experimental Setup

##### Textual Classification of GLUE

We have investigated 7 selected GLUE benchmark tasks [45] including CoLA,
SST-2, MRPC, QQP, MNLI, QNLI, and RTE, which together cover a range of NLP
classification tasks such as sentiment classification, semantic equivalence
checking, grammar checking, entailment prediction etc. The same selection of
GLUE tasks is used in [7]. One key aspect of GLUE tasks is that they do not
have test labels, and for test evaluations, they have to be submitted to their
official website (with a maximum allowed frequency of 3 submits per day). The
textual label tokens are to be directly predicted by the target model. Also,
an instruction is prepended to the input of each training sample. See details
in Table. 2. For each task, a pre-trained T5-base model [35] is trained with
two parameter-efficient finetuning methods: Prompt Tuning [20] and LoRA [14].

Table 2: Training details for the 7 GLUE tasks. The input format column
describes how the inputs are structured before being sent to the model. {S1},
{S2} represents the default input sentence(s) of the corresponding tasks, some
tasks have only one input sentence and some have two. Labels are the tokens to
be predicted by the model for each task.

Task | Input Format | Labels
---|---|---
CoLA |  Classify if the following sentence’s grammar is acceptable or unacceptable: {S1} |  acceptable, unacceptable
SST-2 |  Classify if the following sentence’s sentiment is positive or negative: {S1} |  positive, negative
QQP |  Classify if the following Q1 and Q2 are semantically equivalent, answer yes or no: Q1: {S1} Q2: {S2} |  yes, no
MRPC |  Classify if the following S1 and S2 are semantically equivalent, answer yes or no: S1: {S1} S2: {S2} |  yes, no
MNLI |  Predict whether the premise entails the hypothesis, contradicts the hypothesis,or neither, answer yes, no or maybe: premise: {S1} hypothesis: {S2} |  yes, no, maybe
QNLI |  Determine whether the context sentence S contains the answer to the question Q, answer yes or no: Q: {S1} S: {S2} |  yes, no
RTE |  Classify if S1 entailment S2 or not, answer yes or no: S1: {S1} S2: {S2} |  yes, no

##### Computer Vision Classification with CIFAR100

The well-known CIFAR100 dataset [17] is used to finetune the pretrained ViT
model [8]. The model is finetuned with LoRA [14] following the setup in [13].

###  H.2 Optimisation Experimental Setup

For Prompt Tuning, the Adafactor [21] optimiser is used as baseline (as is
done in [21, 7]). For all other training setups, the AdamW [22] optimiser was
used as the baseline.

Different tasks were trained with a different number of epochs. The validation
accuracy (on the dev set of each task) of the best checkpoint is reported and
sent for test evaluation (for GLUE, the checkpoints are submitted to GLUE
website [45]). Details are shown in Table 3.

Table 3: Optimisation details all involved tasks. Train epochs represent the
number of epochs for which the model is trained. Evaluation frequency
describes the number of update steps between each validation evaluation on the
development set.

Task | Train Epochs | Evaluation Frequency
---|---|---
CoLA | 20 | 100
SST-2 | 20 | 1000
QQP | 5 | 1000
MRPC | 30 | 50
MNLI | 5 | 1000
QNLI | 15 | 1000
RTE | 40 | 30
CIFAR100 | 5 | 1000

For all the GLUE tasks and for each finetuning method (Prompt Tuning or LoRA),
a single set of hyper-parameters were searched for each optimiser and is used
across runs on all 7 GLUE tasks. Three runs with different seeds were
conducted to generate the error bar. For each run, the checkpoint with the
highest validation accuracy was saved for later evaluation (to get test
performance or to be evaluated in the proposed evaluation framework). All
optimisers were trained on a batch size of 32.

Prompt Tuning + T5-base on GLUE:  20 prompts are used for Prompt Tuning. The
trainable parameter size is $15,360$, taking up $0.0069\%$ of the total
parameter size (222,918,912) of the T5-base model. Constant scheduling was
used for Adafactor, SGD and iEF runs. For iEF, EF and SF, a damping of
$\lambda=$1\text{\times}{10}^{-12}$$ was used. For the Adafactor baseline
optimiser, the hyper-parameter provided in [7] was used (which comes from
[21]): weight-decay $1\text{\times}{10}^{-5}$, $\beta_{2}=0.8$, learning rate
$\eta=0.3$ and no parameter scaling. For the SGD method, the learning rate is
$\eta=100$, which was searched from $\\{0.1,1,10,20,50,100\\}$. For the iEF
method, the learning rate was $\eta=50$, which was searched from
$\\{1,10,50,100\\}$. For the EF method, a different scheduling of learning
rate was used to guarantee convergence, due to the inverse scaling of EF
updates. The chosen strategy was a linearly decaying normalised update, with
the first update being normalised to $1$ ({$1\text{\times}{10}^{-3}$,
$5\text{\times}{10}^{-3}$, $1\text{\times}{10}^{-2}$,
$1\text{\times}{10}^{-1}$, $1$, $10$}) and linearly decaying to 0. The SF
method was trained using the same method as EF with the same set of
hyperparameters.

LoRA + T5-base on GLUE:  The LoRA was set to have rank 8 with dropout 0.1
[13]. The trainable parameter size is $884,736$, taking up $0.40\%$ of the
total parameter size (222,918,912) of the T5-base model. Constant scheduling
was used for the AdamW, SGD and iEF runs. For iEF, EF and SF, a damping of
$\lambda=$1\text{\times}{10}^{-7}$$ was used (it is 5 orders of magnitude
larger than that used in Prompt Tuning because the diagonal of the gradient
covariance matrix has a 5 order of magnitude larger norm in LoRA than Prompt
Tuning). For the AdamW baseline optimiser, the hyper-parameters are: weight-
decay $1\text{\times}{10}^{-2}$, and the learning rate of
$1\text{\times}{10}^{-3}$ was searched from {$1\text{\times}{10}^{-3}$,
$5\text{\times}{10}^{-4}$, $1\text{\times}{10}^{-4}$}. For the SGD method, the
learning rate was $\eta=0.1$, which was searched from
$\\{0.1,1,10,20,50,100\\}$. For the iEF method, the learning rate was
$\eta=100$, which was searched from $\\{1,10,50,100\\}$. For the EF method, a
normalised update with linear scheduling was used similar to Prompt Tuning,
and a starting learning rate of 0.01 was searched from
{$1\text{\times}{10}^{-3}$, $5\text{\times}{10}^{-3}$,
$1\text{\times}{10}^{-2}$, $1\text{\times}{10}^{-1}$, $1$}. The SF method was
trained with the same strategy and hyperparameters.

LoRA + ViT on CIFAR100 The setup of using LoRA to finetune CIFAR100 in [13]
was used. The trainable parameter size is 313,344, taking up $0.36\%$ of the
total parameter size (86,862,568) of the ViT model (vit-base-patch16-224). The
same hyperparameters for LoRA + T5 was used for these experiments.

###  H.3 Experiments Compute Resources

All the optimisation experiments and evaluation experiments are run on a cloud
linux machine with 8 A100 GPUs with 80GB GRAM. For all the optimisation
experiments, optimisation time with different optimisers are similar, apart
from the SF optimiser where the additional back-propagation leads to an
additional 60% runtime. For standard SGD/AdamW/Adafactor/EF/iEF optimisers, on
average each complete run takes 10 hours. The slowest task (QQP + LoRA) takes
20 hours and the quickest task (RTE + Prompt Tuning) takes 0.3 hours. In total
420 GPU hours are run for all optimisation experiments. For evaluation runs,
each evaluation is done on 100 batches of 160 training samples. For each
checkpoint and a choice of damping, the evaluation takes on average 30
minutes. For all the evaluation runs (damping evaluated 5 tasks x 3 ckpts x 10
damping = 150 evaluations, standard evaluated 15 tasks x 7 ckpts = 115
evaluated, total 265 evaluation), 133 GPU hours are done. Considering hyper-
parameter tuning, and various preliminary runs to make exact EF/iEF/SF
optimisers and the evaluation framework to run properly, in total 2 times of
additional compute time is required for all experiments.

###  H.4 Licenses for All Used Assets

The license, citation and link to various asset used in this paper is provided
in Table 4.

Table 4: The license, citation and link to various asset used in this paper. Asset Name | Citation | URL | License
---|---|---|---
T5-base | [35] | Model checkpoint | Apache-2.0
ViT-B/16 | [8] | Model checkpoint | Apache-2.0
GLUE | [45] | Website | Multiple
CIFAR100 | [17] | Website | Unknown
Pytorch | [28] | Website | Multiple

###  H.5 Further Evaluation Plots

Due to space limit and presentation concerns, the evaluation results presented
in the main text in experiments (E1) and (E3) are only partial. In this
section, additional evaluation results are presented.

####  H.5.1 Approximation Quality across Tasks and Training Stages

The same indicator evaluation in Fig. 2 is done for all experimented tasks. A
full view of the indicator plots are provided in Fig. 4 and 5. Similar trend
can be found across tasks.

Figure 4: Four (log-scaled) ratios computed for checkpoints at various stages
of training (sampled at the interval of one epoch) for all LoRA tasks,
including T5 + 7 GLUE tasks and ViT + 1 CIFAR100. The figure is drawn in the
same fashion as Fig. 2. Note that the error bar is not presented to improve
presentation clarity.

Figure 5: Four (log-scaled) ratios computed for checkpoints at various stages
of training (sampled at the interval of one epoch) for 7 Prompt Tuning tasks
for T5 + GLUE. The figure is drawn in the same fashion as Fig. 2. Note that
the error bar is not presented to improve presentation clarity.

####  H.5.2 Impact of Damping on Approximation Quality

The same damping analysis in Fig. 3 is applied to several other tasks
(SST2+T5+Prompt Tuning, MRPC+T5+Prompt Tuning, CIFAR100+ViT+LoRA,
RTE+T5+LoRA,as shown in Fig. 6, 7, 8, 9 respectively). Similar trend can be
found across tasks.

Figure 6: Approximation quality (relative to SGD) of EF, SF and iEF methods
w.r.t. damping factor $\lambda$ at different training stages of task
SST2+T5+Prompt Tuning. The figure is drawn in the same fashion as Fig. 3.

Figure 7: Approximation quality (relative to SGD) of EF, SF and iEF methods
w.r.t. damping factor $\lambda$ at different training stages of task
MRPC+T5+Prompt Tuning. The figure is drawn in the same fashion as Fig. 3.

Figure 8: Approximation quality (relative to SGD) of EF, SF and iEF methods
w.r.t. damping factor $\lambda$ at different training stages of task
CIFAR100+ViT+LoRA. The figure is drawn in the same fashion as Fig. 3.

Figure 9: Approximation quality (relative to SGD) of EF, SF and iEF methods
w.r.t. damping factor $\lambda$ at different training stages of task
RTE+T5+LoRA. The figure is drawn in the same fashion as Fig. 3.

###  H.6 Full Optimisation Results

Due to space limit, only a partial test performance table (Table 1) is
presented for the experiment (E2) in the main text. In this section, the full
training, validation and test results for all combinations of structure, tasks
and optimisers are reported for the reader’s information.

Training curves and validation accuracy curves of different optimisers for the
described 15 setups are presented in Fig. 10 and 11. The final training loss
for every task, structure and optimiser combination is presented in Table 5.
The validation performance for every task, structure and optimiser combination
is presented in Table 6. The test performance for all tasks is presented in
Table 7.

Table 5: Final train loss for all the task, structure and optimiser combinations. The average train loss of the final 100 steps of 3 random seed runs is reported, along with the standard deviation. IEF achieves the lowest final training loss for 12 out of 15 tasks. | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | CIFAR100
---|---|---|---|---|---|---|---|---
Prompt Tuning
Adafactor | $0.07_{\pm{0.08}}$ | $0.07_{\pm{0.02}}$ | $0.16_{\pm{0.07}}$ | $\bm{0.02}_{\pm{0.02}}$ | $0.19_{\pm{0.02}}$ | $0.10_{\pm{0.05}}$ | $0.29_{\pm{0.03}}$ | -
SGD | $0.33_{\pm{0.06}}$ | $0.04_{\pm{0.02}}$ | $0.36_{\pm{0.11}}$ | $0.15_{\pm{0.13}}$ | $0.31_{\pm{0.05}}$ | $0.15_{\pm{0.06}}$ | $0.35_{\pm{0.03}}$ | -
EF | $2.28_{\pm{1.11}}$ | $0.99_{\pm{0.35}}$ | $2.80_{\pm{1.19}}$ | $4.29_{\pm{1.04}}$ | $6.11_{\pm{1.10}}$ | $5.12_{\pm{0.88}}$ | $1.79_{\pm{0.16}}$ | -
SF | $0.16_{\pm{0.10}}$ | $0.06_{\pm{0.11}}$ | $0.43_{\pm{0.09}}$ | $0.15_{\pm{0.02}}$ | $0.28_{\pm{0.02}}$ | $0.27_{\pm{0.14}}$ | $0.36_{\pm{0.03}}$ | -
iEF | $\bm{0.01}_{\pm{0.01}}$ | $\bm{0.02}_{\pm{0.01}}$ | $\bm{0.04}_{\pm{0.01}}$ | $\bm{0.02}_{\pm{0.01}}$ | $\bm{0.17}_{\pm{0.07}}$ | $\bm{0.09}_{\pm{0.06}}$ | $\bm{0.02}_{\pm{0.01}}$ | -
LoRA
AdamW | $\bm{0.10}_{\pm{0.03}}$ | $\bm{0.02}_{\pm{0.01}}$ | $0.07_{\pm{0.06}}$ | $\bm{0.07}_{\pm{0.01}}$ | $\bm{0.22}_{\pm{0.06}}$ | $0.10_{\pm{0.03}}$ | $\bm{0.02}_{\pm{0.03}}$ | $0.60_{\pm{0.21}}$
SGD | $0.24_{\pm{0.02}}$ | $0.03_{\pm{0.09}}$ | $0.12_{\pm{0.01}}$ | $0.10_{\pm{0.05}}$ | $\bm{0.22}_{\pm{0.12}}$ | $0.11_{\pm{0.07}}$ | $0.18_{\pm{0.07}}$ | $0.94_{\pm{0.34}}$
EF | $1.92_{\pm{0.49}}$ | $0.90_{\pm{0.03}}$ | $2.541_{\pm{0.78}}$ | $1.03_{\pm{0.51}}$ | $3.42_{\pm{0.33}}$ | $1.31_{\pm{0.89}}$ | $3.08_{\pm{0.40}}$ | $3.10_{\pm{0.11}}$
SF | $0.26_{\pm{0.03}}$ | $0.14_{\pm{0.22}}$ | $0.35_{\pm{0.11}}$ | $0.20_{\pm{0.12}}$ | $0.24_{\pm{0.09}}$ | $0.30_{\pm{0.21}}$ | $0.39_{\pm{0.02}}$ | $0.63_{\pm{0.12}}$
iEF | $0.16_{\pm{0.04}}$ | $\bm{0.02}_{\pm{0.02}}$ | $\bm{0.06}_{\pm{0.04}}$ | $\bm{0.07}_{\pm{0.05}}$ | $0.36_{\pm{0.05}}$ | $\bm{0.08}_{\pm{0.06}}$ | $0.05_{\pm{0.02}}$ | $\bm{0.59}_{\pm{0.30}}$
Table 6: Validation metrics (multiplied by 100) (on development set) for all the task, structure and optimiser combinations. The average of the highest validation accuracy of 3 random seed runs is reported, along with the standard deviation. Task-specific metrics are reported in this table. For SST-2, QNLI, RTE and CIFAR100, accuracy is reported. For CoLA, both accuracy and Matthew’s Corr is reported. For MRPC and QQP, F1-score and Accuracy (in order) are reported. For MNLI, average of accuracy for entire development set is reported. | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | CIFAR100
---|---|---|---|---|---|---|---|---
Prompt Tuning
Adafactor |  $\bm{82.0}_{\pm{0.44}}$ $\bm{53.8}_{\pm{1.01}}$ | $94.2_{\pm{0.25}}$ |  $84.8_{\pm{0.35}}$ $88.0_{\pm{0.73}}$ |  $\bm{90.7}_{\pm{0.01}}$ $\bm{87.7}_{\pm{0.01}}$ | $82.4_{\pm{0.25}}$ | $\bm{91.9}_{\pm{0.09}}$ | $64.7_{\pm{0.34}}$ | -
SGD |  $69.1_{\pm{0.09}}$ $-0.1_{\pm{2.76}}$ | $93.3_{\pm{0.14}}$ |  $70.0_{\pm{0.31}}$ $\bm{80.8}_{\pm{0.14}}$ |  $85.0_{\pm{4.11}}$ $80.6_{\pm{5.03}}$ | $74.5_{\pm{1.17}}$ | $88.7_{\pm{0.20}}$ | $54.8_{\pm{0.34}}$ | -
EF |  $69.0_{\pm{0.01}}$ $-0.55_{\pm{0.58}}$ | $90.3_{\pm{0.23}}$ |  $68.4_{\pm{0.01}}$ $81.2_{\pm{0.01}}$ |  $62.2_{\pm{0.23}}$ $0.11_{\pm{0.07}}$ | $32.9_{\pm{0.09}}$ | $50.9_{\pm{0.59}}$ | $55.6_{\pm{1.30}}$ | -
SF |  $76.1_{\pm{2.27}}$ $39.4_{\pm{7.71}}$ | $93.5_{\pm{0.37}}$ |  $70.3_{\pm{1.23}}$ $81.8_{\pm{0.65}}$ |  $90.3_{\pm{0.03}}$ $87.2_{\pm{0.08}}$ | $77.1_{\pm{0.93}}$ | $64.2_{\pm{0.65}}$ | $57.2_{\pm{0.75}}$ | -
iEF |  $81.7_{\pm{0.16}}$ $50.7_{\pm{1.25}}$ | $\bm{94.4}_{\pm{0.09}}$ |  $\bm{86.2}_{\pm{1.41}}$ $\bm{89.5}_{\pm{0.46}}$ |  $\bm{90.7}_{\pm{0.02}}$ $\bm{87.7}_{\pm{0.05}}$ | $\bm{83.4}_{\pm{0.09}}$ | $\bm{91.9}_{\pm{0.03}}$ | $\bm{74.6}_{\pm{0.85}}$ | -
LoRA
AdamW |  $83.1_{\pm{0.15}}$ $58.7_{\pm{0.55}}$ | $94.9_{\pm{0.07}}$ |  $\bm{88.6}_{\pm{0.51}}$ $\bm{91.9}_{\pm{0.26}}$ |  $90.0_{\pm{0.16}}$ $86.8_{\pm{0.06}}$ | $83.2_{\pm{0.03}}$ | $92.2_{\pm{0.06}}$ | $\bm{83.4}_{\pm{1.06}}$ | $93.7_{\pm{0.32}}$
SGD |  $81.3_{\pm{0.36}}$ $53.6_{\pm{0.97}}$ | $\bm{95.0}_{\pm{0.18}}$ |  $87.3_{\pm{0.28}}$ $90.1_{\pm{0.41}}$ |  $89.9_{\pm{0.69}}$ $86.7_{\pm{0.88}}$ | $\bm{83.3}_{\pm{0.05}}$ | $\bm{92.3}_{\pm{0.09}}$ | $80.9_{\pm{0.72}}$ | $90.6_{\pm{1.02}}$
EF |  $69.1_{\pm{0.01}}$ $10.5_{\pm{2.10}}$ | $91.9_{\pm{0.07}}$ |  $55.5_{\pm{0.28}}$ $71.4_{\pm{0.23}}$ |  $86.3_{\pm{0.28}}$ $82.6_{\pm{0.32}}$ | $61.5_{\pm{0.04}}$ | $89.4_{\pm{0.21}}$ | $52.7_{\pm{0.01}}$ | $30.2_{\pm{1.20}}$
SF |  $79.4_{\pm{0.49}}$ $48.3_{\pm{1.45}}$ | $94.5_{\pm{0.29}}$ |  $76.9_{\pm{0.28}}$ $85.2_{\pm{0.30}}$ |  $\bm{90.1}_{\pm{0.11}}$ $\bm{86.9}_{\pm{0.10}}$ | $81.9_{\pm{0.17}}$ | $91.8_{\pm{0.32}}$ | $71.8_{\pm{0.96}}$ | $92.7_{\pm{0.72}}$
iEF |  $\bm{83.4}_{\pm{0.24}}$ $\bm{59.5}_{\pm{0.64}}$ | $94.9_{\pm{0.21}}$ |  $88.5_{\pm{0.88}}$ $91.8_{\pm{0.55}}$ |  $89.9_{\pm{0.09}}$ $86.6_{\pm{0.12}}$ | $81.2_{\pm{0.02}}$ | $92.2_{\pm{0.11}}$ | $81.7_{\pm{0.55}}$ | $\bm{94.1}_{\pm{0.15}}$
Table 7: Test performance for all task, structure and optimiser combinations. For all tasks, only one test result is reported for the best validation checkpoint across three random seed runs. Task-specific metrics (all multiplied by 100) on the test set are reported in this table. For SST-2, QNLI, RTE and CIFAR100, accuracy is reported. For CoLA, Matthew’s Corr is reported. For MRPC and QQP, F1-score and Accuracy (in order) are reported. For MNLI, matched accuracy and unmatched Accuracy (in order) are reported. | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | CIFAR100
---|---|---|---|---|---|---|---|---
Prompt Tuning
Adafactor | 45.1 | 94.3 |  87.1 82.8 |  71.8 88.8 |  82.9 82.7 | 91.5 | 60.7 | -
SGD | 6.4 | 93.5 |  78.3 66.6 |  71.3 88.9 |  75.7 76.5 | 87.4 | 55.3 | -
EF | $-3.8$ | $90.2$ |  $79.9$ $66.5$ |  $0.3$ $81.6$ |  $33.4$ $33.4$ | $52.4$ | $50.4$ | -
SF | $45.2$ | $93.7$ |  79.7 67.8 |  71.5 88.5 |  77.8 78.5 | $64.1$ | $52.3$ | -
iEF | 50.9 | 94.4 |  88.4 84.2 |  72.0 89.2 |  83.5 83.4 | 91.3 | 68.2 | -
LoRA
AdamW | $\bm{52.2}$ | $\bm{94.5}$ |  $88.6$ $85.1$ |  $71.5$ $88.8$ |  $\bm{83.6}$ $83.1$ | $\bm{92.2}$ | $\bm{71.2}$ | $93.9$
SGD | $47.8$ | $94.0$ |  $79.9$ $66.6$ |  $\bm{71.6}$ $88.8$ |  $83.5$ $\bm{83.8}$ | $91.9$ | $70.1$ | $91.3$
EF | $0.0$ | $92.1$ |  $79.9$ $66.5$ |  $65.4$ $84.4$ |  $61.4$ $62.7$ | $89.4$ | $50.4$ | $31.0$
SF | $42.2$ | $94.2$ |  $84.3$ $75.8$ |  $71.1$ $88.4$ |  $82.1$ $82.3$ | $91.8$ | $64.9$ | $92.8$
iEF | $51.2$ | $94.4$ |  $\bm{89.3}$ $\bm{85.9}$ |  $71.5$ $\bm{88.9}$ |  $81.1$ $80.8$ | $\bm{92.2}$ | $69.1$ | $\bm{94.3}$

Figure 10: Training loss and validation accuracy of all the Prompt Tuning
tasks (7 GLUE tasks). Validation accuracy is reported at the same frequency as
is stated by Table 3. For training loss, 100 points are reported across all
training stages for every task. Each train loss data point represents the
averaged train loss for all train batches between each reported data point.
The error bars represent the standard deviation (1-sigma) for 3 random seed
runs. The training loss for EF always starts to diverge halfway through
training despite the more complicated scheduling. For most tasks, iEF is able
to reach a lower training loss than EF, SF and the well-tuned baseline
Adafactor.

Figure 11: Training loss and validation accuracy of all the LoRA tasks (7 GLUE
tasks and 1 CIFAR100). Validation accuracy is reported at the same frequency
as is stated by Table 3. For training loss, 100 points are reported across all
training stages for every task. Each train loss data point represents the
averaged train loss for all train batches between each reported data point.
The error bars represent the standard deviation (1-sigma) for 3 random seed
runs. The training loss for EF always starts to diverge halfway through
training despite the more complicated scheduling. For most tasks, iEF is able
to reach a lower training loss than EF, SF and the well-tuned baseline AdamW.