Series of Hessian-Vector Products for Tractable Saddle-Free Newton
  Optimisation of Neural Networks
Despite their popularity in the field of continuous optimisation,
second-order quasi-Newton methods are challenging to apply in machine learning,
as the Hessian matrix is intractably large. This computational burden is
exacerbated by the need to address non-convexity, for instance by modifying the
Hessian's eigenvalues as in Saddle-Free Newton methods. We propose an
optimisation algorithm which addresses both of these concerns - to our
knowledge, the first efficiently-scalable optimisation algorithm to
asymptotically use the exact inverse Hessian with absolute-value eigenvalues.
Our method frames the problem as a series which principally square-roots and
inverts the squared Hessian, then uses it to precondition a gradient vector,
all without explicitly computing or eigendecomposing the Hessian. A truncation
of this infinite series provides a new optimisation algorithm which is scalable
and comparable to other first- and second-order optimisation methods in both
runtime and optimisation performance. We demonstrate this in a variety of
settings, including a ResNet-18 trained on CIFAR-10.
Elre T. Oldewage, Ross M. Clarke, José Miguel Hernández-Lobato
# Series of Hessian-Vector Products for Tractable Saddle-Free Newton Optimisation of Neural Networks

###### Abstract

Despite their popularity in the field of continuous optimisation, second-order quasi-Newton methods are challenging to apply in machine learning, as the Hessian matrix is intractably large. This computational burden is exacerbated by the need to address non-convexity, for instance by modifying the Hessian's eigenvalues as in Saddle-Free Newton methods. We propose an optimisation algorithm which addresses both of these concerns -- to our knowledge, the first efficiently-scalable optimisation algorithm to asymptotically use the exact (eigenvalue-modified) inverse Hessian. Our method frames the problem as a series which principally square-roots and inverts the squared Hessian, then uses it to precondition a gradient vector, all without explicitly computing or eigendecomposing the Hessian. A truncation of this infinite series provides a new optimisation algorithm which is scalable and comparable to other first- and second-order optimisation methods in both runtime and optimisation performance. We demonstrate this in a variety of settings, including a ResNet-18 trained on CIFAR-10.

## 1 Introduction

At the heart of many machine learning systems is an optimisation problem over some loss surface. In the field of continuous optimisation, second-order Newton methods are often preferred for their rapid convergence and curvature-aware updates. However, their implicit assumption of a (locally) convex space restricts their usability, requiring the use of mechanisms like damping (Martens, 2010; Dauphin et al., 2014; O'Leary-Roseberry et al., 2021) to avoid degenerate behaviour. In machine learning applications, which are invariably non-convex, high dimensionality further plagues this class of optimiser by creating intractably large Hessian (second-derivative) matrices and a proliferation of saddle points in the search space (Pascanu et al., 2014). These difficulties constrain most practical systems to first-order optimisation methods, such as stochastic gradient descent (SGD) and Adam.

Pascanu et al. (2014) and Dauphin et al. (2014) tackled some of these challenges by proposing Saddle-Free Newton (SFN) methods. In essence, they transform the Hessian by taking absolute values of each eigenvalue, which makes non-degenerate saddle points repel second-order optimisers where typically they would be attractive. Because this transformation would otherwise require an intractable eigendecomposition of the Hessian, they work with a low-rank Hessian approximation, on which this process is achievable, albeit at the cost of introducing an additional source of error.

In this paper, we propose a new route towards SFN optimisation which exploits Hessian-vector products to avoid explicitly handling the Hessian. We use a squaring and square-rooting procedure to take the absolute
value of the eigenvalues without eigendecomposing the Hessian and deploy an infinite series to tractably approximate the expensive square-root and inverse operations. The resulting algorithm is comparable to existing methods in both runtime and optimisation performance, while tractably scaling to larger problems, even though it does not consistently outperform the widely-known Adam (Kingma and Ba, 2015) and KFAC (Martens and Grosse, 2015). To our knowledge, this is the first approximate second-order approach to (implicitly) edit the full Hessian matrix's eigenvalues and be exact in its untruncated form. After summarising previous work in Section 2, we mathematically justify the asymptotic exactness of our algorithm in Section 3 and show its practical use in a range of applications in Section 4. Section 5 concludes the paper.

## 2 Related Work

Although stochastic first-order optimisation methods are the bread and butter of deep learning optimisation, considerable effort has been dedicated to _preconditioned_ gradient methods - methods that compute a matrix which scales the gradient before performing an update step. Newton's method and quasi-Newton methods, which multiply the gradient by the Hessian or an approximation thereof, fall into this category. Other examples include AdaGrad (Duchi et al., 2011) which calculates a preconditioner using the outer product of accumulated gradients, and SHAMPOO (Gupta et al., 2018) which is similar to Adagrad but maintains a separate, full preconditioner matrix for each dimension of the gradient tensor.

Martens (2010) proposes using _Hessian-Free_ (HF) or truncated Newton (Nocedal and Wright, 2006) optimisation for deep learning. The algorithm uses finite differences to approximate the Hessian in combination with the linear conjugate gradient algorithm (CG) to compute the search direction. Like our method, HF implicitly works with the full Hessian matrix and is exact when CG converges.

Pascanu et al. (2014) and Dauphin et al. (2014) present the proliferation of saddle points in high-dimensional optimisation spaces as an explanation for poor convergence of first-order optimisation methods. Various approaches to escaping these saddle points have been proposed. Jin et al. (2017) observe that saddle points are easy to escape by adding noise to the gradient step when near a saddle point, as indicated by a small gradient. Another idea is to normalise the gradient so that progress is not inhibited near critical points due to diminishing gradients (Levy, 2016; Murray et al., 2019).

Saddle points also present a hurdle to second order optimisation, since they become attractive when applying Newton's method. Nevertheless, some work leverages second order information in sophisticated ways to avoid saddle points. For example, Curtis and Robinson (2019) exploit negative curvature information by alternating between classical gradient descent steps and steps in the most extreme direction of negative curvature. Adolphs (2018) builds on this to propose "extreme curvature exploitation", where the eigenvectors corresponding to the most extreme positive and negative eigenvalues are added to the vanilla gradient update step. Anandkumar and Ge (2016) develop an algorithm which finds stationary points with first, second _and third_ derivatives equal to zero, and show that progressing to a fourth-order optimality condition is NP-hard. Truong et al. (2021) project the Newton update step onto subspaces constructed using the positive- and negative-curvature components of the Hessian, allowing them to negate the updates proposed by the latter.

Pascanu et al. (2014) propose the Nonconvex Newton Method, which constructs a preconditioner by decomposing the Hessian and altering it so that all eigenvalues are replaced with their absolute values and very small eigenvalues are replaced by a constant. Unfortunately, explicit decomposition of the Hessian is expensive and does not scale well to machine learning applications. Dauphin et al. (2014) extend this work by proposing the Saddle-Free Newton (SFN) method, which avoids computing and decomposing the exact Hessian by an approach similar to Krylov subspace descent (Vinyals and Povey, 2012), which finds \(k\) vectors spanning the \(k\) most dominant eigenvectors of the Hessian. However, this approach relies on the Lanczos algorithm, which is known to be unstable (Cahill et al., 2000; Scott, 1979). O'Leary-Roseberry et al. (2021) instead invert a low-rank approximation to the Hessian for improved stability. However, their method is susceptible to poor conditioning at initalisation and is limited to very small step sizes in settings with high stochasticity. Consequently, it is unclear how well the algorithm extends beyond the transfer learning settings illustrated.
Instead, our work writes the inverse of the squared and principal square-rooted Hessian as a series, of which we can compute a truncation without explicitly computing or eigendecomposing the Hessian, thereby avoiding instabilities faced by Dauphin et al. (2014) and O'Leary-Roseberry et al. (2021).

There are other examples in machine learning where infinite series are used to motivate approximations to the inverse Hessian (Lorraine et al., 2020; Clarke et al., 2022); we exploit the same construction as Song et al. (2021) to compute the square root of a matrix.

An alternative approach is to precondition the gradient with a curvature matrix that is positive semi-definite by definition, thereby circumventing concerns surrounding saddle points. Notably, the natural gradient method (Amari, 1998) preconditions the gradient with the inverse Fisher information matrix, rather than the inverse Hessian. Whereas the Hessian measures curvature in the model parameters, the Fisher quantifies curvature in terms of the KL-divergence between model and data probability distributions. The natural gradient can be approximated by methods like Factorized Natural Gradient (Grosse and Salakhudinov, 2015) and Kronecker-Factored Approximate Curvature (KFAC) (Martens and Grosse, 2015). In particular, KFAC approximates the Fisher with a block diagonal matrix, which significantly reduces the memory footprint and reduces the cost of inversion. KFAC also leverages several other "tricks", which are relevant for later discussion. We provide a brief overview below and further details in Appendix A.4:

**Moving average of curvature matrix**: KFAC maintains an online, exponentially-decaying average of the approximate curvature matrix, which improves its approximation thereof and makes the method more robust to stochasticity in mini-batches.
**Adaptive learning rate and momentum factor**: KFAC's update rule incorporates a learning rate and a momentum factor which are both computed adaptively by assuming a locally quadratic model and solving for the local model's optimal learning rate and momentum factor at every iteration.
**Tikhonov damping with Levenberg-Marquardt style adaptation.**: KFAC incorporates two damping terms: \(\eta\) for weight regularisation, and \(\lambda\) which is adapted throughout training using Levenberg-Marquardt style updates (More, 1978). The damping constant \(\lambda\) can be interpreted as defining a trust region for the update step. When the curvature matrix matches the observed landscape, the trust region is grown by shrinking \(\lambda\); otherwise damping is increased so that optimisation becomes more SGD-like.

KFAC is arguably the most popular second-order method enjoying widespread use in practice, so we include it as an important baseline in Section 4. Section 4.1 describes our studies incorporating similar adaptive mechanisms into our method, which we now proceed to derive.

## 3 Derivations

Suppose we wish to minimise some scalar function \(f(\mathbf{x})\) over the vector quantities \(\mathbf{x}\), which have some optimal value \(\mathbf{x}^{*}\). Denote by \(\mathbf{g}=\nabla_{\mathbf{x}}f=\nabla f(\mathbf{x})\) and \(\mathbf{H}=\nabla_{\mathbf{x}}(\nabla_{\mathbf{x}}f)^{\mathsf{T}}\) the gradient vector and Hessian matrix of \(f\), respectively, with both quantities evaluated at the present solution \(\mathbf{x}\). We make no assumptions about the convexity of \(f(\mathbf{x})\).

### Preliminaries

Under a classical Newton framework, we can approximate a stationary point \(\mathbf{x}^{*}\) by writing a second-order Taylor series for perturbations around some \(\mathbf{x}\). Assuming \(\mathbf{H}\) is invertible, this recovers

\[\mathbf{x}^{*}\approx\mathbf{x}-\mathbf{H}^{-1}\mathbf{g}, \tag{1}\]

where the RHS is the Newton update to \(\mathbf{x}\). In effect, we have locally approximated \(f\) about \(\mathbf{x}\) by a quadratic function, then set \(\mathbf{x}\) to the stationary point of this quadratic. The invertibility of \(\mathbf{H}\) guarantees that this stationary point is unique. However, if \(\mathbf{H}\) is not positive definite -- for instance, if the function is locally non-convex -- that stationary point may be a maximum or saddle point of the approximated space, rather than a minimum.
To address this limitation, we might consider the eigendecomposition of \(\mathbf{H}\). Since \(\mathbf{H}\) is real and symmetric for non-degenerate loss functions, its eigenvalues are real and its eigenvectors may be chosen to be orthonormal. We can interpret the eigenvectors as the 'principal directions of convexity', and the eigenvalues as the corresponding magnitudes of convexity in each direction (where negative eigenvalues encode concavity). As \(\mathbf{H}^{-1}\) has equal eigenvectors to \(\mathbf{H}\) and reciprocal eigenvalues, we may interpret the product \(\mathbf{H}^{-1}\mathbf{g}\) as a transformation of the gradient vector, with anisotropic scaling governed by the directions and magnitudes of convexity in \(\mathbf{H}\). Moreover, this product gives _exactly_ the updates necessary to move along each principal direction of convexity to the stationary value in that direction, according to the locally quadratic approximation implied by \(\mathbf{H}\). This is illustrated in Figure 1.

As positive eigenvalues are associated with directions of convex curvature, \(\mathbf{H}^{-1}\) selects updates in these directions which _decrease_ the loss function. Conversely, \(\mathbf{H}^{-1}\) selects updates which _increase_ the loss function in the directions associated with negative eigenvalues -- directly opposing our goal of minimising \(f\). Intuitively, we would like to reverse the direction of the latter updates, such that they are decreasing \(f\). This is equivalent to changing the sign of the corresponding eigenvalues.

This intuitive idea was presented by Pascanu et al. (2014), and Dauphin et al. (2014) establish a more direct derivation using a trust region framework, which motivates taking the absolute value of every eigenvalue in a _Saddle-Free Newton_ (SFN) method (Figure 1). However, its implementation in deep learning is challenged by the intractably large Hessian matrices of non-trivial neural networks. Previous work (Dauphin et al., 2014; O'Leary-Roseberry et al., 2021) tackles this by computing a low-rank approximate Hessian, whose eigendecomposition may be calculated directly, changed as required and approximately inverted. While such an approach secures tractability, the cascade of approximations threatens overall accuracy.

### Absolute Values as Square-Rooted Squares

Our proposed method seeks to transform the eigenvalues of \(\mathbf{H}\) without computing its full eigendecomposition. This approach is inspired by the observation that, for scalar \(x\), \(|x|=+\sqrt{x^{2}}\), where we specifically take the

Figure 1: Motivation for Saddle-Free Newton methods. This locally quadratic surface has a saddle point () and its Hessian gives two principal directions of curvature (—, —). From any initial point (), SGD will give an update neglecting curvature () and Newton’s method converges immediately to the saddle point (). Exact Saddle-Free Newton () takes absolute values of the Hessian eigenvalues, negating the components of the Newton update in concave directions (—) and thus changing the saddle point from an attractor to a repeller. Our series-based method () is an approximate Saddle-Free Newton algorithm which converges to the exact Saddle-Free Newton result.
positive square root. In the matrix case, we may define \(\mathbf{S}\) as a square root of a square matrix \(\mathbf{A}\) iff \(\mathbf{A}=\mathbf{S}\mathbf{S}\). For a square, positive semi-definite \(\mathbf{A}\), there is a unique positive semi-definite square root \(\mathbf{B}\), which we term the _principal_ square root of \(\mathbf{A}\); we will write \(\mathbf{B}=\sqrt[n]{\mathbf{A}}\).

If \(\mathbf{A}\) is real and symmetric, we may eigendecompose it as \(\mathbf{Q}\mathbf{\Lambda}\mathbf{Q}^{\mathsf{T}}\), where \(\mathbf{Q}\) is the orthonormal matrix whose columns are the eigenvectors of \(\mathbf{A}\) and \(\mathbf{\Lambda}\) the diagonal matrix whose elements are the corresponding eigenvalues of \(\mathbf{A}\). Then, we have \(\mathbf{B}=\mathbf{Q}\mathbf{\Lambda}^{\frac{1}{2}}\mathbf{Q}^{\mathsf{T}}\). Since raising the diagonal matrix \(\mathbf{\Lambda}\) to the \(k\)th power is equivalent to raising each diagonal element to the \(k\)th power, \(\mathbf{B}\) has the same eigenvectors as \(\mathbf{A}\), but the eigenvalues of \(\mathbf{B}\) are the square roots of those of \(\mathbf{A}\). By taking the principal square root, we guarantee that all the eigenvalues of \(\mathbf{B}\) are non-negative, hence we have taken the positive square root of each eigenvalue in turn.

This reveals a route to transforming our Hessian \(\mathbf{H}\) by taking the absolute value of its eigenvalues. Consider the eigendecomposition \(\mathbf{H}=\mathbf{Q}\mathbf{\Lambda}\mathbf{Q}^{\mathsf{T}}\), noting that \(\mathbf{H}^{2}=\mathbf{Q}\mathbf{\Lambda}^{2}\mathbf{Q}^{\mathsf{T}}\) is positive semi-definite by construction, as its eigenvalues are squares of real numbers. But then \(\sqrt[n]{\mathbf{H}^{2}}\) is the unique positive semi-definite square root of \(\mathbf{H}^{2}\), and each eigenvalue of \(\sqrt[n]{\mathbf{H}^{2}}\) is the positive square root of the square of the corresponding eigenvalue of \(\mathbf{H}\) -- equivalently, its absolute value. Thus, we may take the absolute value of each eigenvalue of \(\mathbf{H}\) by computing the square and then taking the principal square root of \(\mathbf{H}\), i.e. by computing \(\sqrt[n]{\mathbf{H}^{2}}\).

### Inverse Square Root Series

To use this transformed \(\mathbf{H}\) as a second-order preconditioner, we must also invert it, so the matrix of interest is \(\left(\sqrt[n]{\mathbf{H}^{2}}\right)^{-1}\). We now develop a series approximation to this quantity. For scalars \(z\), we may exploit the generalised binomial theorem to write

\[(1-z)^{-\frac{1}{2}}=\sum_{k=0}^{\infty}\frac{1}{2^{2k}}\binom{2k}{k}z^{k} \tag{2}\]

Applying the root test for convergence, a sufficient condition for the convergence of this series is \(\limsup_{n\to\infty}|z^{n}|^{\frac{1}{k}}<1\). We generalise this series to the matrix case by replacing the absolute value \(|\cdot|\) with any compatible sub-multiplicative matrix norm \(\|\cdot\|\) and writing \(\mathbf{I}-\mathbf{Z}\) in place of \(1-z\). Ideally, we would set \(\mathbf{Z}=\mathbf{I}-\mathbf{H}^{2}\) and recover a power series directly, but to ensure convergence we will require a scaling factor \(V\) such that \(\mathbf{Z}=\mathbf{I}-\frac{1}{V}\mathbf{H}^{2}\). With this addition, we have

\[(\mathbf{H}^{2})^{-\frac{1}{2}}=\frac{1}{\sqrt{V}}\sum_{k=0}^{\infty}\frac{1} {2^{2k}}\binom{2k}{k}\left(\mathbf{I}-\frac{1}{V}\mathbf{H}^{2}\right)^{k}. \tag{3}\]

For this matrix series to converge, we require \(\limsup_{n\to\infty}\|\big{(}\mathbf{I}-\frac{1}{V}\mathbf{H}^{2}\big{)}^{n} \|^{\frac{1}{k}}<1\). By Gelfand's formula, this limit superior is simply the spectral radius of \(\mathbf{I}-\frac{1}{V}\mathbf{H}^{2}\) which, this being a real symmetric matrix, is exactly the largest of the absolute value of its eigenvalues. Denoting the largest-magnitude eigenvalue of \(\mathbf{H}^{2}\) by \(\lambda_{\max}\), our convergence condition is thus equivalent to \(V>\frac{1}{2}\lambda_{\max}\). Further, if we strengthen the bound to \(V>\lambda_{\max}\), we have that \(\big{(}\mathbf{I}-\frac{1}{V}\mathbf{H}^{2}\big{)}^{k}\) is positive semi-definite for \(k=0,1,2,\cdots\), so our series, regardless of where it is truncated, produces a positive semi-definite matrix. We are thus guaranteed to be asymptotically targeting the principal square root. Since \(\big{\|}\mathbf{H}^{2}\big{\|}\geq\lambda_{\max}\) for any sub-multiplicative norm \(\|\cdot\|\), a more practical bound is \(V>\big{\|}\mathbf{H}^{2}\big{\|}\). See Appendix D for further analysis of the correctness, convergence and behaviour around critical points of this series.

### Hessian Products, Choice of \(V\) and Series Acceleration

Although we have avoided directly inverting or square-rooting a Hessian-sized matrix, explicitly computing this series remains intractable. Instead, recall that our quantity of interest for second-order optimisation is
\(\left(\sqrt[4]{\mathbf{H}^{2}}\right)^{-1}\mathbf{g}\), and consider the series obtained by multiplying (3) by \(\mathbf{g}\):

\[(\mathbf{H}^{2})^{-\frac{1}{2}}\mathbf{g}=\frac{1}{\sqrt{V}}\sum_{k=0}^{\infty} \frac{1}{2^{2k}}\binom{2k}{k}\left(\mathbf{I}-\frac{1}{V}\mathbf{H}^{2}\right)^ {k}\mathbf{g}. \tag{4}\]

Denoting by \(\mathbf{a}_{k}\) the \(k\)th term of this summation, we have \(\mathbf{a}_{0}=\frac{1}{\sqrt{V}}\mathbf{g}\) and \(\mathbf{a}_{k}=\frac{2k(2k-1)}{4k^{2}}\left(\mathbf{a}_{k-1}-\frac{1}{V} \mathbf{H}\mathbf{H}\mathbf{a}_{k-1}\right)\). With two applications of the Hessian-vector product trick (Pearlmutter, 1994), we can compute \(\mathbf{H}\mathbf{H}\mathbf{a}_{k-1}\) at the cost of two additional forward and backward passes through the model -- a cost vastly smaller than that of storing, manipulating and inverting the full Hessian. By unrolling this recursion, we can thus efficiently compute the summation of a finite number of the \(\mathbf{a}_{k}\).

Under this framework, we have ready access to the product \(\mathbf{H}^{2}\mathbf{g}\), so can use the loose adaptive heuristic \(V\geq\frac{\left\|\mathbf{H}^{2}\mathbf{g}\right\|}{\|\mathbf{g}\|}\), which we found to be the most performant strategy for adapting \(V\).

In practice, we found (4) to converge slowly, and thus benefit from series acceleration. From a variety of strategies, we found the most successful to be a modification due to Sablonniere (1991) of Wynn's \(\epsilon\)-algorithm (Wynn, 1956a). Letting \(\mathbf{s}_{m}\) be the \(m\)th partial sum of (4), the algorithm defines the following recursion:

\[\boldsymbol{\epsilon}_{m}^{(-1)}=0,\qquad\boldsymbol{\epsilon}_{m}^{(0)}= \mathbf{s}_{m},\qquad\boldsymbol{\epsilon}_{m}^{(c)}=\boldsymbol{\epsilon}_{ m+1}^{(c-2)}+\left(\left\lfloor\frac{c}{2}\right\rfloor+1\right)\left( \boldsymbol{\epsilon}_{m+1}^{(c-1)}-\boldsymbol{\epsilon}_{m}^{(c-1)}\right)^ {-1}. \tag{5}\]

We employ the Samelson vector inverse \(\mathbf{a}^{-1}=\frac{\mathbf{a}}{\mathbf{a}^{\dagger}\mathbf{a}}\) as suggested by Wynn (1962). Using these definitions, the sequence \(\boldsymbol{\epsilon}_{m}^{(2l)}\) for \(m=0,1,2,\cdots\) is the sequence of partial sums of the series \(\mathbf{a}_{k}\) accelerated \(l\) times. Thus, we expect the most accurate approximation of (4) to be given by maximising \(l\) and \(m\), acknowledging there is a corresponding increase in computational cost. Pseudo-code for series acceleration is provided in Appendix A.5.

Algorithm 1 incorporates all these elements to form a complete neural network optimisation algorithm. While expanding the series of (4) to a large number of terms may be arbitrarily expensive, we show in the next Section that useful progress can be made on tractable timescales.

```
whiletraining continues do  Compute training loss, gradient \(\mathbf{g}\), Hessian \(\mathbf{H}\) \(V\leftarrow\max\left\{V,\frac{\left\|\mathbf{H}^{2}\mathbf{g}\right\|}{\| \mathbf{g}\|}\right\}\) \(\mathbf{a}_{0},\mathbf{s}_{0}\leftarrow\mathbf{g}\) for\(k\gets 1\)to\(K-1\)do \(\mathbf{a}_{k}\leftarrow\frac{2k(2k-1)}{4k^{2}}\left(\mathbf{a}_{k-1}-\frac{1}{V} \mathbf{H}\mathbf{H}\mathbf{a}_{k-1}\right)\) \(\mathbf{s}_{k}\leftarrow\mathbf{s}_{k-1}+\mathbf{a}_{k}\) endfor  Compute final term \(\hat{\mathbf{s}}_{\infty}\) after \(N\) accelerations of the series \(\mathbf{s}_{K-1-2N},\mathbf{s}_{K-2N},\cdots,\mathbf{s}_{K-1}\) (See Algorithm 2 in Appendix A.5) \(\mathbf{w}\leftarrow\mathbf{w}-\frac{\eta}{\sqrt{V}}\hat{\mathbf{s}}_{\infty}\) endwhile
```

**Algorithm 1** Series of Hessian-Vector Products for Tractable Saddle-Free Newton Optimisation

## 4 Experiments

We now move on to empirical evaluation of our algorithm. For all experiments, we use ASHA (Li et al., 2020) to tune each algorithm and dataset combination on the validation loss, sampling 100 random hyperparameter configurations and setting the maximum available budget based on the model and data combination. Further experimental details and the final hyperparameter settings for all experiments can be found in Appendix A.3, with code available at [http://github.com/rmclarke/SeriesOfHessianVectorProducts](http://github.com/rmclarke/SeriesOfHessianVectorProducts).
We will begin by considering UCI Energy (Tsanas and Xifara, 2012), which is small enough to allow an exact implementation of our algorithm (using eigendecompositions instead of the Neumann series approximation) as a proof of concept, and lends itself to the full-batch setting -- the best case scenario for second-order methods. We then move to a setting without these conveniences, namely Fashion-MNIST (Xiao et al., 2017), which is large enough to require require mini-batching and has too many parameters to allow for exact computation of the Hessian. We go on to increasingly difficult scenarios, in terms of both model and dataset size, by considering SVHN (Netzer et al., 2011) and CIFAR-10 (Krizhevsky, 2009) using ResNet-18 architectures.

For UCI Energy, we generate a random dataset split using the same sizes as Gal and Ghahramani (2016); for Fashion-MNIST, SVHN and CIFAR-10 we separate the standard test set and randomly choose \(\frac{1}{6}\), \(\frac{1}{6}\) and \(\frac{1}{10}\) (respectively) of the remaining data to form the validation set. The numerical data for all experiments can be found in Appendix B.1. While we will usually present wall-clock time on the \(x\)-axis, plots with iteration steps on the \(x\)-axis are available in Appendix B.2.

For all experiments, we present both training and test loss. The optimisation literature often focuses only on the objective function at hand, i.e. the training loss, since a strong optimiser should be able to solve the function it is given. However, in machine learning our target is always to generalise, i.e. to do well on the unseen test set as a measure of generalisation, and the training loss is only a means toward this end. Since we hope to apply our method to deep learning methods, we consider it important to present both these metrics together.

### UCI Energy

We begin with a small-scale experiment on UCI Energy as a proof of concept, training for \(6\,000\) full-batch training epochs. We compare our algorithm to a number of baselines1:

Footnote 1: We also include an L-BFGS (Liu and Nocedal, 1989) baseline in Appendix B.3

**Exact SFN**: Full-Hessian implementation of the absolute-value eigenvalue strategy of Pascanu et al. (2014), where we compute the eigenvalue decomposition and take the absolute value of the eigenvalues. We additionally replace eigenvalues near zero with a small constant and then compute the exact inverse of the resulting saddle-free Hessian. For this method, we tune the learning rate, momentum, the threshold for replacing small eigenvalues, and constant which replaces the small eigenvalues.
**Ours**: Our implementation of Algorithm 1, using tuned learning rate, momentum, series length \(K\) and order of acceleration \(N\). As described in Section 3.4, we adapt \(V\) using the loose bound \(V\geq\frac{\|\mathbf{h}^{2}\mathbf{g}\|}{\|\mathbf{g}\|}\) starting with an initial value of \(100\), as we found minimal benefit to explicitly tuning \(V\). We also considered more accurate approximations to \(V\) that would attain a tighter bound (such as computing the largest eigenvalue using power iteration), but found these held little to no benefit.
**SGD**: Classical stochastic gradient descent, with a tuned learning rate.
**Adam**: (Kingma and Ba, 2015) We tune all the parameters, i.e. learning rate, \(\epsilon\), \(\beta_{1}\) and \(\beta_{2}\).
**KFAC**: **(DeepMind)**: (Martens and Grosse, 2015) We use the implementation of Botev and Martens (2022) which includes adaptive learning rates, momentum, and damping; we tune the initial damping.

The first algorithm above is an exact version of our algorithm, which is tractable in this particular setting. We also considered including an exact implementation of the Newton second-order update but this diverged rapidly, presumably due to the non-convexity of the optimisation task, so we do not include it here.

Figure 2 shows the training and test losses both in terms of wall-clock time and as a function of the number of optimisation steps. Exact SFN achieves the best training loss, as we may hope from it being an exact SFN method. This is encouraging, since it provides proof of concept that our approach is generally sensible in the exact setting. However, it does not converge as quickly as may be desired -- in comparison, KFAC (DeepMind) and Adam make much faster progress, even when considering the change in loss per iteration, rather than wall-time. Our algorithm deflects from the Exact SFN trend, as its approximate nature would suggest, but does not approach the performance exhibited by KFAC DeepMind.
KFAC DeepMind includes clever adaptation mechanisms and smoothing of the curvature matrix which may give it an advantage over the other algorithms. To investigate this, we include additional variants on the baselines:

**Exact SFN (Adaptive)**: Same as Exact SFN, but with adaptive learning rate, momentum and damping strategies as used by KFAC (DeepMind) (see Section 2 and Appendix A.4 for details), as well as an exponential moving average of the curvature matrix. We tune only the initial damping, which subsumes the need for manually replacing small eigenvalues with a constant.
**Ours (Adaptive)**: Our implementation of Algorithm 1, incorporating the adaptive learning rate, momentum and damping used by KFAC (DeepMind). We tune the initial damping, number of update steps and order of acceleration.
**KFAC (Kazuki)**: (Martens & Grosse, 2015) This corresponds to the default settings for KFAC in Osawa. This version of KFAC is not adaptive and does not smooth the curvature matrix by means of averaging. We tune the damping, learning rate and momentum.

Figure 3 shows the training and test loss profiles in wall-clock time. The best test and training losses are now achieved by Exact SFN (Adaptive). We note that this adaptive version of Exact SFN converges considerably faster than the non-adaptive version, reinforcing our and Martens & Grosse's views on the importance of adapting the learning rate, momentum and damping.

Figure 2: Median training (left) and test (right) MSEs achieved over wall-clock time (top) and training iterations (bottom) on UCI Energy by various optimisers in the full-batch setting, bootstrap-sampled from 50 random seeds. Optimal hyperparameters were tuned with ASHA. Note the logarithmic horizontal axes.
In all cases (KFAC, Exact SFN and Ours), the adaptive version of the algorithm performs significantly better than the non-adaptive version. Although our adaptive algorithm matches Exact SFN and beats SGD and both KFAC versions in terms of final test loss, it is still surpassed by Adam and SFN Exact (Adaptive). Notably, our non-adaptive algorithm does not match the performance of SFN Exact, neither does our adaptive algorithm match SFN Exact (Adaptive). Clearly, we sacrifice training performance by using an approximation to Exact SFN and by not smoothing the curvature matrix.

KFAC (DeepMind) achieves the second best training loss, though not test loss. KFAC (Kazuki) diverges quickly at the start of training, which is also unexpected given that its hyperparameters were tuned and that it behaves reasonably on the later, more difficult problems. We hypothesise that adaptive parameters are an important component of its behaviour and that this setting does not lend itself well to fixed parameters (which is supported by the observation that all the adaptive versions performed better than their non-adaptive counterparts).

In this setting, it seems that short of using exact Hessians, Adam is the best choice of optimiser, displaying the second-best training and test losses and completing faster (in wall-clock time) than the second-order methods. However, we are encouraged that our adaptive algorithm's performance is not far off the exact version and continue to more realistic settings in the sections that follow.

### Larger Scale Experiments

Most practical applications are too large to permit full-batch training and so the remainder of our experiments incorporate mini-batching. Since second-order methods may benefit from larger batch sizes, we tune for batch size, choosing from the set \(\{50,100,200,400,800,1600,3200\}\).

We show the best (lowest) losses achieved by each algorithm in each problem setting in Figure 5 as well as the training and test loss profiles in Figure 4. Although KFAC (DeepMind) usually attains the best training loss, there is no clear consistent winner in terms of the best test loss achieved across all problems, despite each algorithm having been tuned specifically for each problem.

Surprisingly, KFAC (DeepMind) performs poorly on Fashion-MNIST, where KFAC (Kazuki) and Adam perform well. First-order optimisers seem well-suited to SVHN, where SGD and Adam achieve the best test losses. On CIFAR-10, Adam and the two KFAC variants perform about the same in terms of training loss, but KFAC (DeepMind) performs significantly better in terms of test loss.

Figure 3: Median training (left) and test (right) MSEs plotted against the log of wall-clock time, with additional optimisers included. Results are on UCI Energy in the full-batch setting and are bootstrap-sampled from 50 random seeds. Optimal hyperparameters were tuned with ASHA.
Figure 4: Median training (left) and test (right) loss achieved on Fashion-MNIST (top), SVHN (centre) and CIFAR-10 (bottom) by various optimisers using the optimal hyperparameters chosen by ASHA. Values are bootstrap-sampled from 50 random seeds.
Our findings seem to validate the widespread use of Adam in practice, given its simplicity as compared to KFAC (DeepMind). However, the performance of KFAC (DeepMind) on CIFAR-10 does indicate that there may be benefit to considering second-order optimisers more seriously.

Although our method is not the best on any of the datasets, its performance is not far from that of the other methods. Where the KFAC variants and Adam occasionally diverge during training (see Figure 4), our method is reasonably stable. Moreover, by leveraging large batch sizes, we converge in fewer epochs and less time than the first-order methods in some settings (e.g. Ours (Adaptive) on Fashion-MNIST is faster than both SGD and Adam), despite the additional complexity of our algorithm.

### Discussion

We posit that the gap between our performance and expected gains is due to error in our series approximation, of which there are two sources. The first is truncation error, which can be reduced to some extent by increasing the number of terms in the series, though the potency of this will depend on how slow the series is to converge. The second is numerical error: if the Hessian is poorly conditioned, then the repeated multiplications required to compute more terms may cause the series to diverge -- even if we have chosen \(V\) appropriately, so that the series should converge in theory.

By increasing the number of terms in the series, we can test whether the error is due to truncating the series or numerical error. From our experiment in Appendix B.4 examining the effect of truncation length, we find that for UCI Energy, increasing the number of terms in the series improves performance. However, for the larger-scale problems, we found that increasing the number of steps in the series to be arbitrarily large did not necessarily lead to improved performance. There is thus a trade-off between choosing a sufficiently high number of steps to approximate the desired matrix and choosing sufficiently few to avoid numerical issues. Strategies to improve conditioning of the Hessian may also help to improve this trade-off.

We consider KFAC, which also approximates the curvature, yet proves quite successful on the benchmark suite.2 KFAC's Kronecker factorisation supports smoothing the curvature estimate with a moving average, which reduces the impact of occasional, poor-quality approximations. Unfortunately, our full-Hessian approximation cannot support such smoothing due to storage requirements. Moreover, KFAC's approximation (which discards the off-diagonal blocks) can be understood intuitively as ignoring the correlations between weights of different layers. In contrast, the rate of convergence of our series varies throughout optimisation, and the impact of truncating the series on the resulting curvature matrix is more difficult to intuit. It may prove fruitful to leverage the same block-diagonal approximation in our method, but with smaller matrices at less risk of ill-conditioning. This would also allow the use of smoothing, which may further improve performance.

Footnote 2: In fact, based on KFAC’s performance in all our experiments, we found it surprising that KFAC is not more widely utilised in practice. That said, we also found its performance to vary widely between implementations, which may explain this observation.

There are links between our series approximation and the conjugate gradient (CG) method (Hestenes and Stiefel, 1952). CG solves a linear system of the form \(\mathbf{Ax}=\mathbf{b}\) iteratively. At the \(k\)-th iteration, CG finds the best \(\mathbf{x}\) in the \(k\)-th Krylov subspace (where the \(k\)-th Krylov subspace is the subspace generated by repeated applications of \(\mathbf{A}\) to the residual \(\mathbf{r}\), i.e. \(\mathcal{K}_{k}=\text{span}\{\mathbf{r},\mathbf{Ar},...,\mathbf{A}^{k-1} \mathbf{r}\}\)) where \(\mathbf{r}=\mathbf{b}-\mathbf{Ax}_{0}\). The inverse Neumann approximation truncated at the \(k\)-th term also finds a vector in the \(k\)-th Krylov subspace, but it is not guaranteed to be the optimal one, and so we may expect the Neumann approximation to be worse than CG.3 Although the series we present in (3) is slightly different, since it is computing the square and square-root at the same time as the inverse, we note that this may provide a clue as to the poor convergence behaviour of the series in general. Future work may consider leveraging insights from the conjugate gradient method to better approximate the inverted saddle-free Hessian.

Footnote 3: However, there is literature showing that Neumann series are more stable than CG in neural networks (Shaban et al., 2019; Liao et al., 2018)

## 5 Conclusions

In this work, we have motivated, derived and justified an approach to implementing Saddle-Free Newton optimisation of neural networks. By development of an infinite series, we are able to take the absolute
Figure 5: Ranking of optimisers according to lowest training (left) and test (right) losses achieved on Fashion-MNIST (top), SVHN (centre) and CIFAR-10 (bottom). Error bars show standard error in the mean. Values are the minimum of the loss profile across time, generated by bootstrap sampling from 50 random seeds.
values of Hessian eigenvalues without any explicit decomposition. With the additional aid of Hessian-vector products, we further avoid any explicit representation of the Hessian. To our knowledge, this is the first approximate second-order method to alter Hessian eigenvalues with an asymptotic exactness guarantee, and whose convergence is limited by compute time rather than available memory. Our algorithm tractably scales to larger networks and datasets, and although it does not consistently outperform Adam or a well-engineered KFAC implementation, its behaviour is comparable to these baselines, in terms of test loss and run time.

Improvements to the inverse approximation such as leveraging Kronecker factorisation or ideas from the conjugate gradient method may provide fruitful avenues of research for future saddle-free Hessian-based optimisation algorithms such as ours. Strategies to reduce numerical error, such as methods to improve the condition number of the Hessian, should also be investigated. Our findings generally support the widespread use of Adam, which performed well on most benchmarks, often beating KFAC despite being a much simpler algorithm. However, the strong performance of KFAC on CIFAR-10, our most complex benchmark, indicates that there may yet be significant gains by applying second-order methods to deep learning.

## Acknowledgements

We acknowledge computation provided by the CSD3 operated by the University of Cambridge Research Computing Service (www.csd3.cam.ac.uk), provided by Dell EMC and Intel using Tier-2 funding from the Engineering and Physical Sciences Research Countil (capital grant EP/P020259/1), and DiRAC funding from the Science and Technology Facilities Council (www.dirac.ac.uk).

Ross Clarke acknowledges funding from the Engineering and Physical Sciences Research Council (project reference 2107369, grant EP/S515334/1).
