Series of Hessian-Vector Products for Tractable Saddle-Free Newton Optimisation of Neural Networks
Authors: Elre Talea Oldewage
Ross M Clarke
José Miguel Hernández-Lobato
Venue: Accepted by TMLR ()

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 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.

Published in Transactions on Machine Learning Research (02/2024)
Series of Hessian-Vectee Products for Tractable Saddle-Free Newton Optimisation of Neural Networks
Elre T. Oldewage*
etv21 @com.ccmmc. uk
Deportment of Engineering University of Cambridge
Ross M. Clarke*
7
Department of Engineering University of Combridge
José Miguel Hernández-Lobato
jmh2330cam.ccc uk
Department of Engineering University of Combridge
(*Equal contribution; randomly ordered)
Reviewed on OpenReview:
Abstract
Despite their popularity in the field of continuous optimisation, second-order quasi-Newton methods are challenging to apply in machine learning, the Hessian matrix is intractably large. This computational burden is exacerbated by the need to address non-come 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-scalablee optimisation algorithm asymptotically use the exact inverse Hessian with absolute-value eigenvalues. Our method frames the problem as a series which principally squaro-roots and inverts the squared Hessian, then uses it to procondition gradient vector, all without explicitly computing 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 optimisation problem 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 (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 (sccond-derivative) matrices and a prolifcration of saddle points in the scarch space (Pascanu et al., 2014). These difficulties constrain most. practical systems first-order optimisation methods, such 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 cigendecomposition of the Hessian, they work with a low-rank Hessian approximation, on which this process achievable, albeit at the cost of introducing an additional source of error.

---

Published in Transactions on Machine Learning Research (02/2024)
In this paper, we propose 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 & Ba, 2015) and KFAC (Martens & Grosse, 2015). To our knowledge, this is the first approximate second-order approach to (implicitly) take the absolute value of 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 update step. Newton's method and quasi-Newton methods, which multiply the gradient by the Hessian or 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 el al., 2018) which is similar to Adagrad but maintains a separate, full preconditioner matrix for each dimension of the gradient tensor.
which scales the gradient before performing
Martens (2010) proposes using ffessian-Free (HF) or truncated Newton (Nocodal & 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. Varions 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 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 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 & 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 propose "extreme curvature exploitation", where the eigenvectors corresponding to the most extreme positive and negative cigenvalues step. Anandkumar & 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 Na-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.
added to the vanilla gradient update
Nesterov & Polyak (2006) show that regularising Newton's method with the cubed update norm secures convergence second-order stationary points, with further developments improving robustness to approximate Hessians (Tripuraneni et al., 2018); considering approximate (Cartis et al., 2011) efficient (Carmon & Duchi, 2019) solutions to the cubio-regularised sub- nub-problem; introducing momentum (Wang et al, 2020); extending the analysis to apply to mini-batched scttings (Wang et al., 2019). Limiting the parameter updates to some trust region in which the Newlon approximation can be relied upon provides an additional analytical perspective (Nocedal & Wright, 2006), with sub-solvers such as that of Curtis et al. (2021) having appealing convergence behaviour. However, the latter note "these methods have often been designed primarily with complexity guarantees in mind and, as a result, represent a departure from the algorithms that have proved to be the most effective in practice".
2

---

Published in Transactions on Machine Learning Research (02/2024)
Pascanu et al. (2014) propose the Nonconvex Newton Method, which constructs a preconditioner by decom- posing 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 Newlon (SFN) method, which avoids computing and decomposing the exact Hessian by an approach similar to Krylov subspace descent (Vinyals Povey, 2012), which finds l vectors spanning the k most dominant eigenvectors of the Hessian. However, this approach relies on the Lanczos algorithm, which has historically been unstable (Cahill et al., 2000; Scott, 1979) and requires careful implementation to avoid these issues (Saad, 2011). "Leary-Roseberry et al. (2021) instead invert 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 slep sizes in settings with high stochasticily. Consequently, it is unclear how well the algorithm extends beyond the transfer learning settings illustrated.
Other saddle-avoiding approaches include detecting non-convexity by comparing optimisation performance to that which would be expected if the function were convex. Carmon el al. (2017) do this by analysing the convergence of desterov-accelerat
gradient descent, while Royer et al. (2020) combine the Conjugate Gradient method with damping to leverage non-convex update steps to the extent allowed by trust region and (Liu& & Roosta, 2023) follows non-convex directions identified by the Minimum-Reste method to escape indefinite regions.
Instead, our work writes the inverse of the squared and principal square-rooted Hessian as a series, of which we can compute truncation without explicitly computing or eigendecomposing the Hessian, thereby avoiding certain instabilities faced by Dauphin et al. (2014) and O'Leary-Roeeberry et al. (2021). This be seen extension of Agarwal et al. (2017)'s use of a Neumann series to approximate the product of the inverse Hessian and gradient vector.
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 matrix.
An alternative approach is to precondition the gradient with curvature matrix that is positive semi-dcfinite 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 KI-divergence between model and data probability distributions. The natural gradient can be approximated by methods like Factorized Natural Gradient (Grosse & Salakhudinov, 2015) and Kronccker-Factored Approximate Curvature (KFAC) (Martens & Grosse, 2015). In particular, KFAC approximates the Fisher with 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-decayin 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 every iteration.
Tikhonov damping with Brrenberg-Marquardee style adaptation. KFAC incorporates two damping lerms: 1/ for weight regularisation, and X which is adapted throughout training using Levenberg- Marquardt style updates (Moré, 1978). The damping constant 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 x; otherwise damping is increased so that optimisation becomes more SGD-like.
3

---

Published in Transactions on Machine Learning Research (02/2024)
KFAC is arguably the most popular second-order method enjoying widespread use in practice, so include it 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(x) over the vector quantities x, which have some optimal value x". Denote by g = Vxf= = fx()) and = vx(Vxf) the gradient vector and Hessian matrix of f, respectively, with both quantities evaluated at the present solution x. We make no assumptiinn about the convexity of f(x).
3.1 Preliminaries
Under classical Newton framework, we can approximate a stationary point x" by writing a sccond-order Taylor series for perturbations around some x. Assuming H is invertible, this recovers
x x~
(1)
where the RHS is the Newton update to x. In effect, we have locally approximated f about x by a quadratic function, then set x to the stationary point of this quadratic. The invertibility of H guarantees that this stationary point is unique. However, if H is not positive definite for instance, if the function is locally non-convex - that stationary point may be a maxinum saddle point of the approximated space, rather than a minimum.
To address this limitation, we might consider the cigendecomposition of H. Since H is real and symmetric for non-degenerate loss functions, 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 correspouding magnitudes of convexity in each direction (where negative eigenvalues encode concavity). As H has equal eigenvectors to H and reciprocal eigenvalues, may interpret the product H gg transformation of the gradient vector, with anisotropic scaling governed by the directions and magnitudes of convexity in H. Morcover, this product gives eractly the updates necessary to move along each principal direction of convexily to the stal ionary value in that direction, according to the locally quadratic approximation implied by H. This is illustrated in Figure 1.
As positive cigenvalues are associated with diroctions of convex curvature, H- selects updates in these directions which decrease the loss function. Conversely, H selects updates which increase the loss function in the directions associated with negative eigenvalues directly opposing goal of minimising f. Intuitively, we would like to reverse the direction of the latter updates, such that they are decrcasing 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-Fre 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; L'Leary-Roseeerry et al., 2021) tackles this by computing 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.
3.2 Absolute Values as Square-
Square-Rooted Rooted Squares
Our proposed method seeks lo transform the eigenvalues of H without computing its full eigendecoomposition. This approach is inspired by the observation that, for scalar x, [x] = +vzz, where we specifically take the positive square root. In the matrix case, we may define as a square root square matrix A iff A ss. For a square, positive semi-definite A, there unique positive semi-definite square root B, which we term the principal square root of A; we will write B = VA.
4

---

Published in Transactions on Machine Learning Research (02/2024)
1
Loss
Parameter 2
A= 0
Newwton
X Neuttn
XzGD
Parameter 2
oour
Exact SPN
Parameter 1
Xours
Exact SPN
Parameter
Figure 1: Motivation for Saddle-Fres Newton methods. This locally quadratic surface has saddle point and its Hessian gives two principal directions of curvature (-, -). From any initial point (o), SGD will give an update neglecting curvature point (+))). Exact Fadd Newton
and Newton's method converges immediately to the saddle
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 repeller. Our scrics-based method (+++) is approximate Saddle-Fre Newton algorithm which converges to the exact Saddle-Fres Newton result.
If A is real and symmetric, may eigendecompose it QAQ where Q is the orthonormal matrix whose columns are the eigenvectors of A and A the diagonal matrix whose elements are the corresponding eigenvalues of A. Then, we have B = QA Since raising the diagonal matrix A to the kth power is equivalent to raising each diagonal element to the kth power, has the same eigenvectors as A, but the cigenvalues of are the square roots of those of A. By taking the principal square root, we guarantee that all the eigenvalues of B are non-negative, hence we have taken the positive square root of each eigenvalue in turn.
This reveals route lo transforming Hessian H by laking the absolute value of its eigenvalues. Consider the rigendecomposition H = QAQT, noting that H = QA²Q is positive semi-definite by construction, its cigenvalucs are squares of real numbers. But then 2 is the unique positive semi-definite square root of H2, and each eigenvalue of VH2 the positive square root of the square the corresponding eigenvalue of H equivalently, absolute value. Thus, we may take the absolute value of each eigenvalue of H by computing the square and then taking the principal square root of H, ie. by computing VH²
3.3 Inverse Square Root Series
To use this transformed H as second-order preconditioner, we must also invert it, so the matrix of interest. We now develop a series approximation to this quantity. For scalars 2, may exploit the generalised binomial thcorom to write
is
H
80
2k
(1 -)
22k
k:
k=0
Applying the root test for convergence, a sufficient condition for the convergence of this series is lim supn-+oo |2"|4 < 1. We generalise this series to the matrix case by replacing the absolute value |- with any compatible sub-multiplicative matrix norm ||+||| and writing I Z in place of Ideally, would set
5

---

Published in Transactions on Machine Learning Research (02/2024)
- H and recover a power series directly, but to ensure convergence will require scaling factor V such that Z = I + H2. With this addition, we have
2k
H² *
H
(3)
VV
22k
k
< By Gelfand's formula, this limit superior is simply the spectral radius of I H2 which, this being a real symmetric matrix, is exactly the largest of the absolute value of its eigenvalues. Denoting the largest-magnitude eigenvalue of H² by Amax: our convergence condition is thus equivalent to V> Amaax Further, if we strengthen the bound to V > Amax:
For this matrix series to converge, we require lim sup, (I- -
have that (I - y )² k is positive semi-definite for k 0,1,2,...
scries, regardless of where it is truncated, produces positive semi-defnitte matrix. We are thus guaranteed to be asymptotically targeting the principal square root. Since | 2 Amax for any sub-multipliativee norm |I-II, a more practical bound is V > | See Appendix D for further analysis of the correctness, convergence and behaviour around critical points of this series.
3.4 Hessian Products, Choice of V and Series Acceleration
Although we have avoided directly inverting square-rooting a Hessian sized matrix, explicitly computing this series remains intractable. Instead, recall that our quantity of interest for second-order optimisation is WH g, and consider the series obtained by multiplying (3) by g:
00
2k
k
(
VV
22k
ke
H2
g.
(4)
Denoting by the kth term of this summation, we have ag = 1 With two applications of the Hessian-vector product trick (Pearlmutter, 1994), we can compute H the cost of two additional forward and backward passes through the model cost vastly smaller than that of storing, manipulating and inverting the full Hessian. By unrolling this recursion, can thus efficiently compute the summation of finite number of the
. and a = 2k(2k-1) 412 (ak-1 HHa
Under this framework, we have ready access to the product H3g, so can use the loose adaptive heuristic [Hg] which we found to be the most performant strategy for adapting V.
V
[
In practice, found (4) to converge slowly, and thus benefit from series acceleration. From variety of strategies, found the most successful to be modification due Sablonniere (1991) of Wynn's e-algorithm (Wynn, 1956a). Letting Sm be the mth partial sum of (4), the algorithm defines the following recursion:
c (-1)== 0, 0,
(0
(e)
(c-2) cm+1
(1.) + ((TTT) -elm-1))
Sms
We employ the Samelson vector inverse a a as suggested by Wynn (1962). Using these definitions, the sequence (21) for m = 2, is the sequence of partial sums of the series accelerated times. Thus, we expect the most accurate approximation of (4) to be given by maximising and m, acknowledging there is a corresponding increase in computational cost. Pscudo-code for series acceleration provided in Appendix A.5.
Algorithm incorporates all these elements to form complete neural network optimisation algorithm. While expanding the scries 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.
4 Experiments
We now move on 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

---

Published in Transactions on Machine Learning Research (02/2024)
Algorithm Series of Hessian- Vecstor Products for Tractable Saddle-Free Newton Optimisation
while training continues do
Compute training loss, gradient g, Hessian H
max ap, So g
V
for k +- to K - do
ag f- (2(-k11)
(ag-1
HHa
162
Sk <- Sk-1 + a
end for
Compute final term S after accelerations of the series S 2K -N,, "N, 5KK1-
(See Algorithm 2 in Appendix A.5)
R &
end while
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 https: : //github.
We will begin by considering UCI Energy (Tsanas Xifara, 2012), which is small enough allow an exact implementation of our algorithm (using eigendecompositions instead of the Neumann series approximation) proof of concept, and lends itself to the full-batch setting the best case scenario for second-order methods. We then move to 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 random dataset split using the same sizes Gal & Ghahramani (2016): for Fashion-MNIST, SVHN and CIFAR-10 separate the standard test set and randomly choose 6 and (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 lest loss. The optimisation literature often focuses only on the objective function at hand, i.e. the training loss, since 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 measure of generalisation, and the training loss is only 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.
4.1 UCI Energy
We begin with a small-scale experiment on UCI Energy as a proof of concept, training for 6000 full-batch training epochs. We compare our algorithm to number of baselines
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 cigenvalues. 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, adapt V using the loose bound V >
|H&]]
starting
g
We also include B-RFGS (Liu & Nocedal, 1989) baseline in Appendix R.3
7

---

Published in Transactions on Machine Learning Research (02/2024)
with an initial value of 100, as found minimal benefit to explicitly tuning V. Noting the inherent instability of attempting to invert near-zero eigenvalues, we add fixed damping coefficient to the Hessian at each use, which we treat as another typerparameter to optimise. 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 tuned learning rate.
Adam (Kingma & Ba, 2015) We tune all the parameters, iee learning rate, e, S1 and B2.
KFAC (DeepMind) (Martens & Grosse, 2015) We use the implementation Botev & Martens (2022) which includes adaptive learning rales, momentum, and damping; we lune the initial damping.
The first algorithm above is an exact version of algorithm, which is tractable in this particular setting. We also considered inchuding 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.
10a
101
Ours
Adam
105
KFMC (DeepMind)
109
Exact SFN
Test Test Loss
10:4_
Training
10
SED
Adam
10
(DerpMind)
Exact SPN
10-7
10-4
10a
10
10a
10a
10-2
10-
10a
10a
Runtime (s)
Runtime (s)
Ours
Ours
101
SGD
Adam
Adam
10a
KFAC (Deephlnn)
KFAC (DeepMind)
SFN
Exact SFN
109
10
Test Loss
Training
10
10-
10
10a2
Figure Median training (left) and test (right) MSEs achieved over wall-clock time (top) and training iterations (bottom) on UCT 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.
Figure 2 shows the training and test losses both in terms of wall-clock time and a function of the number of optimisation steps. Exact SFN achieves the best training loss, 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 - 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.
8

---

Published in Transactions on Machine Learning Research (02/2024)
KFAC DeepMind inchudes clever adaptation mechanisms and smoothing of the curvature matrix which may give it an ad 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 used by KFAC (DeepMind) (see Section 2 and Appendix A.4 for details), as well an exponential moving average of the curvature matrix. We tune only the initial damping, which subsumes the need for manually replacing small eigenvalues with 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 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 considcrably faster than the non-adaptive version, reinforcing our and Martens & Grosse's 's views on the importance of adapting the learning rate, momentum and damping.
In all cases (KFAC, Exact SFN and Ours), the adaptive version of the algorithm performs significantly better than the non-adaptive version. Although our adaplive algorithm matches Exact SFN and beats SGD and both KFAC versions in terms of final test loss, it is still surpassed by Adam and SEN 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 supported by the observation that all the adaptive versions performed better than their non-adaptive counterparts).
In this setting, it scems that short of using exact Hcssians, 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 's performance is not far off the exact version and continue to more realistic settings in the sections that follow.
4.2 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,201 400, 800, 1600, 3200).
We show the best (lowest) losses achieved by each algorithm in each problem setting in Figure 5 as well 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 FaSIS, 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.
9

---

Published in Transactions on Machine Learning Research (02/2024)
10th
10
1033
1024
Training Loss
1030
Ours (Adaptive)
1014
Ours (Adaptive)
SGD
Test Loss
SGD
107
KFAC (Kazuki)
107
KFAC (Kazul))
(DeepMind)
(DeepMind)
101
Exact
10
Exact SFN
Adaptive Exact
Adaptive Exact
10a
10a
2
10-2
10a4
10-1
10a
10c
10a
10a
10-2
10-1
10a
10a
10a
10a
Runtime (s)
Runtime (s)
101
Ours
104
Ours (Adaptive)
SGD
KIAC (DeepMlndd)
10-1
109
Exact SFN (Adaptive)
Test Loss
Training
10-1
10-3
Ours
Ours (Adaptive)
10
10-1
Adam
(DecpMind)
Exact SPN
Exact (Adaptive)
10-
10
10-4
10a
10a
10a
10a
10
10-1
10a
10a
10a
10a
Runtime (s)
Runtime (s)
Figure 3: Median training (left) and test (right) MSEs plotted against the log of wall-clock time. The top row includes all additional optimisers; the bottom row excludes KFAC (Kazuki) for clarity. Results are on UCI Energy in the full-batch set setting and are bootstrap-sampled from 50 random seeds. Optimal hyperparameters were tuned with ASHA.
10

---

Published in Transactions on Machine Learning Research (02/2024)
Ours (Adaptive)
10
Ours (Adaptive)
Adam
SGD
100
KFAC (Kazuki)
KFAC (DeapMind)
KFAC (Kazuki)
Training Loss
KFAC IDeepMindi
Test Loss
60-1
x 10-1
10-4
10
10
10a
10a
10a
10-2
10"3
107
10:
10
Runtime (s)
Runtime
Ours (Adaptive)
101
10a
Adam
(Kazul))
KFAC (DeapMind)
Training Loss
Test Loss
Our
(Adaptive)
SGD
Adam
KFAC (Kacuki
KFAC (DespMind)
Runtime (s)
Runtime (s)
Ours (Adaptive)
SGD
10
Adamm
KFAC (DeegMind)
Test Loss
Training
10-2
Ours
(Adaptive)
SGD
Adam
KFAC (Kacuki)
KFAC (DeepMind)
10-1
10P
10a
Runtime (s)
Runtime (s)
Figure 4: Median training (left) and test (right) loss achieved on MNSSS (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.
11

---

Published in Transactions on Machine Learning Research (02/2024)
Our findings seem to validate the widespread use of Adam in practice, given its simplicity 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 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.
4.3 Discussion
We posit that the gap between our performance and expected gains is due to error 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 appropriatcly, 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, 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.
found that increasing the number of steps in the series to be arbitrarily large did
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 scries 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.
There links between our series approximation and the conjugate gradient (CG) method (Hestenes & Stiefel, 1952). CG solves a linear system of the form =x b iteratively. At the k-th iteration, CG finds the best x in the k-th Krylov subspace (where the k-th Krylov subspace is the subspace generated by repeated applications of A to the residual r, Le. K4 = span(r, Ar,..,, A where = b - Axo. 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 3G. 3 Although the scries 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 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.
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
IIn fact, based on KFAC's performance in all our experiments, we found surprising that KFAC is not more widely utilised in practice. That said, also found its performance to vary widely between implementations, which may explain this observation. However, there literature showing that Neumann series are more stable than CG in neural networks (Shaban al., 2019; Liao et al., 2018)
12

---

Published in Transactions on Machine Learning Research (02/2024)
0.275
Ours
(Adaptive)
0.250
SGD
Adam
Lowest Training Loes
0.225
Lowest Test Loss
0.375
KFAC (Kazuid)
KFAC (DeapMind)
0.200
0.175
0.365
0.150
Our
(Adaptive)
0.125
SGD
Adam
KFAC (Kacuki)
0.100
(DeepMind)
0.355
Ours
Ours (Adaptive)
SGD
Adam
KFAC (Kazukc)
KFAC (DeapMind)
Ours
Ours (Adaptive)
SGD
Adam
KFAC (Kazuki)
KFAC (DeepMind)
0.5
Ours
Ours (Adaptive)
Ours
SGD
0.675
Ours (Adaptive)
Lowest Training Loss
0.4
KFAC
Adimm
KFAC (DeepMind)
0.650
Adam
(Kazukl)
Lowest Test LOSs
0.625
KFAC (DeepMind)
0.600
0.2
0.575
0.530
0.1
0.525
0.500
0.0
0.475
Ours
Durs (Adaptive)
SGD
Adam
KFAC (Kazuki)
KFAC (Despkindd
Ours
Ours
SGD
Adam
KFAC (Kazuki)
KFAC (DeepMind)
(Adaptive)
1.50
1.2
Ours (Adaptive)
Ours (Adaptive)
SGD
Lowest Training Loss
1.0
C
Adam
KFAC (Kazuki)
1.45
Adam
KFAC (DecpMind)
Lowest Test Loss
1.40
KFAC (DeepMind)
0.6
1.35
1.30
0.2
1.25
0.0
KFAC
KFAC (DeepMind)
(Adaptive)
(Deepkind)
[Adeptive]
(Kazuki)
Figure 5: Ranking of optimisers according 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.
13

---

Published in Transactions on Machine Learning Research (02/2024)
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 take the absolute value of Hessian eigenvalues with an asymptotic exactness guarantce, and whose convergence is limited by compute time rather than available memory. Our algorithm tractably scales to larger works and datasels, and although it does not consistently outperform Adam or a well-engineered KFAC implementation, behaviour is comparable these basclines, in terms of test loss and run time.
Improvements to the inverse approximation such as leveraging Kronecker factorisation ideas from the conjugate gradient method may provide fruitful avenues of rescarch 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 are grateful to Richard E. Turner for valuable discussions about this work and its presentation.
We acknowledge computation provided by the CSD3 operated by the University of Cambridge Research Computing Service (www. (ww..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 EP/P020/22 (PP/002229))) and DiRAC funding from the Science and Technology Facilitios Council (www. airaccac.uu uk).
Ross Clarke acknowledges funding from the Engineering and Physical Sciences Research Council (project reference 2107369, grant EP/S515534/1)).
References
Adolphs, T.. Non Convex-Concave Saddle Point Optimization. Master's thesis, ETH Zürich, April 2018. Accepted: 2018-04-
URL https://www collearcccooliciion ethz. wwwpchhhee/...0000000 500. 18T07:38:552 Publisher: ETH Zurich.
http://///.2.00011105225/8//50 11850/258222
Agarwal, N., Bullins, B., and Hazan, E. Second-Order Stochastic Optimization for Machine Learning in Linear Time. Journal of Machine Learning Research, 18(116):1-40, 2017.
Amari, S.-i. Natural Gradient Works Efficiently in Learning. Newral Computation, 10(2):251-25) 276, February 1998.
Anandkumar, A. and Ge, R. Efficient approaches for escaping higher order saddle points in non-convex optimization. In Conference on Learning Theory, pp. 81-102. PMLR, June 2016. ISSN: 1938-7228.
Botev, A. and Martens, J. KFAC-JAX, 2022. URL
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //tp:////tttttttt jomgg.
Cahill, F.,, Irving, A., Johnston, C., and Sexton, J. Numerical stability of Lanezos methods. Nuclear Physics B - Proceedings Supplements, 83-84:825 827, April 2000.
Carmon, Y. and Duchi, J. Gradient Descent Finds the sebic-Regularized Nonconvex Newton Step. SIAM Journal on Optimization, 2114) 2178, January 2019. Publisher: Society for Industrial and Applied Mathematics.
Carmon, Y., Duchi, J. C., Hinder, O., and Sidford, A. "Convex Until Proven Guilty": Dimension-Free Acceleration of Gradient Descent on Non-Convex Functions. In Proceedings of the 34th International Conference Machine Learning, pp. 654 663. PMLR, July 2017. ISSN: 2640-3498.
14

---

Published in Transactions on Machine Learning Research (02/2024)
Cartis, C., Gould, N. I. M., and Toint, P. L. Adaptive cubic regularisation methods for unconstrained optimization. Part I: motivation, convergence and numerical results. Mathemalicol Programming, 127(2): 245 295, April 2011.
Clarke, R. M., Oldewage, E. T., and Heráández-Lobato, J. M. Scalable One-Pass Optimisation of High- Dimensional Weight-Update Hyperparameters by Implicit Differentiation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 2022.
Curtis, F. E. and Robinson, D. P. Exploiting negalive curvature in deterministic and stochastic optimization. Mathematical Programming: Series A and B, 176(1-2):69 94, July 2019.
Curtis, F. E.,, Robinson, D. P., Royer, C. W., and Wright, S. J. Trust-Region Newton-CG with Strong Second-Order Complexity Guarantees for Nonconvex Optimization. SIAM Journal on Optimization, 31(1): 518-544, January 2021. Publisher: Society for Industrial and Applied Mathematics.
Dauphin, Y. N., Pascanu, R., Gulcehre, C., Cho, K., Ganguli, S., and Bengio, Y. Identifying and attacking the saddle point problem in digeniimnniional non-convex optimization. In Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Duchi, J., Hazan, E., and Singer, Y. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research, 12(61):2121 2159, 2011.
Gal, Y. and Ghahramani, Z. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In International Conference on Machine Learning, pp. 1050- 1050, June 2016. ISSN: 1938-7228 Section: Machine Learning.
Graves-Morris, P. R. review of Padé methods for the acceleration of convergence of a sequence of vectors. Applied Numerical Mathematics, 15(2):153 174, September 1994.
Grosso, R. and Salakhudinov, R. Scaling up Natural Gradient by Sparsely Factorizing the Inverse Fisher Matrix. In Proceedings of the 32nd International Conference on Machine Learning, pp. 2304 2313. PMLR, June 2015. ISSN: 1938-7228.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned Stochastic Tensor Optimization. In Proceedings of the 35th International Conference on Machine Learning, pp. 1842 1850. PMLR, July 2018. ISSN: 2640-3498.
He, K., Zhang, X., Ren, S., and Sun, J. Deep Residual Learning for Image Recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, June 2016. ISSN: 1063-6919.
Hestenes, M. R. and Stiefel, E. Methods of Conjugate Gradients for Solving Linear Systems. Journal of Research of the National Bureau of Standards, 49(6):409 43649999,,99-- December 1952.
Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., and Jordan, M. How to Escape Saddle Points Efficiently. Tn Proceedings of the 34th International Conference on Machine Learning, pp. 1724 1732. PMLR, July 2017.
ISSN: 2640-3498.
Kingma, D. P. and Ba, J. Adam: A Method for Stochastic Optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference Learning Representations, ICLR 2015, 2015.
Krizhevsky, A. Learning Multiple Layers of Features from Tiny Images. Master's thesis, University of Toronto, po.
April 2009. URL https: http://ss..rrotoeeeee//////ee/es
Levin, D. Development of non-linear transformations for improving convergence of sequences. International
Journal of Computer Mathematics, 3(1-4)3371-888, January 1972. Publisher: Taylor & Francis _eprint: hupp:://oi.org 10.100000000077163300000000
Levin, D. and Sidi, A. Two New Classes of Nonlinear Transformations for Accelerating the Convergence of Infinite Integrals and Series. Applied Mathematics and Compulation, 9(3):175 215, October 1981.
15

---

Published in Transactions on Machine Learning Research (02/2024)
Levy, K. Y. The Power of Normalization: Faster Evasion of Saddle Points, November 2016. URL http: //arxiv. //arxiv.or//abat wrgg/aa///11..///eeee
arXiv:1611.0483) (0s, math, stat].
Li, L., Jamieson, K., Rostamizadeh, A., Gonina, E., Ben-tzur, J., Hardt, M., Recht, B., and Talwalkar, A. A System for Massively Parallel Hyperparameter Tuning. Proceedings of Machine Learning and Systems, 2: 230-246, March 2020.
Liao, R., Xiong, Y., Fetaya, E., Zhang, L., Yoon, K., Pitkow, X., Urtasun, R., and Zemel, R. Reviving and Improving Recurrent Back-Propagation. In International Conference on Machine Learning, pp. 3082-3091. PMLR, July 2018.
Liu, D. C. and Nocedal, J. On the limited memory BFGS method for large scale optimization. Malhemalionl Programming, 45(1) 528, August 1989.
Liu, Y. and Roosta, F. A Newton-MR algorithm with complexity guarantees for nonconvex smooth 07095. arXiv:2208.07085
unconstrained optimization, September 2023. URL http://arxtw.orrg/e/ee mmath].
Lorraine, J., Vicol, P., and Duvenaud, D. Optimizing Millions of Hyperparameters by Implicit Differentiation. In Proceedings of the Twenty Third International Conference Artificial Intelligence and Statistics, pp. 1540-1552. PMLR, June 2020. ISSN: 2640-3498.
Martens, J. Deep learning via Hessian-free optimization. In Proceedings of the 27th International Conference International Conference Machine Learning, ICML'10, pp. 735-742, Madison, WI, USA, June 2010. Omnipress.
Martens, J. and Grosse, Optimizing Neural Networks with Cronecker-factoree Approximate Curvature. In International Conference on Machine Learning, pp. 2408-2411 June 2015.
Moré, J. J. The Marquartt algorithm: Implementation and theory. In Watson, G. A. (ed.), Numerical Analysis, Lecture Notes in Mathematics, pp. 105-116, Berlin, Heidelberg, 1978. Springer.
Murray, R., Swenson, B., and Kar, S. Revisiting Normalized Gradient Descent: Fast Evasion of Saddle Points. IEEE Transactions on Automatic Control, 64(11):4818 ):4418-4824, 4824, November 2019. Conference Name: IEEE Transactions on Automatic Control.
Nesterov, Y. and Polyak, B. Cubic regularization of Newton method and its global performance. Mathematical Programming, 108(1):177-205, August 2006.
Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., and Ng, Y. Reading Digits in Natural Images with Unsupervised Feature Learning. In NIPS Workshop on Deep Learning and Unsupervised Feabure Leurving 2011, 2011.
Nocedal, J. and Wright, S. J. Numerical Optimization. Springer, 2006.
O'Leary-Roseberry,, T., Alger, N., and Ghattas, O. Low Rank Saddle Free Newton: A Scalable Method for Stochastic Nonconvex Optimization, August 2021. URL wwwpaa/www...////////o/./// 2288000000000 02881. arXiv:2002.02888 [cs, math).
Osada, N. Acceleration methods for vector sequences. Journal of Computational and Applied Mathematics, 38(1):361 371, December 1991.
Osada, N. Vector sequence transformations for the acceleration of logarithmic convergence. Journal of Computational and Applied Mathematics, 66(1):391-400, Jamuary 1996.
Osawa, K. ASD(FGHJK)L- Automatic Second-order Differentiation (for Fisher, Gradient covariance, Hessian, Jacobian and Kernel) Library. URL https://github wtt./////eettiissa/a/a//a/gg///
Pascanu, R., Dauphin, N., Ganguli, S., and Bengio, Y. On the saddle point problem for non-convex optimization, May 2014. URL http:://axxi wwwaaaiwwwra.aaaaeaeeaeeeeo arggabss14055 4604. 4rXiv:1405.4604 [os].
16

---

Published in Transactions on Machine Learning Research (02/2024)
Paternain, S., Mokhtari, A., and Ribeiro, A. A Newton-Bare -asoed Method for Nonconvex Optimization with Fast. Evasion of Saddle Points. STAM Journal on Oplimizalion, 29(1):343 -388, January 2019. Publisher: Society for Industrial and Applied Mathematics.
Pearlmutter, B. A. Fast exact multiplication by the Hessian. Neural Computation, 6(1):147-160, Jamuary 1994.
Royer, C. W., O'Neill, M., and Wright, S. A Newton-CG algorithm with complexity guarantees for smooth unconstrained optimization. Mathematical Programminng, 180(1):451-488, March 2020.
Saad, Y. Numerical Methods for Large Eigenvalue Problems. Classics in Applied Mathematics. Society for Industrial and Applied Mathematics, Jamuary 2011.
Sablonnière, P. Comparison of four algorithms accelerating the convergence of a subset of logarithmic fixed point sequences. Numerical Algorithms, 122):177-197, June 1991.
Schmidt, R. XXXII. On the numerical solution of linear simultaneous equations by an iterative method. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 32(214):369 383, November 1941. Publisher: Taylor & Francis.
Scott, D. S. How to Make the Lanczos Algorithm Converge Slowly. Mathematics of Computation, 33(145): 239-24F 1979. Publisher: American Mathematical Society.
Shaban, A., Cheng, C.-A,, Hatch, N., and Boots, B. Truncated Back-propagatie for Bilevel Optimization. The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1723-1732, April 2019. Section: Machine Learning.
ISSN: 2640-3498
Shanks, D. Non-linens Transformations of Divergent and Slowly Convergent Sequences. Journal of Mathematics and Physics, 34(1-4):1-422 1955. eprint:
Song, Y., Sebe, N., and Wang, W. Fast Differentiable Matrix Square Root. In International Conference Learning Representations, September 2021.
Tripuraneni, N., Stern, M., Jin, C., Regier, J., and Jordan, M. I. Stochastic Cubic Regularization for Fast Nonconvex Optimization. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Truong, T., To, T. D., Nguyen, T. H., Nguyen, T. H., Nguyen, H. P., and Helmy, M. A fast and simple modification of Newton's method helping to avoid saddle points, September 2021. URL http: //arxiv. 2200 01512. arXiv:2006.0L5 0115111 Jes, math, stat).
Tsanas, A. and Xifara, A. Accurate quantitative estimation of energy performance of residential buildings using statistical machine learning tools. Eneryy and Buildings, 49:560 567, June 2012.
Vinyals, O. and Povey, D. Krylov Subspace Descont for Deep Learning. In Proceedings of the Fifteenth International Conference on Artificial Intelligence and Statislics, pp. 1261-1268. PMLR, March 2012. ISSN: 1938-7228.
Wang, 7., Zhou, Y., Liang, Y., and Lan, G. Stochastic Variance-Reduced Cubic Regularization for Nonconvex Optimization. In Proceedings of the Stventy-Seconi International Conference on Artificial Intelligence and Slatistics, pp.
PMLR, April 2019. ISSN: 2640-3498.
Wang, Z., Zhou, Y., Liang, Y., and Lan, G. Cubie Regularization with Momentum for Nonconvex Optimization. Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, pp. PMLR, August 2020. ISSN: 2640-3498.
Wynn, P. On Device for Computing the em(Sn) Transformation. Mathematical Tables and Other Aids to Computation, 10(54):91-90 -69, 1956a. Publisher: American Mathematical Society.
17

---

Published in Transactions on Machine Learning Research (02/2024)
Wynn, P. On Procrustean technique for the numerical transformation of slowly convergent sequences and series. Mathemalical Proceedings of the Cambridge Philosophical Society, 52(4):663 671, October 1956b. Publisher: Cambridge University Press.
Wynn, P. Acceleration Techniques for Iterated Vector and Matrix Problems. Mathematics of Compulation, 16(79):301-93) 1962. Publisher: American Mathematical Society.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-MNIST; Novel Image Dataset for Benchmarking Machine Learning Algorithms, September 2017. URL http://arxiv org/abs/1708...oo stat].
http://arra...rrgaaa////7.. 07747. arXiv:170..07747 [cs,
18

---

Published in Transactions on Machine Learning Research (02/2024)
A Empirical Notes
A.1 Datasets Used
The datasets use are all standard in the ML literature; outline their usage conditions in Table 1.
Table 1: Licences under which we use datasets in this work.
Dataset
Licenoe
Source
Input
Output
Total Size
Creative Commons Attribution 4.0
Tsanas & Xifara (2012);
UCT Energy
8-Vector
Scalar
692
International (CC BY 4.0)
Gal Ghahramani (2016)
TuSIi
MIT
Xiao (2017)
28 Image
Class (from 10)
60
CIFAR-10
None specified
Krizhevsky (2009)
32 x 32 Image
Class (from 10)
60
SVHN
None specified
Netzer al. (2011)
32 x 32 Image
Class (from 10)
99 289
A.2 Computing Resources Used
The experiments presented were performed using hardware shown in Table 2. All runtime comparisons were thus performed on like-for-like hardware. We make use of GPU acceleration throughout, using the JAX library (Bradbury et al., 2018). Our own code is available at https: //tpt//ggtttt com/mmclakke/ BerrleefffesssaaaecooPrrdddd
Table 2: System configurations used to run experiments.
Type
CPU
GPU (NVIDIA)
Python
JAX
CUDA
Cambridge Service for Data Driven Discovery (CSD3)*
AMD EPYC 7763
Ampere A100
3.9.6
0.3.25
11.1
www csd3 cam ac uk
A.3 Experimental Hyperparameters
We outline in Tables and 4 the search ranges chosen for our hyperparameter optimisation using ASHA, as well as the best hyperparameters chosen for each setting and the corresponding final losses. Our network architectures and corresponding time budgets are enumerated below:
UCI Energy (Tsanas & Xifara, 2012): MLP with 7 hidden layers, each of 12 units (budget 5 minutes). A full training run is 6 000 epochs.
Pashion-MNIST (Xiao et al., 2017): MLP with one hidden layer of 50, units (budget 5 minutes). A full Training run epochs.
SVHN (Netzer et al., 2011): ResNet-18 (He et al., 2016) (budget minutes). full training run epochs.
CIFAR-10 (Krizhevsky, 2009): ResNet-18 (He al, 2016) (budget 2 hours). A full training run is 72 epochs.
For KFAC (DeepMind), we set the curvature EMA to 0.95 and did not tune on advice from the library's author. KFAC allows computational savings by using the same damping parameter and inverse curvature estimate for multiple weight updates, but we set both of these update at every iteration to match the setting of Ours and Ours (Adaptive). For KFAC Kazuki we set the curvature EMA to zero to turn off the moving average. For KFAC Kazuki, the learning rate, momentum and initial damping were tuned and then fixed (i.e. not adapted).
Since the SEN Exact variants were only applied to UCI Energy, we enumerate those settings here rather than the tables. For SFN Exact, the optimal settings of tuned parameters were learning rate of 0.0049 and momentum of 0.0638. The threshold for replacing small eigenvalues was 0.00052 and the replacement constant -4 For SFN Exact (Adaptive), we only need to tune the initial damping, for which the
19

---

Published in Transactions on Machine Learning Research (02/2024)
Table 3: Details our hyperparameter search strategy (Part 1). Ranges shows the search spaces considered for each hyperparameter as uniform range, except those marked log, which are sampled from a log- log-uniform range. Other rows show the optimal hyperparameters chosen to minimise validation loss and the corresponding losses obtained. Note that initial random sceds during tuning losses achieved by each method must be made with care. All values are rounded to three significant figures.
not controlled, comparisons of the
Ours
Number
Batch
Learning
Order of Training
Test
Validation
Setting
Momentum
Damping
Series Terms Acceleration
Range
2
[og 3 9.5] log(10 0.95
log(10 8 8,1]
[1,20]
[0,
UCI Energy
1.79
0.717
2.44e-2
18
8
4.55c-4
8.62e-4
1.29c-3
Fashion- MNIST
0.305
0.428
5.77e-2
13
2. 2.15e-1
3.72n-1
3.38e-1
SVHN
200 3200
0.108
0.327
1.87e-1
0
9.01e-2
5.95e-1
5.20e-1
CIFAR-10
0.562
0.115
2.44e-3
1.37
1.51
1.52
Ours (Adaptive)
Number of Series Terms (K)
Batch Size
Initial Dampiug
Order of Acceleration
Training
Test Loss
Validation
Setting
Loss
Loss
Range
50 x 2 21.77
log[10 -5 .110
[1.20]
[0.
K-1
UCI Energy
1.46e-6
9
7.15e-3
8.17e-3
8.45e-3
Fashion-MN
800
9.33e-2
14
2
2.45o-
3.65e- 1
3.35e-1
SVHN
100
4.14e-7
4.48o-1
7.10e-1
6.49e-1
CIFAR-10
3200
1.04e-1
5
2
8.91e-1
1.56
1.35
SGD
Learning Rate
Training Loss
Test
Validation
Setting
Size
Loss
Loss
Range
50 x 20-77
log[10-* 10
UCI Energy
9.45e-2
7.15e-3
8.17e-3
8.45e-3
Fashiou-MNISTT
50
1. 1.88--2
3.81e-1
3.69e-1
3.42e-1
SVHN
1.33e-2
2.30e-1
4.87e-1
4.47o-1
CIPAR-10
3200
4.11o-3
1.01
1.57
1.56
Adam
Batch Size
Learning
Training
Test
Validation
Setting
1 B1
1 Ag
Rate
Losss
Loss Loss
Range
211.7| log(10 6, 100) log(10 10, 100)
Iog[10 8, 100)
log[10 4, 100)
UCI Energy
1.71e-2
2.73e-10
0.255
0.0156
3.99e-1
8.01e-1
7.83e-4
Fashion- MNIST
200
6.55e-4
5.99e-3
0.475
0.0195
2,33e-1
3.61e-1
3.31e-1
SVHN
1600
1.84e-3
3.35e-6
0.126
0.0253
1.84e-1
5.50e-1
5.1le-1
CIFAR-10
800
7.70e-5
3.61e-4
0.002 66
0.000 100
1.40
1.41
1.45
optimal value was 0.0125. We set the decay rate of the exponential moving average 0.95, as for KFAC (DeepMind).
In the tables below, Ranges shows the scarch spaces considered for cach hyperparameter as a uniform range, except those marked log, which are sampled from a log-uniform range. We sampled 100 random configurations for each algorithm and dataset combination. For Adam, we tuned 1 1- B1 and 1- 32 using the ranges below and then computed the corresponding values for 31 and f2. We show the optimal hyperparameters chosen to minimise validation loss and the corresponding losses obtained. Note that initial random seeds during tuning were not controlled, comparisons of the losses achieved by each method must be made with care. All values are rounded to three significant figures.
A.4 KFAC Adaptive Heuristics
Here, we give a more technically specific overview of the key adaptive heuristics deployed by KFAC Martens & Grosse (2015), which we presented in Section 2 and employ in our Adaptive experimental settings.
20

---

Published in Transactions on Machine Learning Research (02/2024)
Table 4: Details of our hyperparameter search strategy (Part 2), comments as for Table 3
KFAC (DeepMind)
Batch Sixe
Initial Damping
Training Loss
Test Losss
Validation Loss
Setting
Range
2
log 10))
UCI Energy
6.97e-5
2.44e.4
1.19e.3
8.41e-4
Fashion- MNTST
3200
5.27e-1
9.98o.2
4.20c-1
SVHN
1600
5.76e-1
2.35e-2
6.72e-1
6.10e-1
CIFAR-10
800
2.17e-1
1.19
1.34
1.30
KFAC (Kazuki)
Setting
Batch
Learning
Momentum
Damping
Training Loss
Test
Validation
Size
rate
Loss
Loes
Range
20.7.
log (10010,,000 10"]
log(10-9, 0.95
log[10-*, 10"1
UCI Energy
9.09e-2
2. 2.64e-1
5.63e-6
2.19e-1
7.58e-1
9.79e-4
Aishiou-MNISS
8.68e-2
4.06e-1
2.20e-1
2.50e-1
3.37e-1
3.27e-1
SVHN
800
8.83e-3
2.22e-1
3.79c-4
1.38c-1
7.25c-1
6.65e-1
CIFAR-10 10
3200
2.29n-2
6.70o-3
2.21o-3
9.01e-1
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 For curvature matrix C and decay factor (0, 1),
more robust to stochasticity in we have
C, + BC + (1-3)C,-,.
(6)
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. When the local approximate model has curvature matrix C, gradient g and our proposed update direction is A, we compute the learning rate n and momentum f by
? Pc
ACC. C(x - -x-1i (xi xt-1) TC(x: xx-1)]
ATC(x,- - X(-1) - g (x: x)-
87 A
=
(7)
Tikhonov damping with Sevenberg-Marquardee style adaptation. KFAČ incorporates two damping terms: 17 for weight regularisation, and X which adapted throughout training using Levenberg- Marquardt style updates (Moré, 1978). The damping constant can be interpreted defining trust region for the update step. When the curvature matrix matches the observed landscape, the trust region is grown by shrinking X and vice versa. This level of "mismatch" is captured by the ratio of the actual change in lass to the change predicted by the locally quadratic model. If the ratio is near one and thus the local quadratic model matches the observed losses well, then the curvature matrix useful approximation to the local landscape and damping is decreased (i.e. the trust region is increased). Conversely, the ratio is far from one (implying the local model is not accurate), the damping increased so that optimisation becomes more SGD-like (i.e. the trust region is reduced),
In notation, if the objective function is f(x), f(x) is our local quadratic estimate of f(x) and we have some adjustment factor w E (0, 1), KFAC updates the damping X by the following rule:
d
Pt
f(xx) f(x(-1)
Pt f(xx) f(x(-1)
Ae+1
w/
Pe otherwise
(8)
A
21

---

Published in Transactions on Machine Learning Research (02/2024)
A.5 Series Acceleration
Recall the Sablonnière (1991) (acel)-acclerrated Wynn e-algorithm (Wynn, 1956a) applied to the series of mlh partial sums Sm gives the recursion
e 0,
Sm
(c) - m++1
(c-2)
(12) +1) Em+1
e(c-1))
(9)
This definition sufficient compute the accelerated series, but a naive implementation requires all the terms to be stored in memory, which rapidly becomes problematic for larger networks. By carefully defining the order in which these terms computed, may substantially reduce the intermediate memory storage required. Such a strategy was outlined by Wynn (1962), but combination of changing conventions and unclear formatting make it difficult to interpret; we present our own derivation of the same process in Algorithm 2.
Wynn e-Algorithm with Samelson Inverse Require: Sequence Po, , P2N and acceleration order N for m +- 0 to 2N do
Algorithm Sablonnière-
,
(0) t Pm.
(-1)
m t 0
end for
m + 0, C e 1
while m 2N do
while m 2 do
cm
m+1
(c22)
(L +1)
-1 6m c(c-1))
Delete mm+
(c-2)
+ - 1, c+1
end while
Delete cm+1
(c-2)
m &- c- c+1
end while (2N)
return 6
22

---

Published in Transactions on Machine Learning Research (02/2024)
Table 5: We provide the median of the training loss, test loss and validation loss for bootstrap-samples from 50 random seeds. Modians values are rounded to four significant figures; errors are rounded two significant figures.
Setting
Algorithm
Training Loss
Test Loss
Validation Loss
Ours
0.006459 + 0.005
0.007087 = 0.0066
0.0104 + 0.0083
Ours (Adaptive)
0.001029 + 4e-05
0.001678 + 9.5e-05
0.002165 + 0.00012
SGD
0.001947 0.0002
0.002361 + 0.00025
0.003191 + 0.00038
UCI Energy
Adam
0.000657 + 4e-06
0.00113 + 8.9e-05
0,001571 + 0.00011
KFAC (Kazuki)
0.5018 + 0.00056
881.8 6.2e--02
904.1 + 8.1e--02
KFAC (DeepMind)
0,000714 1 0.00012
0.006663 +: 0.0024
0.008849 + 0.0031
Exact SFN
0.0005019 +: 2.4c-05
0.0015 + 9.2c-05
0,001994 + 9.7c-05
Exact SFN (Adaptive)
0.0004941 + 1.4e-05
0.001045 + 7.-00
0.001407 + 3.4c-05
LBFGS
0.002619 + 0.00051
0.002929 = 0.0008
0.004161 i 0.00069
Ours
0.237 + 0.0061
0.3691 + 0.00081
0.3423 + 0.00065
Ours (Adaptive)
0.233 + 0.0054
0.3684 + 0.0009
0.3429 + 0.0011
MusHion-MNIST
SGD
0.2762 + 0.0024
0.3693 + 0.00062
0.3436 0.00053
Adam
0.2425 + 0.0054
0.3582 + 0.00086
0.3312 + 0.00036
KFAC (Kazuki)
0.2375 + 0.0019
0.3566 0.0006
0.3328 + 0.00078
KFAC (DeepMind)
0.09567 + 0.0027
0.4423 + 0.0023
0.4222 +: 0.0022
Ours
0.1573 + 0.0049
0.583 0.0033
0.5205 + 0.0021
Ours (Adaptive)
+ 0.4757
0.6901 0.0041
0.6335 0.0036
SVHN
SGD
0.1979 + 0.0036
0.5113 0.0024
0.4503 +-2 0.0024
Adam
0.2082 + 0.0042
0.5029 0,0045
0,4514 0.004
KFAC (Kazuki)
0.09535 0.006
0.7036 0.0047
0.6288 0.0048
KFAC (DeepMind)
0.007294 + 0.00078
0.6854
0.6288 0.0035
Ours
1.275 + 0.009
1.497 0.00098
1.504 0.0019
Ours (Adaptive)
0.7614 0.02
1.595 0.0084
1.596 0.0076
CIFAR-10
SGD
0.9871 + 0.0027
1.597 + 0.017
1.617 + 0.014
Adam
0.01125 + 0.0012
3.803 0.0091
3.761 1 0.012
KFAC (Kazuki)
0.03947 1 0.0011
3.29 1 0.0075
3.226 1 0.013
KFAC (DecpMind)
0.01466 1 0.0012
1.727 1 0.019
1.683 1 0.019
B Additional Results
B.1
Tabulated Results
We present the results from Section 4 in tabular form in Table 5.
B.2
Test and Training Trajectories Per Iteration
In Figures 6 and 7, we present complementary plots for the experiments in Section 4, showing the median training and test losses plotted as a function of weight update steps rather than time.
B.3 L-BFGS Baseline
In this soction, we include plots for UCI Energy with L-BFGS (Liu & Nocodal, 1989) included an additional baseline. Numerical results for L-BFGS are included in Table well.
We used the version of L-BFGS in the JAX library (version 0.3.14). We set the number of optimisation steps in the main loop to 20, the maximum number of function evaluations to 25 and the maximum number of Jacobian evaluations to 100. Since larger values are almost always better for all these parameters, we set. them to the largest values within our hardware constraints that allowed for a comparable runtime.
B.4 Effect of Truncation Length
In practice, we wish to avoid computing many terms from the series approximation to the inverse saddle-free Hessian in Equation (4). We may then ask: how many terms sufficient? Here, investigate that question
23

---

Published in Transactions on Machine Learning Research (02/2024)
101.
109
Ours
Outs
(Adeptive)
Ours (Adeptive!
10
10
Adam
KPAC (Kazuk)
10R
KFAC (Kazukii
10
KFAC (Dcepttind)
KFAC (DecpMind)
Training Loss
Exact SPN
1.0
10
SFN
Loos
101
10a
10a
step
step
(Adsptive)
Ours
SGD
Durs (Adaptive)
Adam
SGD
100
KFAC (Kazuk)
KPAC (DeepMindd)
Training Loss
(Desptiind)
5:000
10-4
4:300
168
10
102
109
30a
101
10"
10'
104
step
Figure 6: Median training (left) and test (right) loss achieved on UCI Energy (top) and Fashion-MNIST (bottom) plotted per iteration of training. Values are bootstrap-sampled from 50 random seeds. Optimal hyperparameters were tuned with ASHA.
24

---

Published in Transactions on Machine Learning Research (02/2024)
Out
Ourt (Acaptive)
106
SGD
(Karuk)
KPAC (Desp
Training Loss
10t
Adam
KFAC (Kazubil
KFAC (DespMind)
1.0
10
10
10a
10
16a
10
10a
102
step
step
OuTs
Outs (Adaptive)
105
S60
3.K.1FF
Adem
(Kszuk)
Training Loss
10-1
2x
(Acaptive!
Adam
10
KFAC (Kazuki
(DespNind)
10*
10
19'
step
step
Figure 7: Modian training (left) and test (right) loss achieved on SVHN (top) and CIFAR-10 (bottom) plottod per iteration of training. Values are hootstrap-sample from 50 random seeds. Optimal hyperparameters were tuned with ASHA.
25

---

Published in Transactions on Machine Learning Research (02/2024)
10a
101
105
Training Loss
Test Loss
10-4
10-9
10a4
SGD
SGD
Adam
Adam
KFAC (DeepMind)
10
KFAC (DeepMind)
10 -1
Exact
LBFGS
LBFGS
1
10-4
100
10a
10a
10a
10
10-1
10
10a
10a
10a
Runtime (s)
Runtime (s)
Ours
Ours
SGD
SGD
KFAC (DeepMind)
KFAC (DeepMind)
0.008
Lowest Training Loss
0.008
LBFGS
LBF GS
Lowest Training
0.006
0.006
0.004
0.004
0.002
0.002
0.000
0.000
Ours
SGD
Adam
KFAC [DeepMind)
Exact SFN
LBFGS
Ours
SGD
Adam
KFAC
Exact SFN
LBFGS
(DespMind)
Figure 8: Median training (left) and test (right) MSEs achieved on UCI Energy by various optimisers including L-BFGS in the full-batch setting, bootstrap-sampled from 50 random seeds. Optimal hyperparameters were tuned with ASHA. The x-axes are wall-clock time the top row, and training iteration in the bottom row; note the log-scaling both cases.
empirically by applying our method UCT Energy, but varying the number of terms, used approximate the scries. As shown Figure 9, we see clear improvement the number of computed terms increasod, but even computing only three for us to reach reasonable loss values. We consider further theoretical justification for this in Appendix D.2.
provides a sufficiently close approximation to the saddle-free Hessian
B.5 Comparison of Series Accelerators
While developing our algorithm, we considered range of series accelerators:
Shanks transformation (Schmidt, 1941; Shanks, 1955), which is implemented by Wynn's ealgorithm (Wynn, 1956a)
Sablonnière (1991) modification of the Wymn e-algorithm
Levin-t transform (Levin, 1972), which we found more stable than the related u and v transforms
Padé approximants (Graves-Morris, 1994)
We also investigated the vector- and topological-p accelerators (Wynn, 1956b; Osada, 1991) and the d(2) transformation (Levin & Sidi, 1981; Osada, 1996), but found these to be markedly less robust, so do not. show results here.
To investigate the relative merits of these accelerators, we randomly populate a 100-dimensional Hessian matrix H with independent. draws from standard normal distribution, from which we compute the exact.
26

---

Published in Transactions on Machine Learning Research (02/2024)
K =3
km3
=5
K k-5
10
-
K
Training Loss
=9
=10
Test Loss
=10
10n7
10-4
10ah
Runtime (5)
Runtime
Figure 9: We consider the effect of varying the number of series terms K used to approximate the saddle- free Hessian in (4). The number of steps varies from three to ten. All other scttings are as in Section 4.1, including the acceleration order N = 1. The results above are boolstrap sampled from 50 different random seeds.
1.0
14
Target
Original
1xShanks
2 2xShanks
Cosine Similarity
2xlevin
0.6
Target
2-Norm
2xPade
Orginaal
IxShanks (Sablonnière)
1xShanks
2xShanks (Sablonnière)
2xShanks
1xlevin
2xLevin
1 1xPadé
0.7
2xPade
1xShanis (Sablonniére)
2xShanks (Sablonnière)
0.0
200
400
600
800
1000
200
400
600
800
1000
Series Steps
Series Steps
Figure 10: Comparison of series acceleration techniques. From a randomly-initialise Hessian, we compute an exact saddle-free update vector (-), then compare its direction (by cosine similarity; left) and magnitude (by 2-norm; right) with those vectors found by our unmodified series (-) and a variety of accelcrators applied to that series, as the latter vary when progressively more series terms considered.
27

---

Published in Transactions on Machine Learning Research (02/2024)
vector V
g random starting point. With this exact target in mind, compute 1 1.00 steps of our approximating series, then apply each acceleration algorithm in turn up to four times, comparing the resulting directional error (by cosine magnitude) and magnitude error (by 2-norm) of our update step. Since the differences between accelerators dwarfed those between different acceleration orders of the same accelerator, we show only the acceleration orders acceleration and its Sablonnière modification are indistinguishable at the scale of these plots.
1,2 in our results (Figure 10). Note that Shanks
From these plots, observe Shanks acceleration, and its Sablonnière modification, to reliably converge faster towards the correct update direction than the other accelerators, although Padé acceleration marginally beats these when we compare update magnitudes. Every accelerator makes progress faster than the original series, though Levin-t acceleration seems insufficiently robust for our purposes.
As no accelerator comes particularly close to our target vector, there clearly remains some improvement be made managing this series' convergence. Noting that the greatest acceleration benefit seen for early series steps, choose to focus this window. Further, since the update magnitude is generally underestimated, prioritise the update direction through the cosine similarity metric, as expect too-short steps in the correct direction to retain stable optimisation behaviour. Subjectively, Sablonnière's ' modification of Shanks' algorithm was slightly more stable in our experiments, so we select this accelerator to use in this paper.
28

---

Published in Transactions on Machine Learning Research (02/2024)
C Detailed Derivations
In this Section, we provide a more verbose derivation of the key results of Section 3.
C.1 Scalar Inverse Square-Root Root Series
The generalised binomial theorem provides a means of writing the quantity (xr + u) as the infinite series
(x +) =
x h
(10)
k=
where the generalisation admits any complex r using the definition
r
r(r 1)(r 2)
(r k + 1)
(11)
k
k!
In particular, we have
?
(-z)a
(12)
k-
(-))( -1)(-} -2) ..3- k 1)
1-1)=22
(13)
k!
(2) +1)(2 +2).... .)...(.. ( k
(-1)2k 2k
(14)
k!
(4) + 1)( 2)
......
2 + k -1) zk
(15)
k!
(1(1+ 2) + 4) ... (2k 1)
26 2
(16)
k!
(2k - 1)!
k!(2k: - 2) 4)k 6). ... (2) 2 2/
(17)
(2k - 1)!
(18)
AI(k 1)(k: 2) - 3) ... (1) 2* 12
(2k- 1)!
(19)
k!(K 1)! 22k-
(2R) k 1
(20)
k!k! 2k 22k 1
2k 22k
(21)
k=0
C.2 Series Convergence
Denote by a the kth term of the summation of (21). For this series to be convergent, it suffices that
lim sup |a,. 1 <
(22)
29

---

Published in Transactions on Machine Learning Research (02/2024)
by the root test. Applying this test to our series yields
lim sup Jan 1 lim sup
2n
(23)
2
72
(2n)
lim sup n-+0c
[22]]
(24)
2
n!n!
(2n))
lim sup n-+00
12
(25)
4
n!m!
(2m)2n
2m
S lim sup n-100
12"1
(26)
4
n2n
4n²
lim sup
|="1
(27)
n
n-+00
lim sup
<
(28)
Thus, for the series to converge, it is sufficient that lim sup 12 <I
C.3 Matrix Extension and Scaling
This series extends naturally to the matrix case by choosing square matrix to substitute for and replacing with the appropriately-sizer identity matrix I. Ideally, we would choose = I-H²² and immediately recover a series expression for (H²) - However, as we will observe, such z will not allow the series to converge for arbitrary H, so we will instead introduce a scaling factor V and write 2 = 1 H2.
The scalar convergence condition lim supx-4o |=-| < generalises naturally to the matrix case. Let |I-I] denote any compatible sub-multiplicative matrix norm that is, one which satisfies ||AB|| < ||
and matrices A, B and vectors This definition includes all
||A||| < | |x|| for all dimeesiionllly-compaaeo
matrix norms induced by vector norms. Then, the convergence condition becomes lim sup || < 1.
Collecting these extensions, we recover the series
2k
(I
(T- H2 K
k )
(29)
22h
k=0
and the convergence condition
| < 1.
lim sup ||="|| lim sup
I H²²
(30)
n=000
00000
Gelfand's formula gives that, for any matrix norm, sum -I- - "1 is equal to the spectral radius of J H.. Since we are working with real, symmetric matrices, their eigenvalues are all real, whence the spectral radius is simply the largest of the absolute values of the eigenvalues of H
Let À be arbitrary cigenvalue of H2. By reference to the cigendecomposition of H2, the corresponding eigenvalue of - H22 for all eigenvalues A of H² that
2 is 1 - X.. Thus, for the spectral radius of I - + H2 to be less than unity, require
-1 - <1
(31)
= 0 < X 2V.
(32)
Now, since 22 positive semi-definite by construction, we have A 2 0, and our implicit assumption of the invertibility of H (and hence H2) gives X * 0, whence we recover À as required. For the upper bound, it
30

---

Published in Transactions on Machine Learning Research (02/2024)
suffices to consider only the largest eigenvalue of H which denote by Amax. We thus secure convergence by the condition
V
2 max
(33)
We would prefer to compute this bound on V without explicit reference to the largest cigenvalue of H which
may be expensive to compute in general. Instead, let Umax be the corresponding eigenvector of Amax- Then, by sub-multillicativity of the matrix norm, we have
H
Ummx
2 H Umaax]
(34)
= || max Umax
(35)
Amax [mmax]]
(36)
H 2 Amax-
(37)
So for convergence of the scries, it is sufficient that, for any sub-multiplicative norm ||-I|:
1 V
H
(38)
2
C.4 Principality of Square Root
Recall that the principal square root of H² is positive semi-definite by construction. The inverse of the principal square root, where it exists, must also then be positive semi-definite. So if the result of our series is positive semi-definite, it must have computed the principal square root.
Consider again our series from (29):
06
2k
(1- (1- *
- V-
(39)
22k
k
k=0
we have that the eigenvalues of I - 1 H2 all fall within (-1,1). However, if we strengthen the bound on V to V > Amaxe reprising the argument of the previous section gives that the eigenvalues of I - H must fall within [0, 1], making - H positive semi-definite. But then (I 1 H²)* is also positive semi-definite for k = 0, 1,2. .... This means our scries is a linear combination of positive semi-definite matrices with positive coefficients, so the sunmation even when truncated a finite number of terms must be positive semi-definite. Thus, our construction has computed the inverse of the principal square root of H², as required, when we the stronger condition
Under our convergence condition V 1 Amax < 1/2 |
> Amax
(40)
(41)
31

---

Published in Transactions on Machine Learning Research (02/2024)
D Algorithm Analysis
D.1 Choice of Scaling Factor
Since V exists only to suitably scale the matrix H², we have some freedom in its choice of value. We hypothesise that a smaller V (representing a tighter fit of our convergence bound) would best mitigate any issues with numerical precision, as this avoids rescaling values more than necessary. Although we also hypothesise, based on results for the scalar series, that larger V would ensure more rapid convergence of the series, subsequent rescaling outside the summation most likely climinates any gains here. Thus, we seek V which satisfies our bound V Amax without excessive computational cost.
| tightly possible, but which may be calculated
A naïve approach is to note that Ir(H²) is the sum of the (guaranteed non-negalive) eigenvalues of H², 80 is certainly an upper bound on the largest. Denoting the dimensionality of by m x m, we then have
tr(H') tr (Ia) s (tr H) It)
(42)
m tr ( m (tr H)²
(43)
tr( S ml trr H)²
(44)
so it suffices to set V = V=m(tr H)². Since the diagonal elements of H are the unmixed second derivatives, we can compute them efficiently by differentiating every element of the gradient vector with respect to its corresponding weight parameter, and thus compute H without explicitly computing H. However, we find this bound be extrem loose in practice, and thus detrimental to performance.
Another approach to lower bound is to note that sub tubtiiuiipliiattiiityy of the matrix norm gives
2 g
H
(45)
H'gll
(46)
|g]|
Since our algorithm already computes H'g g, this allows us to efficiently compute lower bound on V based condition:
H'g
(47)
In practice, the algorithm is initialised with some initial value of V (specifically V = 100 in our experiments) which is then increased Lo [H's] whenever the bound in (47) is violated.
D.2 Justification of the Truncated Series
We have shown that our infinite series (29) converges to the required transformed Hessian, but clearly we will forced to truncate the series to K terms in practical implementation. In this subsection, we informally justify the appropriateness of this truncation.
Restating (4),
1
2k
(I
ig
H
(48)
VV
22k
k
k=
and recalling we denote the kth term of the summation by ak; we have from Algorithm 1 that
4k2
2k
1 4 H2
ap g:
a
-
ak.
(49)
41:2
Now, for == 1,....... we have 1k2-2k convergence. It follows that
1, and we have
||
1 H22
by construction in order to secure
4R
)1
(1- HP2
<1
(50)
32

---

Published in Transactions on Machine Learning Research (02/2024)
for k == 0, Thus, we have ||a++ < ||ax|| for such k as suggested by the convergence property of our series, and we can describe the sequence of terms of the summation to be monotonically decreasing in magnitude. It is thus justifiable to suppose that, if we wish to take finitely many terms of the series, we should prioritise the earlier terms (smaller k)), since these will have the greatest impact on the summation.
To develop further insight into this behaviour, recall we exploited the real, symmetric nature of H eigendecompose it as H QAQ". Substituting this decomposition into our series gives
20
2k
- V H² H g
k
(51)
VV
2
k:
k=0
20
1
(2k)
1 - (1=(QQQT))) g
(52)
VV k=0
22k
k
D0
1
(2k
8 (QQ - )
VV
22k
k
g
(53)
k=0
VV
30
22R
1
2k
k
Q -- A2²
Q h 12
g
(54)
k=0
1
1
(2k)
[ (1
A? V
QTg.
(55)
VV
2
Since I and A are diagonal matrices, this series is actually a parallel combination of independent scalar series, and we can consider each diagonal component individually. For an arbitrary eigenvalue X, this gives
2k
(1- k (1- xxe
Σ 1m
(56)
22k
Now, we specifically chose V to be larger than the greatest eigenvalue magnitude of H2. Since we also assumed H has only positive eigenvalues, we can say 0 < 1 - 1 < 1. This common ratio will be near zero for the largest eigenvalues X, so we will see the most rapid convergence of these components of the series. Similarly, the common ratio will be near unity when A is near zero, so we will see the slowest convergence in these components.
This result allows us to consider the high- and low-eigenvalue components of the transformed Hessian independently. High-curvature directions in the space, indicated by large eigenvalues, will converge relatively quickly, so we expect the earlier terms of the series to be of most use in approximating these curvatures. As k increases, the main contribution of each term is towards progressively smaller cigenvalucs, representing lower-curvalure regions of the space. Thus, the more-impactful higher-curvalure information is addressed predominantly towards the start of the series, so even if we only consider finitely many terms, we can be sure none of the first K terms could more optimally be replaced by later term.
D.3 Convergence and Escape
We follow the proof of Paternain al. (2019) prove that in the neighbourhood of a critical point, our method will converge to the critical point in locally convex directions and move away from the critical point in locally concave directions, Let xe be any critical point, and define the immediate vicinity of x by the closed 3'lball Q = E ||x xell < 88} for some Br. We require the following assumptions:
Assumption 1. Over Q, the loss function f(x) is twice continuously differentiable, and further the gradient g(x) and Hessian Hxx) are Lipschilz continuous. Specifically, there exist constants M, L > 0 such that for any x,y RR.
lg(x) g(y) < M |x yll H(x) H(y) & L X yll
(57)
(58)
33

---

Published in Transactions on Machine Learning Research (02/2024)
The hession H(x) is invertible over Q. Specifically, there erists 8 > O such that > (H(x))) 8 for all x E C and i = 1,,,.... P. This additionally implies non-degenemacy of the saddle point.
Assumption 2.
We note that Paternain et al. (2019) also require Assumption 1, though they assume a weaker form of Assumption 2, namely that the (x())))) > 6 must hold at all local minima and saddle points, rather than in a ball around local minima and saddle points. In practice, applying damping to the Hessian ensures that this assumption holds.
We also assume that our series approximation to the inverted saddle-free Hessian in Equation (4) has converged. We use the notation A to denote the matrix obtained by taking the absolute value of each eigenvalue of A and note that the saddle-free Hessian, is thus written as H(x) = (QΙQT) where QAQ is the eigendecompasition of Hx). Without loss of generality, we shall assume the eigenvalues be arranged in ascending order, such that \ A2 ... S Ap and the ith column of Q is the eigenvector associated with eigenvalue d
Recall the critical point of interest is X We let g+ (x) denote the gradient at x projected onto the subspace of H(xc)'s 's eigenvectors associated with the positive cigenvalues of H(xc). Similarly, let g-(x) denote the projection of g(x) onto the subspace defined by the eigenvectors of Hxxc) associated with negative eigenvalues. We now go on to prove that given the assumptions above, for a point x, that is in the neighbourhood Q of a critical point xc, our method will converge in the subspace corresponding to the positive eigenvalues of (xx)) and escape in the subspace corresponding to the negative eigenvalues of Hxxc). In other words, we show that lg+ (x++1)ll converges to zero and that (x++()++1)] will grow.
Theorem 1. Given Assumptions and 2, suppose that ||x: - xcll < $ []g(())]| where xe is critical point and let H(x), (x) and (x) be defined as above. Let ||-|| denote sub-multiplicatioe norm. Then both the Sollowing inequalilies hold:
Ig++ (x++1)]l S D |
(59)
Ig (x++)) 2 2 [g (xx) D l
(60)
where D +C2+4
Proof- We split the the proof into two cases, one for positive eigenvalues, corresponding to (59) and one for negative eigenvalues, corresponding to (60). We start with the negative case.
Case 1: Negative Eigenvalues
Noting that g(x: + +Axx is an anti-derivative of H(x, + AAx) with respect to 0 we can write
g(x,+1) g + /. H(x, + Ax.x)xxd.
(61)
where Ax= = X,+1 Xx. Now, the update rule of our method is given by
X++1 = -
(62)
so that -(xx))-g(xx)) = Ax. Using this fact, we note that g(x,) H(x) -H(x) We
add and subtract g(x,) from (61) as follows:
g(x+11)= = g(x,) + +g(x)) g(x,) + / H(x, + Axx)
j
= 2g((x) + / H(x, + 0Ax) do
-2g(x;) + 1. (H(x, + +0Ax) + H(x))) Ax
(63)
34

---

Published in Transactions on Machine Learning Research (02/2024)
We continue in the manner of Paternain et al. (2019) to add and subtract H(xc), H(x:), and H(x)) inside the integral and shuffle the terms to arrive at:
g(x++1) 2g((x)) [' (H(x; + +Ax) - H(xx)) Ax de
+ ! (H(x)) - H(xx)) Ax + / (H(x)) - H(xx)) Ax
+ [/ (H(xc) + H(x))) Ax
= -2g(x;) + /' (H(x; + 0Ax) - H(xx)) Ax do
+ (H(x)) - H(xc)) Ax + (((x)) H(xc))Ax + (H(xc) + H(xc) Ax.
(64)
Let Q- denote the matrix of cigenvoctors corresponding negative cigenvalues of H(xc). We pro-multiply the left and right of Equation (64) by Q__ and consider each of the last four terms separately.
For the integrand, note that Q since the columns normalised eigenvectors. Moreover, since ((x) is Lipschitz by Assumption 1, we have that [[((x, 0Ax) - H(xx)l < L |xx + AAx - x.ll = L0 Ax so that
Q + 0Ax) H(x:)) Ax| < [[x + 0Ax) H(x))l | I 0L |Ax||
(65)
We handle the next two terms a similar way, applying the Lipschitz assumption:
Q (H(x)) H(xc)) Ax| < (x;)) H(x))|l || S ||x xc || ||xl||
(66)
Q H(xc)) Ax|l < [(x)) H(x)) |
- L ||x: - xcll ||
(67)
Finally, we show that the last term in (64) becomes zero. Using the eigendecompositionn of H(x)), we observe that
(H((x) + H(xc)) (QA(x))) QA(xc))7) QQ[[))=)++++(C+C+CCCC)
Suppose there are d negative eigenvalues. Then Q_Q [L., 0axm-d] i.e. the eigenvectors of Q corresponding to positive eigenvalues mapped to zero, and those corresponding to negative eigenvalues are mapped to unit basis vector. This is because the columns of Q are orthonormal, so the inner product of columns is unily if the columns are equal and zero otherwise. Furthermore, A(xc) + A(xc) is diagonal where the first d elements are zero and the remaining elements double (due to negative eigenvalues cancelling out with their positive counterparts in A(xc) and positive eigenvalues being added to their positive counterparts in A(xc)). But then the product Q (A((() + A(xc)) 0, because the zero components of each term complement each other.
We recall the following identity from the reverse triangle inequality: ||a + bll = ||a (-b)1l 2 ||l|||
|bb||
||a]| - ||b]| and combine it with (63), (66) and (67) to lower bound Equation (64) as follows:
||g (x |l 2 2 lgg (x,) I. ||x
e dd 2L 1|x, x ||x|]
(68)
We use the definition of the update step to bound | || as follows:
X x; H(x,) g(x;) so that
Ax| < [[x,) -1 [x)
[QAQ lg(()))l
< A ||((x))]]
(69)
35

---

Published in Transactions on Machine Learning Research (02/2024)
is bounded because max A(A > 3 by Assumption 2 For |1-11: this bound is In the general case, we denote the bound by Cs (where Cs may also depend on the dimensionality of the problem for some choices of [1-I1)- This, along with our assumption that x, is near x gives us the final bound:
Now:
LC2 llg(x))ll 2LC&3 llx(x))ll
llg (x++1)l 2 2 2 (x)]l
I.C, + 4LC,3 llg(x))ll
2 2 -x(g(xx)]l
(70)
2
Case 2: Positive Eigenvalues
As in the negative case, start with
g(x++1) = g(xx) + / H(x; 0Ax) d0.
(71)
This time, we substitute g(x,) = x ((xx)=((xx))) g((x)) = H(xx) directly to obtain
g(x++1) = / (H(x; + +AAx) H(x,)) Ax
(72)
We proceed in the negative case, adding and subtracting H(xc), H(x)), and H(xc) to obtain
g(x++1) = //' (H(x; + +AA) H(x:)) Ax + (H())) H(xc))
+ (H(xc) - H(x,)())))))))))) Ax (H(x)) H(xc) Ax.,
(73)
noting that the last two lerms are different to the negative case. Let Q denote the matrix of eigenvectors corresponding to positive eigenvalues of H(xc). Multiply the left and the right hand side of Equation (73) by Q and apply the triangle equality to obtain
gg (x++1)| - 1//- (H(x, + 0Ax) H(x))) Ax de
+ +1 (H((x))) H(x))) Ax
+ (H((x) H(xx)) Ax|| + ((obc)-HH(x)Aas Ax||
(74)
Using Equations (65), (66) and (67) as in the negative case, we can bound the first three terms. The bound the last term follows similar reasoning before:
Q (H(xc) H(xx)) ([(x)))-(xx)))=----- -
(75)
where Q.Q(A - A(xc)) = 0 because Q Q maps the eigenvectors of Q that correspond to negative eigenvalues to zero and A(xc) A(xc) produces diagonal matrix where the positive eigenvalues cancel out and the negative eigenvalues double.
We thus arrive at the bound
LC2 2
gg+ (x++111
lg(x;))l 2LC53 [lg(x))]1
LC3
4LCs3 ||g(x)))?
(76)
"While do not consider this work, note that the of canonical second-order damping methods, which replace curvature matrix by and thus increase every eigenvalue by X, allows to relax Assumption hold for the damped (saddle-free) Hessian, and thus admit arbitrary Hessians suitable choice
36

---

Published in Transactions on Machine Learning Research (02/2024)
Reprising the arguments in Corollary 3.3 and Proposition 3.4 of Paternain et al. (2019), (59) gives that lig (x++1)ll converges quadratically to zero if the greatest contribution to [g(x))ll = (x,+ g (x+] is from the (x,) term (as for a local minimum), and (60) gives that 1x-(x))]] grows by multiplicative factor 2 2-D D (where we may choose the free parameter D such that 0 < D < 1) if llg(())]l is negligible compared to (gg((x)) |. In combination, these results justify our claim to converge to local minima and repel saddle points.
D.4 Rate of Convergence
In this subsection, we provide a brief analysis of the rate of convergence of our algorithm to critical points, following similar proof pattern to that of classical Newton methods. Throughout, we will assume every term of our modified- Hesssian summation used, such that our Hessian transformation is exact.
For brevity, denote by H the matrix obtained by taking the absolute value of every eigenvalue of H. With this shorthand, recall the exact version of our update rule is
X++1 u(x:) x H(xx))'g(x))
(77)
where we will now explicitly denote the points at which the Hessian H and gradient g are calculated.
Let x be an arbitrary critical point of the objective function f, such that g(xc) 0. The latter fact gives u(xc) xc, and thus x, xe u(x;-1) u(xc).
Consider taking a Taylor expansion of our update rule u(x) about xe:
u(x:) = u(xc) + Vu(xc)] (x: (xx=xc)+p - xc) VVu(xx)(xx -
-xc) xc) +...
(78)
Denote by = X the error between our critical point and Note this is unrelated to any discussion of Wynn's e-algorithm (Wymn, 1956a); we have chosen to reflect standard notation by overloading here.
Now, by direct differentiation of (77), we have
Vu(xc) = I VH(xc) g(xc) - -H(xc) H(xc). (H(x))
(79)
Noting that H and H have the same eigenvectors, we may gain insight into the final product by eigendecom- posing
H(xc) (ixc))"((xx)) = QA QQQQQ =Q 'QQ..
(80)
(81)
Since [^] and are corresponding eigenvalues of H and H, the result of this product is a matrix with the [X] = sign X. Consequently, we recover different dynamics for positive and negative eigenvalues equivalently, positive and negative curvatures in the space (recall our assumption of invertibility of H provides À + 0).
same eigenvectors as H, but with eigenvalues A
We proceed to analyse cach case individually, effectively ereating two complementary subspaces of the optimisation space. We will use the subscripts + and - to denote the positive- and negative-curvature subspaces, respectively. Note that the orthogonality of these subspaces (ensured by the real, symmetric nature of H giving orthogonal eigenvectors) justifies our independent analysis.
For the positive-curvalure subspace, A4 = -A, whence - (xc) =I This gives
(x4 --
(82)
which collapses our Taylor series to
u (xx) = u4 (xc))) +
2 c, VVu, (xc) t...
(83)
u. (x,) - u4 (xc) = c+11+ = =((,,))
(84)
37

---

Published in Transactions on Machine Learning Research (02/2024)
where by O(e)+) mean to indicate that the positive-stbspace error between our current point and critical point varies quadratically with time, the truncated terms are in higher-order products of c,
We on to repeat this argument for the negative-curvature subspace where A_ = -A-, so that
H (xc) " (xc)= = -I. Then, we recover
(xc== 21,
(85)
which collapses our Taylor series in a different way;
1
(xx) = (xc) 26,,- 2 ct VVu-(xc)e, +
(86)
(x,) u (xc) 1,+1,- 2e, + O(E
(87)
that is, that the negative-subepace error between our current point and a critical point diverges exponentially with time.
This derivation proves that, over time, algorithm will converge to some critical point xc. But derivation in Appendix D.3 shows that our algorithm escopes from non-degenerate saddle points and local maxima, and is attracted to local minima. Thus, any convergence to critical point must be to local minimum; the results of these two sections combine to give that our algorithm converges quadratically along positive-curvature directions and escapes exponentially from negative-curvature directions.
38


Reviews:
[{"review_id": "7MniX9t21e", "reviewer": "Reviewer_Ji7s", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The paper is very well written. While the contribution is a bit incremental, it is elegant, and well motivated. The paper does a good job of summarizing other people's work. The authors present many different algorithmic considerations such as various hyperparameter optimizations and other accelerations that leads me to believe the numerical work is thorough.\n\n\nWeaknesses: The idea is somewhat incremental, and essentially about approximating Hessians. I am not sure that a sequence of matrix-vector products is a practical means of approximating Hessians in modern machine learning computing. Can the mat-vecs all be done simultaneously? This would port better to GPUs and make a better argument for using an idea like this (which I would be more in favor of if these HPC considerations could be commented on).\n\nThe authors claim to essentially solve instability issues of earlier SFN works, but I do not totally follow the argument. Hessians of large neural networks are typically rank degenerate and require damping or some other means for their (approximate) inversion. The authors claim that the inverse of the absolute Hessian can be computed from a series expansion and this avoids the pitfalls of these earlier methods that are due to damping and low rank truncation. From my understanding of Algorithm 1 the authors are not using damping, which to me would seem to lead to instabilities, since the Hessian matrix is itself low-rank. However the appendices to list damping parameters for the method, so perhaps this can be articulated better in the main body to clarify this. \n\nThe numerical results should report generalization accuracies for the classification problems so that readers can understand if the regime the results are in is actually competitive for these large scale ML problems. For example, its well known that second order methods have better stability properties with larger batch sizes, this however often defeats the purpose as it leads to worse generalization accuracies. Small batch sizes tend to lead to better generalization accuracies, but also introduce additional statistical sampling errors that can cause instability in Newton methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes computing a saddle free Newton type update via a matrix series expansion of the absolute Hessian instead of using direct matrix factorizations (e.g., Lanczos or eigenvalue decomposition). Both the proposed methods of this paper, as well as previous cited works can approximate the absolute Hessian matrix free where Hessian-vector products can be easily computed in automatic differentiation frameworks as $H(w)v = \\frac{d}{dw}(g(w),v)$\n\nThe method is compared against other optimization methods such as Adam, SGD, KFAC, and other saddle free Newton methods, and shows comparable performance. \n\nThe appendices of the paper analyze the algorithm, giving a mathematical justification to the idea. Additionally convergence of the method and saddle point escape is analyzed.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would like to see a few things:\n\n1. If possible some analysis of the rate of approximation of the Hessian by series expansion would greatly aid the central contribution of this work. And if this is infeasible, showing some numerical plots of this would be very interesting. This would really help illustrate the power of the method and make a much stronger case for publication.\n\n2. I suspect that the optimal hyperparameters (e.g., large batch sizes) are masking some potential issues with the method. These issues are endemic to many Newton methods and not unique to this method. If the series expanded Hessian truly does solve the instability issues of earlier SFN type methods then the authors must demonstrate this numerically in the regimes where those methods had trouble: namely smaller batch sizes. These are also the typical setups for achieving good performance on datasets like CIFAR10 with architectures like ResNet. Please report generalization accuracy as a percentage, as this is the standard way of assessing the quality of a trained classifier. If the generalization accuracies are where they should be (e.g. >90%), this will give the reader more confidence in the method being practically useful. See for example this link: https://github.com/kuangliu/pytorch-cifar.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "GgShhgTCE1", "reviewer": "Reviewer_DFUg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Let me begin by highlighting the strengths:\n* I like the idea of $|A|=\\sqrt{A^2}$; it is well-motivated as it helps to not be attracted to the first-order stationary points. \n* The experiments are performed and presented in the professional way with multiple seed fine-tuning and providing detailed specifications, which should help to the reproducibility of the results. \n\nHowever, I would like to point out some weaknesses and limitations from my perspective:\n* The main problem for me is the absence of a comparison with the main second-order methods for non-convex optimization: cubic regularized Newton [1-5] and Trust-region Newton [7-8]. Both of these methods have a long history of research, they are both probably optimal for non-convex functions. Both of them have efficient first-order subsolvers: Trust-Region-CG [9] and Cubic Newton + GD [6]. Both of them could escape from the first-order stationary point.  From this point, some claims in the paper are too strong and misleading, for example: “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”. \n* There is no convergence theorems, which would show the convergence rate of the method.\n* From the practical point of view, it seems that the method has too many parameters and additional heuristics to perform comparably with the much more simple methods. “Our implementation of Algorithm 1, using tuned learning rate, momentum, series length K and order of acceleration N”. 4 parameters in total. In the main figures 2-3, the fine-tuned 4 parameters method with momentum and Hessian-vector information is losing to SGD with 1 parameter. This questions the practical performance of the proposed method for me. \n\n[1] - Nesterov, Yurii, and Boris T. Polyak. \"Cubic regularization of Newton method and its global performance.\" Mathematical Programming 108.1 (2006): 177-205.\n\n[2] - Tripuraneni, Nilesh, et al. \"Stochastic cubic regularization for fast nonconvex optimization.\" Advances in neural information processing systems 31 (2018).\n\n[3] - Cartis, Coralia, Nicholas IM Gould, and Philippe L. Toint. \"Adaptive cubic regularisation methods for unconstrained optimization. Part I: motivation, convergence and numerical results.\" Mathematical Programming 127.2 (2011): 245-295.\n\n[4] - Wang, Zhe, et al. \"Cubic regularization with momentum for nonconvex optimization.\" Uncertainty in Artificial Intelligence. PMLR, 2020.\n\n[5] - Wang, Zhe, et al. \"Stochastic variance-reduced cubic regularization for nonconvex optimization.\" The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.\n\n[6] - Carmon, Yair, and John Duchi. \"Gradient descent finds the cubic-regularized nonconvex Newton step.\" SIAM Journal on Optimization 29.3 (2019): 2146-2178.\n\n[7] - Conn, Andrew R., Nicholas IM Gould, and Philippe L. Toint. Trust region methods. Society for Industrial and Applied Mathematics, 2000.\n\n[8] - Nocedal, Jorge, and Stephen J. Wright, eds. Numerical optimization. New York, NY: Springer New York, 1999.\n\n[9] - Curtis, Frank E., et al. \"Trust-region Newton-CG with strong second-order complexity guarantees for nonconvex optimization.\" SIAM Journal on Optimization 31.1 (2021): 518-544.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In the paper, the new second-order optimization method is presented for non-convex optimization with applications to neural networks. The main idea of the method is using that $|A|=\\sqrt{A^2}$ and applying the series technique to compute its inverse. The extensive experimental results are presented for big variety of optimization setups.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* I would recommend adding the clear contribution subsection with the main claims from the paper.\n* I would recommend adding Cubic Regularized Newton and Trust Region Newton to the literature review.\n* I would recommend adding some versions of Trust-Region-CG and Cubic Newton + GD to the practical comparison.\n* In Figure 3, KFAC(Kazuki) seems to ruin the visualization of the results, it is clear that it somehow diverged, but it took almost the whole space in Figure 3.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Rdv8bC5vZs", "reviewer": "Reviewer_2zn1", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths：\n\nThe paper maintains a smooth flow, complemented by clear illustrations, rendering it easy to read.\n\nThe methods employed to approximate the infinite series for obtaining reasonable updating directions are commendable.\n\n\nWeaknesses：\n\nWhile the paper concentrates on proposing saddle-free algorithms, there is limited discussion on other saddle-free algorithms that utilize curvature information without necessitating eigenvalue decomposition, e.g., [1], [2], [3]. Especially, the experiments lack thorough comparisons in this regard.\n\nThe performance of the proposed algorithm does not consistently outperform the compared algorithms.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a novel second-order approach to escaping saddle points. Unlike Saddle-Free Newton (SFN) methods, this approach computes updating directions by leveraging the inverse of the squared and principal square-rooted Hessian through an infinite series. Notably, it achieves this without necessitating eigenvalue decomposition of the Hessian matrix. The authors then introduce a truncated Hessian-free algorithm, utilizing series acceleration to approximate these directions. Experimental results demonstrate that the proposed algorithm performs comparably to SGD, Adam, KFAC and the exact SFN.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "LiSSA [4] also utilizes an infinite series to compute second-order updating directions, subsequently approximating them via truncation. Some discussion on this approach would be beneficial.\n\nThe convergence rate of Lanczos methods is well-documented [5] and can be enhanced by applying reorthogonalization for numerical stability. Additionally, efficient Hessian-free Newton methods like Newton-CG [2] and Newton-MR [3] variants employ negative (non-positive) curvature detections with complexity guarantees, enabling efficient escape from saddle points. It would be valuable to include discussion on these methods.\n\n[1] https://proceedings.mlr.press/v70/carmon17a/carmon17a.pdf\n\n[2] https://link.springer.com/article/10.1007/s10107-019-01362-7\n\n[3] https://arxiv.org/abs/2208.07095\n\n[4] https://www.jmlr.org/papers/volume18/16-491/16-491.pdf\n\n[5] https://www-users.cse.umn.edu/~saad/eig_book_2ndEd.pdf", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces an approximate Newton method that makes uses a Taylor series expansion of the matrix square root. This avoid issues with negative eigenvalues, and it does not require computing an SVD. 

Furthermore the paper is well written, the claims are clear and substantiated, and now with the expansion of the background to include cubic Newton, trust region methods, saddle-avoiding methods, Lissa ..etc I also find the background to now be good. I agree with DFUg that reporting test accuracy (instead of loss) would be helpful. But since you make no claims of having a SOTA test accuracy on these benchmakrs,  I think this omission is reasonable. 

Overall, because of this novelty, clarity and good background, I recommend the paper be accepted.

I also find reading Figure 2 difficult because of the divergence of one of the KFAC method. May I recommend that you make use of extra space in the camera ready to duplicate this figure, one with and one without this divergence KFAC method
Decision: Accept (poster)