ISAAC Newton: Input-based Approximate Curvature for Newton's Method
Authors: Felix Petersen
Tobias Sutter
Christian Borgelt
Dongsung Huh
Hilde Kuehne
Yuekai Sun
Oliver Deussen
Venue: ICLR 2023 poster (2023)

Abstract: We present ISAAC (Input-baSed ApproximAte Curvature), a novel method that conditions the gradient using selected second-order information and has an asymptotically vanishing computational overhead, assuming a batch size smaller than the number of neurons. We show that it is possible to compute a good conditioner based on only the input to a respective layer without a substantial computational overhead. The proposed method allows effective training even in small-batch stochastic regimes, which makes it competitive to first-order as well as second-order methods.

ISAAC NEWTON: INPUT-BASED APPROXIMATE CURVATURE FOR NEWTON'S METHOD
Felix Petersen 12 Tobias Sutter?, Christian Borgelt), Dongsung Huh
Hilde Kuehne Yuekai Sun", Oliver Deussen"
Stanford University. University of Konstanz, 3 University of Salzburg,
*MIT-IBM Watson AI Lab, University of Frankfurt, University of Michigan petersenEcs. st stanford.edu
ABSTRACT
We present ISAAC (Input-baSed ApproximAte Curvature), novel method that conditions the gradient using selected second-order information and has an asymp- totically vanishing computational overhead, assuming batch size smaller than the number of neurons. We show that is possible to compute a good conditioner based only the input to respective layer without substantial computational over- head. The proposed method allows effective training even in small-batch stochastic regimes, which makes it competitive to first-order as well second-order methods.
INTRODUCTION
While second-order optimization methods are traditionally much less explored than first-order methods in large-scale machine learning (MI.) applications due to their memory requirements and prohibitive computational cost per iteration, they have recently become more popular in MI mainly due their fast convergence properties when compared to first-order methods (T). The expensive computation of an inverse Hessian (also known as pre-conditioning matrix) in the Newton step has also been tackled via estimating the curvature from the change in gradients. Loosely speaking. these algorithms are known as quasi-Newton methods; for a comprehensive treatment, see Nocedal & Wright (2). Various approximations to the cre-conditinn matrix have been proposed in recent literature 3 (.). From theoretical perspective, second-order optimization methods are not nearly as well understood as first-order methods. is an active research direction to fill this gap (7. 88)
Motivated by the task of training neural networks, and the observation that invoking local curvature information associated with neural network objective functions can achieve much faster progress per iteration than standard first-order methods -- several methods have been proposed. One of these methods, that received significant attention, is known as coonecker-factore Approximate Curvature (KFFCC) [12], whose main ingredient is sophisticated approximation to the generalized Gauss-Newton matrix and the Fisher information matrix quantifying the curvature of the underlying neural network objective function, which then can be inverted efficiently.
Inspired by the K-FAC approximation and the Tikhonov regularization of the Newton method, we introduce novel two parameter regularized Cronecker-factorize Newton update step. The proposed scheme disentangles the classical Tikhonov regularization and in a specific limit allows us to condition the gradient using selected second-order information and has an asymptotically vanishing computational overhead. While this case makes the presented method highly attractive from the computational complexity perspective, we demonstrate that its empirical performance on high-dimensional machine learning problems remains comparable to existing SOTA methods.
The contributions of this paper can be summarized as follows: (i) we propose a novel two parameter regularized K-FAC approximated Gauss-Newton update step: (ii) we prove that for an arbitrary pair of regularization parameters, the proposed update direction is always a direction of decreasing loss; (iii) in the limit, one regularization parameter grows, obtain efficient and effective conditioning of the gradient with an asymptotically vanishing overhead; (iv) we empirically analyze the method and find that efficient conditioning method maintains the performance of its more expensive counterpart; (v) we demonstrate the effectiveness of the method in small-batch stochastic regimes and observe performance competitive to first-order as well quasi-Newton methods.

---

2 PRELIMINARIES
In this section, we review aspects of second-order optimization, with a focus on generalized Gauss- Newton methods. In combination with Kronecker factorization, this leads to new regularized update scheme. We consider the training of an L-layer neural network fz; f) defined recursively as
24 - a--1i( (pre-activations),
ag (Z,) (activations),
(1)
where dn the vector of inputs and az = [ax; 0) is the vector of outputs. Unless noted otherwise, we assume these vectors to be row vectors (i.e., in R1x*) as this allows for a direct extension to the (batch) vectorized case (i.e., in Rbx=) introduced later. For any layer i. let W E 1 be weight matrix and let & be an element-wise nonlinear function. We consider convex loss function Cly. that measures the discrepancy between y and y/. The training optimization problem is then
arg min Evy, (=; v),
where 0 = [0(2) ,..., ,...., 0(4) with g(0) vec(W()).
The classical Newton method for solving (2 is expressed as the update rule
0' 0 - H (o 0).
where rj >> 0 denotes the learning rate and Hg is the Hessian corresponding to the objective function in ( The stability and efficiency of an estimation problem solved via the Newton method can be improved by adding a Tikhonov regularization term [33 leading regularized Newton method
8' 17 HHg + AI) - (o ()
where À is the so-called Tikhonov regularization parameter. It is well-known [14] [15] that under the assumption of approximating the model ff with its first-order Taylor expansion, the Hessian corresponds with the so-called generalized Gauss-Newton (GGN) matrix Gg. and hence can be expressed as
0' 0 1j (Go + Vv 0),
major practical limitation difficulty is known as Kronecker-Factton the block-diagonal (i.e., liser empirical Hessian or GGN matrix. Inspired by K-FAC, there have been other works discussing approximations of Ge and its inverse (15). In the following, discuss a popular approach that allows for (moderately) efficient computation.
5 is computation of the inverse term. A method that alleviates this Approximate Curvature (K-FAC) [12) which approximates
The generalized Gauss-Newton matrix G3 is defined
G. E (J+ff 6)) [ Jaf
(6)
where J and 2 denote the Jacobian and Hessian matrices, respectively. Correspondingly, the diagonal block of Gg corresponding to the weights of the ith layer W
Gwo w (x;; VzL((( 0), Jwro [(;a
According to the backpropagation rule J Jwi ((; 0) = fl; a.-1, T = a @ b, and the mixed-product property, can rewrite Gwtt
Gw
=B (JJ 2. I 0) a. (Ft.
a).,
(7)
=B [[g aa-1) 1-11 ] [6 a-1) (-) E s (-]
(8)
where
g = (J=,(( (V) ())
(9)
Remark (Monte- Marlo Low-Rank Approximation for g g). As g is a matrix of shape m x d; where m the dimension of the output of f. 9 is generally expensive compute. Therefore, [12] use a low-rank Monte-Carla approximation to estimate V2,((f 0),3) y) and thereby ğ g. For this, we need to use the distribution underlying the probabilistic model of our loss C (e.gg. Gaussian for MSE loss, categorical distribution for cross entropy). Specifically, by sampling from this distribution
2

---

Pf(x)) defined by the network output fx; 0). get estimator of Vff 0), via the identity
0), v) ()
(x ( (L)((() (,C(( V)C(((=(0),))
(10)
An extensive reference for this (as well as alternatives) can be found in Appendix A.2 of Dangel et al. 15). The respective rank-1 approximation (denoted by A of (()
()C(((1;0),)) & v,L(((x; 1) v,
where y ~ (z()). Respectively, we can estimate 9 9 using this rank-/ approximation with
g - (Jz, (x: () v., ((;
(11)
In analogy to g. we introduce the gradient of training objective with respect to pre-activations 2; as g = (J=, ( V,C({(x;0),)) V. ((fx(())))
(12)
In other words, for given layer, letg E Bixd, denote the gradient the loss between an output and the ground truth and let Rmxx: denote the derivative of the network f times the square root of the Hessian of the loss function (which may be approximated according to Remark [ each of them with respect the output 2; of the given layer i. Note that ş is not equal to g and that they require one backpropagation pass each (or potentially many for the case of g). This makes computing 9 costly.
Applying the K-FAC [12] approximation (8) the expectation of Kronecker products be approximated as the Kronecker product of expectations as
= (g 9) @ (a aa))
= E(g g) & B(aTa).
(13)
where, for clarity, we drop the index of a in 8 and denote it with similarly we denote Gwe as G. While the expectation of Kronecker products generally not equal the Kronecker product of expectations, this K-FAC approximation (3) has been shown to be fairly accurate in practice and to preserve the "coarse structure" of the GGN matrix [12]. The K-FAC decomposition in (13) is convenient as the Kronecker product has the favorable property that for two matrices A, B the identity (A & B) A- 1 @ B-1 which significantly simplifies the computation of an inverse.
In practice, E(g g) and IB(a a) can be computed by averaging over a batch of size as
B(g T g) = - g//g/b
[(a = a a/ab
(14)
where denote batches of g. 9 and as g E bx g E Rbbb and E gbxd;-1, where our layer has d inputs, d; outputs, b is the batch size, and r is either the number of outputs m or the rank of an approximation according to Remark Correspondingly, the K-FAC approximation of the GGN matrix and its inverse are concisely expressed as
G
(a a))/b²
G-
gg
9 2 (aaaaaa
(15)
Equipped with the standard terminology and setting, we now introduce the novel, regularized update step. First, inspired by the K-FAC approximation 333, the Tikhonov regularized Gauss-Newton method (5) can be approximated by
000r (00) "( g/b AI) & (a a/b AI) -V poc [(I(x;))) 0)),
(16)
with regularization parameter
key observation, which is motivated by the structure of the above update, is to disentangle the two occurrences of À into two independent regularization parameters Ag. Aa > 0. By defining the Kronecker-factorizeed Gauss-Newton update step as
Na g/ + A²1) 3 (a T a/b + AaI) 1V oco (()(x)))),
(17)
we obtain the concise update equation
0 0(0) S'¹5.
(18)
This update (88 is equivalent to update (16) when in the case of 1)' Ae and A A= Ag Aa. This
equivalence does not restrict 17", .dg. Aa in any way, and changing Ag X does not mean that change our learning rate or step size v7". Parameterizing Ç in (17) with the multiplicative terms AgAn makes the formulation more convenient for analysis.
3

---

In this paper, we investigate the theoretical and empirical properties of the iterative update rule (8 and particular show how the regularization parameters Ag. A affect the cronecker-factoriine Gauss-Newton update step Ç. When analyzing the Kromecker-factorize Gauss-Newton update step ç. particularly useful tool is the vector product identity,
o
& a a)
vec(g a) = vec (g g a aa)
(19)
where the gradient with respect to the weight matrix is g a.
THEORETICAL GUARANTEES
In this section, we investigate the theoretical properties of the croneckerffactriizee Gauss-Newton update direction Ç defined in (17). We recall that Ç introduces Tikonov regularization, it is commonly done in implementations of second order-based methods. Not surprisingly, we show that by decreasing the regularization parameters >g. a the update rule (18) collapses (in the limit) the classical Gauss-Newton method, and hence in the regime of small Xg. aa the variable Ç describes the direction. Moreover, by increasing the regularization strength, we converge (in the limit) to the conventional gradient descent update step.
The key observation is that, we disentangle the regularization of the two Kronecker factors g g and a and consider the setting where only one regularizer is large (Ag oo to be precise), we obtain an update direction that can be computed highly efficiently. We show that this setting describes approximated Causs-Newton update scheme, whose superior numerical performance is then empirically demonstrated Section 4
Theorem 1 (Properties of (). The K-FAC based update step defined in (17 can be expressed as
1 1----g g"
1 bAR
1 bλa
1
5= Im bdg
1b
99
9 ...
Ib
aa
Ib
aa
bda
(20)
Moreover, S admits the following asymptotic properties:
(i) In the limit of Ag, Xa - 0, Je is the K-FAC approximation of the Gauss-Newtord step, i.e. lim Aa +0 Xgda G G-I g(f 0)), where = denotes the K-FAC approximation (5)
(ii) In the limit of Ag. An - $ is the gradient, i.e., lim da -100 V ((;(;;))
The Proof is deferred to the Supplementary Material.
We want to show that Ç is well-defined and points in the correct direction, not only for Ag and An numerically close to zero because want to explore the full spectrum of settings for Ag and Xa- Thus, we prove that Ç is a direction of increasing loss, independent the choices of Ag and )a-
Theorem 2 (Correctness of Ç is independent of Ag and Xa). Ç is a direction of increasing loss. independent of the choices of Xug and Xa-
Proof. Recall that (AgIm +9 g/b) and (ALI. a/b) are positive semi-definite (PSD) matrices by definition. Their inverses (AgIm +9 g/b)- -1 and (AaI, + a a/66-1 are therefore also PSD. As the Kronecker product of PSD matrices is PSD, the conditioning matrix ((*gIm. + g/b)- @ (Aa aaab)) G is PSD, and therefore the direction of the update step remains correct.
This leads us to our primary contribution: From our formulation of Ç, we can find that, the limit for Xg 20, Equation (2T) does not depend on g. This computationally very beneficial computing g is costly as it requires one or even many additional backpropagation passes. In addition, allows conditioning the gradient update by multiplying 6 b matrix between g and a, which is very fast.
Theorem 3 (Efficient Update Direction / ISAAC). In the limit of >g oc, the update step converges to lim < S", where
(I,
S* = g" I,
aa T Ib
aa
a.
(21)
bia
bia
4

---

(i) Here, the update direction S" is based onty the inputs and does require computing (which would require a second backpropagation pass), making it efficient.
(ii) The computational cost of computing the update S" lies in O(bn² + b2n + 12),, where is the number of neurons each layer. This comprises the conventional cost of computing the gradient n 1 b²). V = g x lying in Obbr²), and the overhead of computing 5 instead of V lying in O(b
The overhead is vanishing, assuming n > b. For > n the complexity lies in Obm + n²)).
Proof. We first show the property (1.). Note that according to (22) >g (AgIm + g/b) con- verges in the limit of Ag + oo to Im and therefore (2I) holds.
(i) The statement follows from the fact that the term 9 does not appear in the equivalent characteriza- tion (2T) of 5*
(ii) We first note that the matrix aa is of dimension b, and can be computed in Ob2)n) time. Next, the matrix
I T aa
Ib
aa'
bda
is of shape b x b and can be multiplied with a in O(Pn)) time.
Notably, can be computed with vanishing computational overhead and with only minor modifications to the implementation. Specifically, only the g a expression has to be replaced by (21) in the backpropagation step, As this can be done independently for each layer, this lends itself also to applying it only to individual layers.
As we see in the experimental section, in many cases in the mini-batch regime (i.e., b < n), the optimal (or a good) choice for Ag actually lies in the limit to co. This is surprising result, leading to the efficient and effective 5" = Sxy- optimizer.
Remark (Relation between Update Direction Ç and (*). When comparing the update direction S in (20) without regularization (i.e., Ag + 0, >a + 0) with 5 (i.e., Xg + 20) as given in (21). can be directly seen that Ç" corresponds to particular pre-conditioning of C, since Ç = MC M- bbg 9 g.
As the last theoretical property of our proposed update direction (", we show that in specific networks k" coincides with the Gauss-Newton update direction.
Theorem 4 (S" is Exact for the Last Layer). For the case of linear regression or, more generally, the last layer of networks, with the mean squared error, Ç" is the Gauss-Newton update direction.
Proof. The Hessian matrix of the mean squared error loss is the identity matrix. Correspondingly, the expectation value of g g is I. Thus, 5" =Ç
The direction 5" corresponds to the Gauss-Newton update direction with approxima- tion of G that can be expressed as G se E [I (a a)]
Remark 3.
Remark 4 (Extension to the Natural Gradient). In some cases, it might be more desirable to use the Fisher-based natural gradient instead of the Gauss-Newton method. The difference to this setting is that in 5) the GGN matrix G is replaced by the empirical Fisher information matrix F.
We note that our theory also applies to F, and that Ç also efficiently approximates the natural gradient update step F The i-th diagonal block of (F gco E[ ga) @ (4-1)
has the same form as a block of the GGN matrix G (Geta = E [ g.) o (a] Thus, we can replace ğ with g in our theoretical results obtain their counterparts for F.
4 EXPERIMENTS!
In the previous section, we discussed the theoretical properties of the proposed update directions Ç and Ç" with the aspect that Ç* would actually be "free" to compute the mini-batch regime. In this section, we provide empirical evidence that 5" is good update direction, even in deep learning. Specifically,, we demonstrate that
Cooee will he made available fetthhuhccom//eiii-eeeessseeeeee
5

---

[h]
2.955
2292
2.00
18%
1AM
1072
1AM0
Figure 1: Logarithmic training loss (top) and test accuracy (bottom) on the MNIST classification task. The axes are the regularization parameters 0g and Aa logarithmic scale with base 10. Training with 5-layer ReLU activated network with (left, a, e). 400 (center, b, c, f. g).
(right, h) neurons per layer. The optimizer is SGD except for (c, g) where the optimizer is SGD with momentum. The top-left sector is the top-right column is Ç", and the bottom-right corner is V (gradient descent). For each experiment and each of the three sectors, comparison between the methods; within each sector the learning rate constant. can observe that in limit of Ag + oo (i.e., in the limit to the right) the performance remains good, showing the utility of ("
learning rate, i.c., 5. V have their own learning rate to make fair
(EI) C* achieves similar performance to K-FAC, while being substantially cheaper to compute. The performance of proposed method can be empirically maintained in the mini-batch
(E2)
regime (n >> b).
(E3) Ç* may be used for individual layers, while for other layers only the gradient V is used. This still leads improved performance.
(E4) J* also improves the performance for training larger models such as BERT and ResNet.
(E5) The runtime and memory requirements of 5* are comparable to those of gradient descent.
E1: IMPACT OF REGULARIZATION PARAMETERS
For (E1), we study the dependence of the model's 's performance the regularization parameters Xg and Ag. Here, we train a 5-layer deep neural network on the MNIST classification task [16] with batch size of 60 for total of 40 epochs 40 000 steps.
The plots in Figure [ demonstrate that the advantage of training by conditioning with curvature information be achieved by considering both layer inputs and gradients with respect random samples y. but also using only layer inputs a. In the plot. we show the performance of Ç for different choices of Ag and Aa, each in the range from 10-6 to 106. The right column shows Ç", i.e., Ag oc, for different Aa- The bottom-right corner is gradient descent, which corresponds to >g = oo and Aa 20.
Newton's 's method or the general K- K-FAC approximation corresponds to the area with small Ag and Aa- The interesting finding here is that the performance does not suffer by increasing >g toward co, i.e. from left to right in the plot.
In addition, in Figure 3 we consider the case of regression with an auto-encoder trained with the MSE loss on MNIST [16] and MaSSI [17]. Here, we follow the same principle above and also find that Ç" performs well.

---

Gradient descent
Gradien descent
0.06
0.06
KPAC fa-1-
KMAC 14-011
Training (m3
0.05
0.04
Training
0.04
0--4-011
Me Jg 0.011
0.03
0.03
0.02
0.02
0.02
0.00
0.00
2000
4000
6000
8000
10000
500
1000
1500
2000
2500
3000
3500
4000
training time [s]
Steps
Figure 2: Training loss of the MNIST auto-encoder trained with gradient descent, K-FAC, and Ç", Comparing the performance per real-time (left) and per number of update steps (right). Runtimes are for CPU core
In Figure compare the loss for dif- ferent methods. Here, we distinguish between loss per time (left) and loss number of steps (right). We can ob- serve that, for À = 0.1, K-FAC, Ç. and 5" are almost identical per update step (right), while Ç* is by a large margin the fastest, followed by Ç. and the con- ventional K-FAC implementation is the slowest (left). On the other hand, for 0.01 we can achieve a faster con- vergence than with À 0.1, but here only the K-FAC and ( methods are nu- merically stable, while (" is unstable in this case. This means in the regime of very small X, 5" FAC and Ç. however, achieves good performance with small but moderate À like A = 0.1. For À < 0.01, also K-FAC and Ç become numerically un- stable in this setting and, in general, we observed that the smallest valid À for K-FAC is 0.01 or 0.001 depending on model and task. Under consideration of the runtime, 5" performs best as it is almost as fast as gradient descent while performing equivalent Specifically, a gradient descent step is only about 10% faster than Ç*.
(4)
-10
-5
Figure 3: Training an auto-encoder MNIST (left) and Fashion- model is same used by Botev el al. (18). i.e., it is ReLU-activated 6-layer fully connected model with di- mensions 784 4 - 1100-1000--0 000-500- 30 10-50001100--88 Displayed is the logarithmic training loss.
MNIST (right).
as robust K-
5-
3-
2-
+1-2-3445544
100
19
K-FAC and Ç.
Figure 4: Training 5-layer ReLU network with 400 neurons per layer on the MNIST classification task (as in Figure I but with the Adam optimizer [9
E2: MINIBATCH REGIME
Gradent daccer
Brkkyes
For (E2), in Figure I we can see that training performs well for
for layers
{100, 400, 1 600} neu- rons per layer at a batch size of only 60. Also, in small batch sizes
enur
for loyera
-8
all other experiments, of between 8 and 100.
-10
E3: 5* IN INDIVIDUAL LAYERS
-14
In Figure 5 we train the 5-layer fully connected model with 400 neurons per layer. Here, we consider the setting that we use Ç" in some of the layers while using the default gradient V in other layers. Specifically, settings, where all, the first, the final, the first three, the final three, the odd numbered, and the numbered layers are updated by 5". We observe that all settings with Ç" perform better than
-18
750
1000
1250
1500
1750
Figure Training the MNIST classification task using 5" only in selected layers. Runtimes are CPU.
consider the

---

Table 1: BERT results for fine-tuning pre-trained BERT-Bas Base (B-B) and BERT-Mini (B-M) models the COLA, MRPC, and STSB text classification tasks. Larger values are better for metries. MCC the Matthews correlation. Results averaged over 10 runs.
Method Setting
CaLA (B-B)
CoLA (B-M)
MRPC (B-B)
STS-n (B-M)
Metric
MCC
MCC
Acc.
F
Pearson
Spearman
Gradient baseline
54.20 + 7.56
21.08 2.88
82.32 + 1.22
87.88 0.74
76.98
1.10 76.88 0.79
57.62 +: 1.59
24.67
2.62 83.28 + 0.89
88.28 + 0.70
81.09 + 1.58
80.82 = 1.57
plain gradient descent, except for "Ç" for layers 3,4,5" which performs approximately equivalent to gradient descent.
E4: LARGE-SCALE MODELS
BERT To demonstrate the utility of 5" also in large-scale models, we evaluate it for fine-tuning BERT 20 on three natural language tasks. In T we summarize the results for the BERT fine-tuning task. For the "Corpus of Linguistic Acceptability" (CoLA) [21] data set, we fine-tune both the BERT-Base and the BERT-Mini models and find that we outperform the gradient descent baseline in both cases. For the "Microsoft Research Paraphrase Corpus" (MRPC) [22 data set, fine-tune the BERT-Base model and find that we outperform the baseline both in terms of accuracy and Fi-score. Finally, the "Semantic Textual Similarity Benchmark" (STS-B) B) [23 data fine-tune the BERT-Mini model and achieve higher Pearson and Spearman correlations than the baseline. While for training with CoLA and MRPC, were able use the Adam optimizer (19) (which is recommended for this task and model) in conjunction with Ç* in place of the gradient, for STS-B Adam did work well. Therefore, for STS-B, we evaluated it using the SGD with momentum optimizer. For each method, we performed a grid search over the hyperparameters. We note that we use a batch size of 8 in all BERT experiments.
ResNet In addition, we conduct an experiment where we train the last layer of a ResNet with S", while the remainder of the model is updated using the gradient V. Here, we train a ResNet- 18 (24)
10 [25] using SGD with batch size of 100. In Figure 6 we plot the test accuracy against number of epochs. The times for each method lie within 1 5 of each other. We consider three settings: the typical setting with momentum and weight decay, setting with only momentum, and a setting with vanilla SGD without momentum. The results show that the proposed method outperforms SGD each of these cases. While the improvements are rather small in the case the default training, they are especially large in the case of no weight decay and no momentum.
0.84
Gradintt dencers w nomentam wnight dady C" ast layer 520 memenaar wight docay
0.82
$40
0
25
50
100
125
150
175
200
Epochs
Figure 6: ResNet-18 trained on CIFAR-10 with image augmentation and cosine learning rate schedule. To ablate the optimizer. specifically, without weight decay and without momen- tum. Results are averaged over 5 runs and the standard deviation is indicated with the colored arcas.
additional settings added,
E5: RUNTIME AND MEMORY
Finally, we also evaluate the runtime and memory requirements of each method. The runtime evaluation is displayed in Table 2 We report both CPU and GPU runtime using Py (26) and (for K-FAC) the backpack library [15). Note that the CPU runtime more representative of the pure computational cost, as for the first rows of the GPU runtime the overhead of calling the GPU dominant. When comparing runtimes between the gradient and Ç" on the GPU, we can observe that we have overhead of around 2.5 independent of the model size. The overhead for CPU time also very small at less than 1% for the largest model, and only 1.3 s for the smallest model. In contrast, the runtime of S" is around 4 times the runtime of the gradient, and K-FAC has substantially larger runtime. Regarding memory, Ç* (contrasting the other approaches) also requires only small additional footprint.
8

---

Table 2: Runtimes and memory requirements for different models. Runtime is the training time per epoch on MNIST at batch size of 60. i.e. for 000 training steps. The K-FAC implementation is from the ba backpack library [5]. The GPU is an Nvidia A6000.
Godlient
K-FAC
Model
CPU time GPU time
Memory
CPC time
GPU
Memoey
CPU time
GPU
Memury
CPUs.
GPU
Memory
layers 100
2.055
78
17.
1155M
L6MB
LOMB
layers
23
x
218.48 x
32.00
22.4 MB
38
12.02
7.7
x
4.19
4.9MR
layers w/1 n. a.
187 87
1.93
51.0MD MB
6985.44 48
158 48
212.2
12.53
85.8MB
291.01
4.49
51.4
6 loyes w/ 400 (34899 59
8.22
691.0 MB
1820.81
$165.3 MB
9678
31.87
1197.8 MB
8451.61
10:24
602. MB
Aute-Encoder
78.61
16.2MB
1207.58
74.09
70.7
198.25
33.8B
ST.
4.98
16.5 MB
Remark 5 (Implementation). The implementation of can be done by replacing the backpropagation step of a respective layer by (2T). As all "ingredients" are already available in popular deep learning frameworks, it requires only little modification crniiiionnnottttt (contrasting K-FAC and Ç. which require at least one additional backpropagation)
We will publish the source code of our implementation. In the appendix, we give PyTorch 26 implementation of the proposed method (C").
5 RELATED WORK
K-FAC uses the approximation Our methods are related to K-FAC by Martens and Grosse 12
(33) to approximate the blocks of the Hessian of the empirical risk of neural networks. most implementations of K-FAC, the off-diagona blocks of the Hessian are also set to zero. One of the main claimed benefits of K-FAC is its speed (compared to stochastic gradient descent) for large-batch size training. That said, recent empirical work has shown that this advantage of K-FAC disappears once the additional computational costs of hyperparameter tuning for large batch training is accounted for. There is line of work that extends the basic idea of K-FAC to convolutional layers (27]. Botev al. (18) further extend these ideas to present KFLR, a Kronecker factored low-rank approximation, and KFRA, Kronecker factored recursive approximation of the Gauss-Newton step. Singh and Alistarh (88) propose WoodFisher, a Woodbury matrix inverse-based estimate of the inverse Hessian, and apply it to neural network compression. Yao al. [29] propose AdaHessian, second-order optimizer that incorporates the curvature of the loss function via an adaptive estimation of the Hessian. Frantar al. (6) propose M-FAC, a matrix-free approximation of the natural gradient through queue of the (e.g., 000) recent gradients. These works fundamentally differ from our approach in that their objective is to approximate the Fisher or Gauss-Newton matrix inverse vector products. contrast, this work proposes to approximate the Gauss-Newtar matrix by only one of its Kronecker factors, which we find to achieve good performance at a substantial computational speedup and reduction of memory footprint. For overview of this area, we refer to Kunstner et al. (30) and Martens (31). For an overview of the technical aspects of backpropagation of second-order quantities, we refer to Dangel et al. 15 32
Taking a step back, K-FAC is one of many Newton-type methods for training neural networks. Other prominent examples of such methods include subsampled Newton methods (33). 34 (which approximate the Hessian by subsampling the terms in the empirical risk function and evaluating the Hessian of the subsampled terms) and sketched Newton methods (3) (which approximate the Hessian by sketching, e.g., by projecting the Hessian to a with random matrix). Another quasi-Newter method (35) proposes approximating the Hessian by a block-diagonal matrix using the structure of gradient and Hessian to further approximate these blocks. The main features that distinguish K-FAC from this group of methods are K-FAC's superior empirical performance and K-FAC's lack of theoretical justification.
space by multiplying
6 CONCLUSION
In this work, we presented ISAAC Newton, novel approximate curvature method based on layer- inputs. We demonstrated it to be a special case of the eggularizzaionngeeerzae Gauss-Newtae method and empirically demonstrate its utility. Specifically, our method features an asymptotically vanishing computational overhead in the mini-batch regime, while achieving competitive empirical performance on various benchmark problems.

---

ACKNOWLEDGMENTS
This work was supported by the IBM-MIT Watson AI Lab, the DFG in the Cluster of Excellence EXC 2117 "Centre for the Advanced Study of Collective Behaviour" (Project-ID 390829875), the Land Salzburg within the WISS 2025 project IDA-Lab (20102-F1901 166-KZP and 20204-W18S//25/1997 2019), and the National Science Foundation (NSF) (grants no. 1916271, 2027737, and 2113373).
REFERENCES
[1] N. Agarwal, B. Bullins, and Hazan, "Second-order stochastic optimization for machine learning in linear time," Journal on Machine Learning Research, vol. 18, no. 1, pp. 4148-4181
2017.
[2] J. Nocedal and S. Wright, Numerical Optimization, 2e, New York, NY, USA: Springer, 2006.
[3] A. Gonen and S. Shalev-Shwartz, "Faster SGD using sketched conditioning," arXiv preprint,
arXiv:1506 20600490..00644, 2015.
[4] M. Pilanci and M. J. Wainwright, "Newton sketch: A near linear-time optimization algorithm
with linear-quadratic convergence," SIAM Journal on Optimization, vol. 27, 2017.
[5] M. A. Erdogdu and A. Montanari, "Convergence rates of sub-sampled Newton methods," in
Proc. Neural Information Processing Systems (NeurIPS), 2015.
[6] E. Frantar, E. Kurtic, and D. Alistarh, "M-FAC: Efficient matrix-free approximations of second-order information." in Proc. Neural Information Processing Systems (NeurIPS), 2021.
[7] N. Doikov and Y. Nesterov, "Convex Optimization based on Global Lower Second-order Models," in Proc. Neural Information Processing Systems (NeurIPS), Curran Associates, Inc., 2020.
[8] Y. Nesterov and B. T. Polyak, "Cubic regularization of Newton method and its global perfor- mance," Mathematical Programming, vol. 108, 2006.
[9] S. Becker and Y. Lecun, "Improving the convergence of back pack-propagation learning with second-order methods," 1989.
[10] T. Schaul, S. Zhang, and Y. LeCun, "No more pesky learning rates," in International Conference on Machine Learning (ICML), 2013.
[11] Y. Ollivier, "Riemannian metrics for neural networks i: Feedforward networks," Information and Inference, vol. 4, pp. 108-153, Jun. 2015.
[12] J. Martens and R. Grosse, "Optimizing neural networks with Kronecker- approximate curvature," in International Conference on Machine Learning (ICML), 2015.
[13] A. N. Tikhonov and V. Y. Arsenin, Solutions of Ill-posed problems. W.H. Winston, 1977.
[14] P. Chen, "Hessian matrix vs.
Hessian matrix," " SIAM Journal on Numerical
Analysis, 2011.
[15] F. Dangel, Kunstner, and Hennig, "Backpack: Packing more into backprop," in Interna-
tional Conference Learning Representations, 2020.
[16] Y. LeCun, C. Cortes, and C. Burges, "MNIST Handwritten Digit Database," " A7T Labs, 2010.
[17] Xiao, K. Rasul, and R. Vollgraf, "Fashion-MNIST: novel image dataset for benchmarking machine learning algorithms," arXiv, 2017.
[18] A. Botev, H. Ritter, and D. Barber, "Practical
optimisation for deep learning," in International Conference on Machine Learning (ICML), 2017.
[19] D. Kingma and J. Ba. "Adam: A method for stochastic optimization." in International Confer- ence on Learning Representations (ICLR), 2015.
[20] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, "Bert: Pre-training of deep bidirectional transformers for language understanding." in North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), 2018.
[21] A. Warstadt, A. Singh, and S. R. Bowman, "Neural network acceptability judgments," Trans- actions of the Association for Computational Linguistics, vol. 7, 2019.
[22] W. B. Dolan and Brockett, " "Automatically constructing a corpus of sentential paraphrases." in Proceedings of the Third International Workshop on Paraphrasing (ZWP2005), 2005.
10

---

[23] D. Cer, M. Diab, E. Agirre, I. LLpez-Gazpin and L. Specia, "SemEval-200" task 1: Semantic textual similarity multilingual and crosslingual focused evaluation," " in Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2017), Vancouver, Canada: Association for Computational I Linguistics, 2017.
[24] K. He, X. Zhang, S. Ren, and J. Sun, "Deep residual learning for image recognition." in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
[25] Krizhevsky, V. Nair, and G. Hinton, "Cifar- (Canadian Institute for Advanced Research),"
2009.
[26] Paszke, S. Gross, Massa, et al., "Pytorch: An imperative style, bigh-perffrmance deep learning library." in Proc. Neural Information Processing Systems (NeurIPS), 2019.
[27] R. Grosse and J. Martens, "A Kronecker-factoree approximate Fisher matrix for convolution layers," in International Conference on Machine Learning (ICML), 2016.
[28] S. P. Singh and D. Alistarh, "Woodfisher; Efficient second-order approximation for neural network compression," " Proc. Neural Information Processing Systems (NeurlPS), 2020.
[29] Yao, A. Gholami, S. Shen, M. Mustafa, K. Keutzer, and M. Mahoney, "Adahessian: An adaptive second order optimizer for machine learning," in AAAI Conference Artificial Intelligence, 2021.
[30] F. Kunstner, L. Balles, and P. Hennig, "Limitations of the empirical Fisher approximation for natural gradient descent," in Proc. Neural Information Processing Systems (NeurIPS), 2019.
[31] J. Martens, "New insights and perspectives on the natural gradient method," Journal of Machine Learning Research, 2020.
[32] F. Dangel, S. Harmeling, and Hennig, "Modular block-diagonal curvature approximations for feedforward architectures," in International Conference on Artificial Intelligence and Statistics (AISTATS). 2020.
[33] F. vergent Algorithms," arXiv: 1601.04737, 2016.
and M. W. Mahoney, "Sub-Sampled Newton Methods Globally Con-
[34] P. Xu, J. Yang. F. Roosta, C. Ré, and M. W. Mahoney, "Sub-sampled Newton Methods with Non-uniform Sampling," in Proc. Neural Information Processing Systems (NeurIPS), 2016.
[35] D. Goldfarb, Y. Ren, and A. Bahamou, "Practical Quasi-Newton Methods for Training Deep Neural Networks," " Proceedings of the 34th International Conference on Neurul Information Processing Systems, 2020.

---

A PYTORCH IMPLEMENTATION
We display PyTorch [26] implementation of ISAAC for fully-connected layer below. Here, mark the important part (i.e., the part beyond the boilerplate) with a red rectangle.
import torch
class Isactiinerrunction (torch. autograd. Ftn.orrad.Funntiioooo. Function)
Ostaticmethor
def forward (ctx, input, weight, bias, la, inv_type):
ctx.
weight, bias)
ctx.la = la if inv_type
ahholesky_inversee
ctx. inveree = wwwphhchlllsyyinne inverss 'inverse':
elif inv inv_type
inveree = torch. inverse
else:
raise NotImplementedBror (inv_type)
return input @ weight (bias if bias is not None else 0)
estaticmethod
def backward (ctx, grad_output):
input, weight, bias ct.. saxee_teeeesssrs if www.nedds_inutt_rree
0 = puad_output @ weight
else:
grad_0 0 = None
if www.nneed_inuuuueoe
aaT input input. / www.tuu..ssaapp..o [0] I torch. (aaT. shape [0], device=aaT. device, dtype-aaT.dtypal
aaT IaaT inv aaT @ etx. www.ineeeeeeee (aaT / ctx.la Ib)
grad_1 = grad_output.7
I_b 1. ctx.la aaт jar_TaaT_ime
input
else:
grad_1 None
return
grad_0,
grad_i, grad_output. mean (0, keepdim-Trus if bias is not None else None, None, None, None,
class ISAACLinear (torch. nn .i
def __nit__(ee in_features, out_features, la, inv_sype='inveree **wwargs)
super (ISAACInnaar, s iii_init____ selfeatuure==i__eattreee out oit_faatureeroottffaaureeee **kwargs
self. la
self. inlttinv_type = inv_type
def forward input: torch. Tensor) -> Tensor::
return www.iieearuuuiiionnmmmmm input, self. fl weight,
12

---

self www.:baaa..uusn... (0) if self.bias is not None else None, self.la, self.inv_ty:e
B IMPLEMENTATION DETAILS
Unless noted differently, for all experiments, we tune the learning rate on a grid of (1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001). We verified this range to cover the full reasonable range of learning rates. Specifically, for every single experiment, we made sure that there is no learning rate outside this range which performs better.
For all language model experiments, we used the respective Huggingface Py PyTorch implementation.
All other hyperparameter details are given in the main paper.
C ADDITIONAL PROOFS
Proof of TTherrmm 7
We first show, that Ç as defined in 17 can expressed as in (20). Indeed by that
using (19), the Woodbury matrix identity and by regularizing the inverses,
< = \g g/b + Ag (a a//b Aa1)
g T
AgdA
(Ag..
+ + g/b)
g (Ma a aTa// aa/b)
1
Agda:
m
bdg 2
I
99
9
A
bigg
1
1
g
Ln
a'
Ib
+
aa
Aa
b\a
bba
I
g
Ib
99
9
g
bdg
bdg
I
1
T
Ib
1
bda
bd
I
g T
Ib
+
99
9
bāg
bdg
1
a
aa
T
I
+
aa
bAn
bdn
I
g" T
Ib
+
99
g
.g.
bdg
bdg
L
aa
T
I +
aa
b
bba
To show Assertion (i), we note that according to (7
lim 5
Ag +0 Ag aa
lim (g Tg/b + AxI) @ (a
a/bb
-
70
8 (a T -1 Tg
G g
13

---

where the first equality uses the definition of S in (7). The second equality is due the continuity of the matrix inversion and the last approximate equality follows from the K-FAC approximation (15)
To show Assertion (ii), consider l and lim +20 independently, that is
lim AR (AgIm g
(22)
Ag +00
lim tx
Im + bd
g g'g
Im:
and
lim da (datn + aaaaa aTa/b)
(23)
-1
lim
In +
aTa
a
I,-
bda
This then implies
lim
Ag (AgIm + +g +g"g/b) g
(24)
A 300
a An (aann + aa/b)
In g I, a
which concludes the proof.
D ADDITIONAL EXPERIMENTS
Gruden:
Grademi Cam
0.08
CANN LA -0-3
(A
0.06
CAA Murig
2.05
K-PAC
0.05
Mecig E.3
0.04
0.04
0.0%
560 fo 4ec-
Ue
0.03
560 joc. 6000
Adum
2000
4000
5000
8000
10000
10
training time
Epochs
Bradent REScArt
(4-
-0.32
Du -1
EA, 2611
Ag E31
Mu-Ay E01
5G0
Nomeriam
560 A' Memannur
0.03
560
0.03
0.62
0.02
0.01
3.00
0.00
2000
4000
5000
8000
10000
5
10
25
30
35
40
Eowiting time [s]
Epochs
Figure 7: Training loss of the MNIST auto-encoder trained with gradient descent. K-FAC, C. as well as SGD w/ momentum, SGD with 10x larger batch size (600), K-FAC with 10x larger batch size (600), and Adam. Comparing the performance per real-time (left) and per number of epochs (right). We display both the training loss (top) as well as the test loss (bottom) Runtimes are for a CPU core.
14

---

0.990
0.985
Test accuraty
Gredien: descent
5" for layers 3,4,,
lnyers
C" for layers
layers
0.965
for layers 3.4,3
for layers 2.4
0.960
250
500
750
1000
1250
1500
1750
2000
training
Figure 8: Test accuracy for training the MNIST classification task using C² only in selected layers. Runtimes are for CPU.
(4)
(43
[d]
4-
4-
4-
3-
16
Ih]
2.955
2.94
1078
18M
6
Figure 9: Reproduction of the experiments Figure but with the Fisher-based natural gradient formulation from Remark 4 For a description of the experimental settings, see the caption of Figure [ We observe that, for large Ag. the hchavior is similar to Figure 1, which expected as they are the same in the limit of Ag - co. Further, observe that (in this case of the Fisher-based not only in the limit of Xx of Av oc good performance better optimal performance compared to Ag oe, but X, - oo is more sensitive to changes Xg compared the sensitivity of the case of Ag -> ao wrt. changes in Aa. This phenomenon was also (to lesser extent) visible in the experiments of Figure I We would like to remark that the case of xg oo (i.e., (") is computationally more efficient compared Aa - oC.
but also in the limit be achieved. Moreover, in this specific experiment, A,
has slightly
15


Reviews:
[{"review_id": "qOKPn7pWWm", "reviewer": "Reviewer_U2Re", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a new approximation to the Gauss-Newton update step, in similar fashion to the well studied Kronecker-factored approximate curvature (K-FAC). The authors start from the generalised Gauss-Newton matrix $\\mathbf{G}_\\theta$ and follow a Monte-Carlo low-rank approximation to it before applying the K-FAC to approximate the expectation of the Kronecker with the Kronecker of the expectations.\nThen, they introduce independent Tikhonov regularisation to each of the terms in the Kronecker. By taking the limit of one of the regularisations to infinity, they prove that the \"expensive\" gradient term vanishes and the resulting gradient step is just a preconditioned gradient that accounts for the curvature of the space. The authors have experimentally evaluated the proposed update step in various set-ups, proving the effectiveness and the efficacy of it.\n", "summary_of_the_review": "In my opinion, the paper provides a good theoretical result and also a nice practical optimisation step which can be used as an easy approximation to second order optimisation.\n", "strength_and_weaknesses": "The paper explores a new pathway to approximate a Gauss-Newton step. I found this idea and the analysis simple and novel enough. The related work is sufficiently covered and, the presented approach is nicely motivated from the well studied K-FAC and the differences with which the authors intend to attack the problem.\n\nSee next section for some of my questions/issues I raised.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "F-PL7uOCysG", "reviewer": "Reviewer_xNi3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper seems only adding a $\\lambda  \\cdot I$ to the KFAC. \nThe theory in this paper seems only an easy extension of the work of KFAC.\nThe most useful part of this paper is the case $\\lambda_g \\to \\infty$. \nThis paper gives a descent direction without the expensive cost of computing $\\bar{g}$.\nThe experiments show the proposed algorithms are computation efficient.", "summary_of_the_review": "This paper seems only adding a $\\lambda  \\cdot I$ to the KFAC. \nThe most useful part of this paper is the case $\\lambda_g \\to \\infty$ which avoids the expensive cost of computing $\\bar{g}$.\nThe experiments show the proposed algorithms are computation efficient.\n\nI am not very familiar with training algorithms of  neural networks.\nThus, I am not sure the value of the proposed algorithms.\n", "strength_and_weaknesses": "This paper proposes novel second-order algorithms for training neural networks. \nThe experiments show the proposed algorithms are computation efficient.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "5mqT7mHz6Ow", "reviewer": "Reviewer_2oTp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a method that conditions the gradient using selected second-order information, which uses less computation than naively running Newton. ", "summary_of_the_review": "The paper seems like a minor modification of K-FAC by adding some stochasticity and ridge penalties. ", "strength_and_weaknesses": "* Strength\n\nThe paper shows that their proposed method works well in practice.\n\n* Weakness\n- There are rather few experiments\n- The main change in the paper is adding ridge penalties, which seems rather incremental", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster