ISAAC Newton: Input-based Approximate Curvature for Newton’s Method
Anonymous Author(s)
Abstract
We present ISAAC (Input-baSed ApproximAte Curvature), a novel method thatconditions the gradient using selected second-order information and has an asymptotically vanishing computational overhead, assuming a batch size smaller thanthe number of neurons.
We show that it is possible to compute a good conditionerbased on only the input to a respective layer without a substantial computationaloverhead.
The proposed method allows effective training even in small-batchstochastic regimes, which makes it competitive to first-order as well as quasiNewton methods.
Introduction
While second-order optimization methods are traditionally much less explored than first-ordermethods in large-scale machine learning (ML) applications due to their memory requirements andprohibitive computational cost per iteration, they have recently become more popular in ML mainlydue to their fast convergence properties when compared to first-order methods [1].
The expensivecomputation of an inverse Hessian (also known as pre-conditioning matrix) in the Newton step hasalso been tackled via estimating the curvature from the change in gradients.
Loosely speaking, thesealgorithms are known as quasi-Newton methods and a comprehensive treatment can be found inthe textbook [2].
In addition, various new approximations to the pre-conditioning matrix have beenproposed in the recent literature [3]–[6].
From a theoretical perspective, second-order optimizationmethods are not nearly as well understood as first-order methods.
It is an active research direction tofill this gap [7], [8].
Motivated by the task of training neural networks, and the observation that invoking local curvatureinformation associated with neural network objective functions can achieve much faster progressper iteration than standard first-order methods [9]–[11], several methods have been proposed.
Oneof these methods, that received significant attention, is known as Kronecker-factored Approximate
Curvature (K-FAC) [12], whose main ingredient is a sophisticated approximation to the generalized
Gauss-Newton matrix and the Fisher information matrix quantifying the curvature of the underlyingneural network objective function, which then can be inverted efficiently.
Inspired by the K-FAC approximation and the Tikhonov regularization of the Newton method, weintroduce a novel two parameter regularized Kronecker-factorized Newton update step.
The proposedscheme disentangles the classical Tikhonov regularization and allows us to condition the gradientusing selected second-order information and has an asymptotically vanishing computational overhead.
While this property makes the presented method highly attractive from the computational complexityperspective, we show that its achieved empirical performance on complicated high-dimensional
Machine Learning problems remains comparable to existing state-of-the-art methods.
The contributions of this paper can be summarized as follows: (i) we propose a novel two parameterregularized K-FAC approximated Gauss-Newton update step; (ii) we show that asymptotically—as both regularization parameters vanish—our method recovers the classical K-FAC scheme and inthe opposite setting—as both regularization parameters grow—our method asymptotically reducesto classical gradient descent; (iii) we prove that for an arbitrary pair of regularization parameters,the proposed update direction is always a direction of decreasing loss; (iv) in the limit, as oneregularization parameter grows, we obtain an efficient and effective conditioning of the gradient withan asymptotically vanishing overhead; (v) we empirically analyze the presented method and find thatour efficient conditioning method maintains the performance of its more expensive counterpart; (vi)we demonstrate the effectiveness of the presented method in the setting of small-batch stochasticregimes and observe that it is competitive to first-order as well as quasi-Newton methods.
Preliminaries
In this section, we review aspects of second-order optimization, with a focus on generalized GaussNewton methods.
In combination with Kronecker factorization, this leads us to a new regularizedupdate scheme.
We consider the training of an L -layer neural network f ( x ; θ ) defined recursively as
[Equation]
where a 0 = x is the vector of inputs and a L = f ( x ; θ ) is the vector of outputs.
Unless noted otherwise,we assume these vectors to be row vectors (i.e., in R 1 × n ) as this allows for a direct extension to the(batch) vectorized case (i.e., in R b × n ) introduced later.
For any layer i , let W ( i ) ∈
R
d i − 1 ×
d i be aweight matrix and let ϕ be an element-wise nonlinear function.
We consider a convex loss function
L ( y, y ′ ) that measures the discrepancy between y and y ′ .
The training optimization problem is then
[Equation]
where θ = (cid:2) θ (1) , . .
, θ ( L ) (cid:3) with θ ( i ) =
vec( W ( i ) ) .
The classical Newton method for solving (2) is expressed as the update rule
[Equation]
where η > 0 denotes the learning rate and H θ is the Hessian corresponding to the objective functionin (2).
The stability and efficiency of an estimation problem solved via the Newton method can beimproved by adding a Tikhonov regularization term [13] leading to a regularized Newton method
[Equation]
where λ > 0 is the so-called Tikhonov regularization parameter.
It is well-known [14], [15], thatunder the assumption of approximating the model f with its first-order Taylor expansion, the Hessiancorresponds with the so-called generalized Gauss-Newton (GGN) matrix G θ , and hence (4) can beexpressed as
[Equation]
A major practical limitation of (5) is the computation of the inverse term.
A method that alleviates thisdifficulty is known as Kronecker-Factored Approximate Curvature (K-FAC) [12] which approximatesthe block-diagonal (i.e., layer-wise) empirical Hessian or GGN matrix.
Inspired by K-FAC, therehave been other works discussing approximations of G θ and its inverse [15].
In the following, wediscuss a popular approach that allows for (moderately) efficient computation.
The generalized Gauss-Newton matrix G θ is defined as
[Equation]
where J and H denote the Jacobian and Hessian matrices, respectively.
Correspondingly, the diagonalblock of G θ corresponding to the weights of the i th layer W ( i ) is
[Equation]
According to the backpropagation rule J θ ( i ) f ( x ; θ ) =
J z i f ( x ; θ ) a i − 1 , a ⊤ b = a ⊗ b , and themixed-product property, we can rewrite G W ( i ) as
G W
[Equation]
i
[Equation]
[Equation] [Equation]
7) =
E
[Equation]
g
[Equation]
a
[Equation]
1 ) ⊤
[Equation]
g
[Equation]
a
[Equation] [Equation]
=
E
[Equation]
(¯ g ⊤ ¯ g ) ⊗ ( a ⊤ i − 1 ⊗ a i − 1 ) (cid:3) , (8) where
[Equation]
Remark 1 (Monte-Carlo Low-Rank Approximation for ¯ g ⊤ ¯ g ) .
As ¯ g is a matrix of shape m × d iwhere m is the dimension of the output of f , ¯ g is generally expensive to compute.
Therefore, [12] usea low-rank Monte-Carlo approximation to estimate H f L ( f ( x ; θ ) , y ) and thereby ¯ g ⊤ ¯ g .
For this, weneed to use the distribution underlying the probabilistic model of our loss L (e.g., Gaussian for MSEloss, or a categorical distribution for cross entropy).
Specifically, by sampling from this distributionp f ( x ) defined by the network output f ( x ; θ ) , we can get an estimator of H f L ( f ( x ; θ ) , y ) via theidentity
[Equation]
An extensive reference for this (as well as alternatives) can be found in Appendix A.2 of Dangel etal.
The respective rank-1 approximation (denoted by ≜ ) of H f L ( f ( x ; θ )) is
[Equation]
where ˆ y ∼ p f ( x ) .
Respectively, we can estimate ¯ g ⊤ ¯ g using this rank-1 approximation with
[Equation]
In analogy to ¯ g , we introduce the gradient of training objective with respect to pre-activations z i as
[Equation]
In other words, for a given layer, let g ∈ R 1 ×
d i denote the gradient of the loss between an output andthe ground truth and let ¯ g ∈ R m ×
d i denote the derivative of the network f times the square root ofthe Hessian of the loss function (which may be approximated according to Remark 1), each of themwith respect to the output
z i of the given layer i .
Note that ¯ g is not equal to g and that they require onebackpropagation pass each (or potentially many for the case of ¯ g ).
This makes computing ¯ g costly.
Applying the K-FAC [12] approximation to (8) the expectation of Kronecker products can beapproximated as the Kronecker product of expectations as
[Equation]
where, for clarity, we drop the index of a i − 1 in (8) and denote it with a ; similarly we denote G W ( i )as G .
While the expectation of Kronecker products is generally not equal to the Kronecker productof expectations, this K-FAC approximation (13) has been shown to be fairly accurate in practiceand 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 theidentity ( A ⊗ B ) − 1 =
A − 1 ⊗ B − 1 which significantly simplifies the computation of an inverse.
In practice, E (¯ g ⊤ ¯ g ) and E (a ⊤ a) can be computed by averaging over a batch of size b as
[Equation]
[Equation]
where we denote batches of g , ¯ g and a , as g ∈ R b × d i , ¯ g ¯ g ¯ g ∈ R rb × d i and a ∈ R b × d i − 1 , where our layerhas d i − 1 inputs, d i outputs, b is the batch size, and r is either the number of outputs m or the rank ofan approximation according to Remark 1.
Correspondingly, the K-FAC approximation of the GGNmatrix and its inverse are concisely expressed as
[Equation]
Equipped with the standard terminology and setting, we now introduce the novel, regularized updatestep.
First, inspired by the K-FAC approximation (13), the Tikhonov regularized Gauss-Newtonmethod (5) can be approximated by
[Equation]
with regularization parameter λ > 0 .
A key observation, which is motivated by the structure ofthe above update, is to disentangle the two occurrences of λ into two independent regularizationparameters λ g , λ a > 0 .
By defining the Kronecker-factorized Gauss-Newton update step as
[Equation]
we obtain the concise update equation
[Equation]
This update (18) is equivalent to update (16) when in the case of η ∗ =
η λ g λ a and λ = λ g
= λ a .
Thisequivalence does not restrict η ∗ , λ g , λ a in any way, and changing λ g or λ a does not mean that wechange our learning rate or step size η ∗ .
Parameterizing ζζζ in (17) with the multiplicative terms λ g λ amakes the formulation more convenient for analysis.
In this paper, we investigate the theoretical and empirical properties of the iterative update rule (18)and in particular show how the regularization parameters λ g , λ a affect the Kronecker-factorized
Gauss-Newton update step ζζζ .
When analyzing the Kronecker-factorized Gauss-Newton update stepζζζ , a particularly useful tool is the vector product identity,
[Equation]
[Equation]
where the gradient with respect to the weight matrix is g ⊤ a .
Theoretical Guarantees
In this section, we investigate the theoretical properties of the Kronecker-factorized Gauss-Newtonupdate direction ζζζ as defined in (17).
We recall that ζζζ introduces a Tikonov regularization, as it iscommonly done in implementations of second order-based methods.
Not surprisingly, we show thatby decreasing the regularization parameters λ g , λ a the update rule (18) collapses (in the limit) to theclassical Gauss-Newton method, and hence in the regime of small λ g , λ a the variable ζζζ describes the
Gauss-Newton direction.
Moreover, by increasing the regularization strength, we converge (in thelimit) to the conventional gradient descent update step.
The key observation is that, as we disentangle the regularization of the two Kronecker factors ¯ g ¯ g ¯ g ⊤ ¯ g ¯ g ¯ gand a ⊤ a , and consider the setting where only one regularizer is large ( λ g → ∞ to be precise),we obtain an update direction that can be computed highly efficiently.
We show that this settingdescribes an approximated Gauss-Newton update scheme, whose superior numerical performance isthen empirically demonstrated in Section 4.
Theorem 1 (Properties of ζζζ ) .
The K-FAC based update step ζζζ as defined in (17) can be expressed as
[Equation]
Moreover, ζζζ admits the following asymptotic properties:
(i)
In the limit of λ g , λ a → 0 ,
[Equation]
ζζζ is the K-FAC approximation of the Gauss-Newton step, i
e.,
[Equation]
) L ( f ( x ; θ )) , where ≈ denotes the K-FAC approximation (15) .
(ii)
In the limit of λ g , λ a → ∞ , ζζζ is the gradient, i
e., lim λ g ,λ a →∞ ζζζ = ∇ θ ( i ) L ( f ( x ; θ )) .
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 λ g and λ anumerically close to zero because we want to explore the full spectrum of settings for λ g and λ a .
Thus, we prove that ζζζ is a direction of increasing loss, independent of the choices of λ g and λ a .
Theorem 2 (Correctness of ζζζ is independent of λ g and λ a ) .
ζζζ is a direction of increasing loss,independent of the choices of λ g and λ a .
Proof.
Recall that ( λ g I m +¯ g ¯ g ¯ g ⊤ ¯ g ¯ g ¯ g/b ) and ( λ a I n + a ⊤ a /b ) are positive semi-definite (PSD) matrices bydefinition.
Their inverses ( λ g I m + ¯ g ¯ g ¯ g ⊤ ¯ g ¯ g ¯ g/b ) − 1 and ( λ a I n + a ⊤ a /b ) − 1 are therefore also PSD.
As the
Kronecker product of PSD matrices is PSD, the conditioning matrix ( ( λ g I m + ¯
g ¯ g ¯ g ⊤ ¯ g ¯ g ¯ g/b ) − 1 ⊗ ( λ a I n +a ⊤ a /b ) −
≈ G − 1 ) is PSD, and therefore the direction of the update step remains correct.
From our formulation of ζζζ , we can find that, in the limit for λ g → ∞ , Equation (21) does not dependon ¯ g ¯ g ¯ g .
This is computationally very beneficial as computing ¯ g ¯ g ¯ g is costly as it requires one or evenmany additional backpropagation passes.
In addition, it allows conditioning the gradient update bymultiplying a b × b matrix between g ⊤ and a , which can be done very fast.
Theorem 3 (Efficient Update Direction) .
In the limit of λ g → ∞ , the update step ζζζ converges tolim λ g →∞ ζζζ = ζζζ ∗ , where
[Equation]
(i)
Here, the update direction ζζζ ∗ is based only on the inputs and does not require computing ¯ g ¯ g ¯ g
(which would require a second backpropagation pass), making it efficient.
(ii)
The computational cost of computing the update ζζζ ∗ lies in O ( bn 2 + b 2 n + b 3 ) , where n is the
number of neurons in each layer.
This comprises the conventional cost of computing the gradient
∇ =
g ⊤ x lying in O ( bn
and the overhead of computing ζζζ ∗ instead of ∇ lying in O ( b 2 n + b 3 ) .
The overhead is vanishing, assuming n ≫ b .
For b > n the complexity lies in O ( bn 2 + n 3 ) .
Proof.
We first show the property (21).
Note that
according to (22), λ g · (cid:0) λ g I m + ¯ g ¯ g ¯ g ⊤ ¯
g ¯ g ¯ g/b (cid:1) − 1 converges in the limit of λ g → ∞ to I m , and therefore (21) holds.(i)
The statement follows from the fact that the term ¯ g ¯ g ¯ g does not appear in the equivalent characterization (21) of ζζζ ∗ .(ii
We first note that the matrix aa ⊤ is of dimension b × b , and can be computed in O ( b 2 n ) time.
Next, the matrix
[Equation]
is of shape b × b and can be multiplied with a in O ( b 2 n ) time.
Notably, (21) can be computed with a vanishing computational overhead and with only minormodifications 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 toapplying it only to individual layers.
As we see in the experimental section, in many cases in the mini-batch regime (i.e., b < n ), theoptimal (or a good) choice for λ g actually lies in the limit to ∞ .
This is a surprising result, leading tothe efficient and effective ζζζ ∗ = ζζζ λ g →∞ optimizer.
Remark 2 (Relation between Update Direction ζζζ and ζζζ ∗ ) .
When comparing the update directionζζζ in (20) without regularization (i.e., λ g → 0 , λ a → 0 ) with ζζζ ∗ (i.e., λ g → ∞ ) as given in (21) , itcan be directly seen that ζζζ ∗ corresponds to a particular pre-conditioning of ζζζ , since ζζζ ∗ =
Mζζζ for
M = 1 bλ g ¯ g ¯ g ¯ g ⊤ ¯ g ¯ g ¯ g.
As the last theoretical property of our proposed update direction ζζζ ∗ , we show that in specific networksζζζ ∗ coincides with the Gauss-Newton update direction.
Theorem 4 ( ζζζ ∗ is Exact for the Last Layer) .
For the case of linear regression or, more generally, thelast 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 ¯ g ⊤ ¯ g ¯ g ¯ g is I .
Thus, ζζζ ∗ = ζζζ .
Remark 3.
The direction ζζζ ∗ corresponds to the Gauss-Newton update direction with an approximation of G that can be expressed as G ≈ E (cid:2) I ⊗ (a ⊤ a) (cid: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 isthat 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 naturalgradient update step F − 1 ∇ .
The i -th diagonal block of F ( F θ ( i ) =
E (cid:2) (g ⊤ i g i ) ⊗ ( a ⊤ i − 1 ⊗ a i − 1 ) (cid:3) ),has the same form as a block of the GGN matrix G ( G θ ( i )
=
E (cid:2) (¯ g ⊤ i ¯ g i ) ⊗ ( a ⊤ i − 1 ⊗ a i − 1 ) (cid:3) ).
Thus, we can replace ¯ g ¯ g ¯ g with g in our theoretical results to obtain their counterparts for F .
+1+2+3+4+5+6 log 10 g
-6 -5 - -3 -20 +
+ +3 +
+
+ l o g 10 a (a)
+1+2+3+4+5+6 log 10 g (b)
+1+2+3+4+5+6 log 10 g (c) +1+2+3+4+5+6 log 10 g (d)
[Figure] [Figure]
+
+ +3 +
+
+ l o g 10 a (e)
[Figure] [Figure]
[Figure] [Figure]
[Figure]
The optimizer is SGD except for (c, g) where the optimizer is SGD
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 in the mini-batch regime.
Inthis section, we provide empirical evidence that ζζζ ∗ is a good update direction, even in deep learning.
Specifically, we demonstrate that
(E1) ζζζ ∗ achieves similar performance to K-FAC, while being substantially cheaper to compute.
(E2)
The performance of our proposed method can be empirically maintained in the mini-batchregime ( n ≫ b ).(E3)
ζζζ ∗ may be used for individual layers, while for other layers only the gradient ∇ is used.
Thisstill leads to improved performance.(E4) ζζζ ∗ also improves the performance for training larger models such as BERT and ResNet.(E5)
The runtime and memory requirements of ζζζ ∗ are comparable to those of gradient descent.
E1: Impact of Regularization Parameters
For (E1), we study the dependence of the model’s performance on the regularization parameters λ gand λ a .
Here, we train a 5-layer deep neural network on the MNIST classification task [16] with abatch size of 60 for a total of 40 epochs or 40 000 steps.
The plots in Figure 1 demonstrate that the advantage of training by conditioning with curvatureinformation can be achieved by considering both layer inputs a and gradients with respect to randomsamples ¯ g ¯ g ¯ g , but also using only layer inputs a .
In the plot, we show the performance of
Newton’s method or the general K-FAC approximation corresponds to the area with small λ g and λ a .
The interesting finding here is that the performance does not suffer by increasing λ g toward ∞ , i.e.,from left to right in the plot.
[Figure]
In addition
[Figure]
in Figure
[Figure]
we consider the case of regression with an auto-encoder trained with the
MSE loss on MNIST [16] and Fashion-MNIST [17].
Here
[Figure]
we follow the same principle as aboveand also find that ζζζ ∗ performs well
[Figure] [Figure]
In Figure 7, we compare the loss for different methods.
Here, we distinguishbetween loss per time (left) and lossper number of steps (right).
We can observe that, for λ = 0 .
1 , K-FAC, ζζζ , andζζζ ∗ are almost identical per update step(right), while ζζζ ∗ is by a large marginthe fastest, followed by ζζζ , and the conventional K-FAC implementation is theslowest (left).
On the other hand, forλ = 0 .
01 we can achieve a faster convergence than with λ = 0 .
1 , but hereonly the K-FAC and ζζζ methods are numerically stable, while ζζζ ∗ is unstable inthis case.
This means in the regime ofvery small λ , ζζζ ∗ is not as robust as KFAC and ζζζ , however, it achieves goodperformance with small but moderateλ like λ = 0 .
For λ < 0 .
01 , also
K-FAC and ζζζ become numerically unstable in this setting and, in general, weobserved that the smallest valid λ for
K-FAC is 0 .
01 or 0 .
001 depending onmodel and task.
Under considerationof the runtime, ζζζ ∗ performs best as it isalmost as fast as gradient descent whileperforming equivalent to K-FAC and ζζζ .
Specifically, a gradient descent step isonly about 10% faster than ζζζ ∗ .
E2:
Minibatch Regime
[Figure]
For (E2), in Figure 1, we can see that trainingperforms well for n ∈ { 100 , 400 , 1 600 } neurons per layer at a batch size of only 60 .
Also, inall other experiments, we use small batch sizesof between 8 and 100 .
E3: ζζζ ∗ in Individual Layers
In Figure 5, we train the 5-layer fully connectedmodel with 400 neurons per layer.
Here, weconsider the setting that we use ζζζ ∗ in some ofthe layers while using the default gradient ∇in other layers.
Specifically, we consider the
[Table]
settings, where all, the first, the final, the first three, the final three, the odd numbered, and theeven numbered layers are updated by ζζζ ∗ .
We observe that all settings with ζζζ ∗ perform better thanplain gradient descent, except for “ ζζζ ∗ for layers 3,4,5” which performs approximately equivalent togradient descent.
E4: Large-scale Models
BERT To demonstrate the utility of ζζζ ∗ also in large-scale models, we evaluate it for fine-tuning
BERT [20] on three natural language tasks.
In Table 1, we summarize the results for the BERTfine-tuning task.
For the “Corpus of Linguistic Acceptability” (CoLA) [21] data set, we fine-tuneboth the BERT-Base and the BERT-Mini models and find that we outperform the gradient descentbaseline in both cases.
For the “Microsoft Research Paraphrase Corpus” (MRPC) [22] data set, wefine-tune the BERT-Base model and find that we outperform the baseline both in terms of accuracyand F1-score.
Finally, on the “Semantic Textual Similarity Benchmark” (STS-B) [23] data set, wefine-tune the BERT-Mini model and achieve higher Pearson and Spearman correlations than thebaseline.
While for training with CoLA and MRPC, we were able to 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 not work well.
Therefore, for STS-B, we evaluated it using the SGD withmomentum optimizer.
For each method, we performed a grid search over the hyperparameters.
Wenote that we use a batch size of 8 in all BERT experiments.
[Figure]
ResNet In addition, we conduct an experimentwhere we train the last layer of a ResNet withζζζ ∗ , while the remainder of the model is updated using the gradient ∇ .
Here, we train a
ResNet-18 [24] on CIFAR-10 [25] using SGDwith a batch size of 100 in a vanilla setting, i.e.,without additional tricks employed in by He etal.
[24] and others.
Specifically, we use (i) aconstant learning rate for each training (optimalfrom (1 , 0 . 3 , 0 . 1 , 0 .
03 , 0 . 01) ) and (ii) vanilla
SGD and not momentum-based SGD.
The reason behind this is that we want a vanilla experiment and with aspects such as extensively tuningmultiple parameters of learning rate schedulerwould make the evaluation less transparent; however, therefore, all accuracies are naturally lower than SOTA.
In Figure 6, we plot the test accuracyagainst time.
The results show that the proposed method outperforms vanilla SGD when appliedto the last layer of a ResNet-18.
To validate that the learning rate is not the cause for the betterperformance, we also plot the neighboring learning rates and find that even with a too small or toolarge learning rate ζζζ ∗ outperforms gradient descent with the optimal learning rate.
E5: Runtime and Memory
Finally, we also evaluate the runtime and memory requirements of each method.
The runtimeevaluation is displayed in Table 2.
We report both CPU and GPU runtime using PyTorch [26] and(for K-FAC)
the backpack library [15].
Note that the CPU runtime is more representative of thepure computational cost, as for the first rows of the GPU runtime the overhead of calling the GPUis dominant.
When comparing runtimes between the gradient and ζζζ ∗ on the GPU, we can observethat we have an overhead of around 2 .
5 s independent of the model size.
The overhead for CPU timeis 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 ζζζ ∗ is around 4 times the runtime of the gradient, and K-FAC has an evensubstantially larger runtime.
Regarding memory, ζζζ ∗ (contrasting the other approaches) also requiresonly a small additional footprint.
Remark 5 (Implementation) .
The implementation of ζζζ ∗ can be done by replacing the backpropagationstep of a respective layer by (21) .
As all “ingredients” are already available in popular deep learningframeworks, it requires only little modification (contrasting K-FAC and ζζζ , which require at least oneadditional backpropagation.)
[Table]
We will publish the source code of our implementation.
In the appendix, we give a PyTorch [26]implementation of the proposed method ( ζ ζ ζ ∗ ).
Related Work
Our methods are related to K-FAC by Martens and Grosse [12].
K-FAC uses the approximation(13) to approximate the blocks of the Hessian of the empirical risk of neural networks.
In mostimplementations of K-FAC, the off-diagonal blocks of the Hessian are also set to zero.
One of themain claimed benefits of K-FAC is its speed (compared to stochastic gradient descent) for large-batchsize training.
That said, recent empirical work has shown that this advantage of K-FAC disappearsonce the additional computational costs of hyperparameter tuning for large batch training is accountedfor.
There is a line of work that extends the basic idea of K-FAC to convolutional layers [27].
Botev etal.
[18] further extend these ideas to present KFLR, a Kronecker factored low-rank approximation,and KFRA, a Kronecker factored recursive approximation of the Gauss-Newton step.
Singh and
Alistarh [28] propose WoodFisher, a Woodbury matrix inverse-based estimate of the inverse Hessian,and apply it to neural network compression.
Yao et al.
[29] propose AdaHessian, a second-orderoptimizer that incorporates the curvature of the loss function via an adaptive estimation of the Hessian.
Frantar et al.
[6] propose M-FAC, a matrix-free approximation of the natural gradient through a queueof the (e.g., 1 000 ) recent gradients.
These works fundamentally differ from our approach in that theirobjective is to approximate the Fisher or Gauss-Newton matrix inverse vector products.
In contrast,this work proposes to approximate the Gauss-Newton matrix by only one of its Kronecker factors,which we find to achieve good performance at a substantial computational speedup and reduction ofmemory footprint.
For an 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.
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] (whichapproximate the Hessian by subsampling the terms in the empirical risk function and evaluating the
Hessian of the subsampled terms) and sketched Newton methods [3]–[5] (which approximate the
Hessian by sketching, e.g., by projecting the Hessian to a lower-dimensional space by multiplying itwith a random matrix).
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.
Conclusion
In this work, we presented ISAAC Newton, a novel approximate curvature method based on layerinputs.
We demonstrated it to be a special case of the regularization-generalized Gauss-Newtonmethod and empirically demonstrate its utility.
Specifically, our method features an asymptoticallyvanishing computational overhead in the mini-batch regime, while achieving competitive empiricalperformance on various benchmark problems.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms to them?[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs of all theoretical results?
[Yes]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimentalresults (either in the supplemental material or as a URL)?
[Yes] /
[No] We include a
Python / PyTorch implementation of the method in the supplementary material.
We willpublicly release full source code for the experiments.
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how they werechosen)?
[Yes](c)
Did you report error bars (e
g., with respect to the random seed after running experimentsmultiple times)?
[Yes](d) Did you include the total amount of compute and the type of resources used (e
g., type ofGPUs, internal cluster, or cloud provider)?
[Yes]4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a)
If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
(IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spenton participant compensation?
[N/A]
A
PyTorch Implementation
We display a PyTorch [26] implementation of ISAAC for a fully-connected layer below.
Here, wemark the important part (i.e., the part beyond the boilerplate) with a red rectangle.import torch class ISAACLinearFunction(torch.autograd.
Function): @staticmethoddefforward(ctx, input, weight, bias, la, inv_type): ctx.save_for_backward(input, weight, bias) ctx.la = la if inv_type =
= 'cholesky_inverse':
ctx.inverse = torch.cholesky_inverse elif inv_type
=
= 'inverse': ctx.inverse = torch.inverse else: raise NotImplementedError(inv_type) return input @ weight.
T + (bias if bias is not None else 0) @staticmethoddefbackward(ctx, grad_output): input, weight, bias = ctx.saved_tensors if ctx.needs_input_grad[0]: grad_0=grad_output@weight else: grad_0 =
None if ctx.needs_input_grad[1]: aaT
= input @ input.
T / grad_output.shape[0]
I_b =
torch.eye(aaT.shape[0], device=aaT.device, dtype=aaT.dtype) aaT_IaaT_inv =
aaT @
ctx.inverse(aaT / ctx.la + I_b) grad_
= grad_output.
T @ (
I_b - 1. /
ctx.la * aaT_IaaT_inv ) @ input else: grad_
=
None return ( grad_0, grad_1, grad_output.mean(0, keepdim=
True) if bias is not None else None,
None, None, None, ) class ISAACLinear(torch.nn.
Linear): def __init__(self, in_features, out_features, la, inv_type='inverse', **kwargs): super(ISAACLinear, self).__init__( in_features=in_features, out_features=out_features, **kwargs ) self.la
= la self inv type inv_type def
: return self.bias.unsqueeze(0) if self.bias is not None else None, self.la, self.inv_type )
B Implementation Details
Unless noted differently, for all experiments, we tune the learning rate on a grid of(1 , 0 . 3 , 0 . 1 , 0 .
We verified this range to cover the full reasonable range oflearning rates.
Specifically, for every single experiment, we made sure that there is no learning rateoutside this range which performs better.
For all language model experiments, we used the respective Huggingface PyTorch implementation.
All other hyperparameter details are given in the main paper.
The code will be made publicly available.
C Additional Proofs
Proof of Theorem 1.
We first show, that ζζζ as defined in (17) can be expressed as in (20).
Indeed byusing (19), the Woodbury matrix identity and by regularizing the inverses, we can see that
[Equation]
To show Assertion (i)
[Equation]
we note that according to (17)
[Equation]
where the first equality uses the definition of ζζζ in (17).
The second equality is due to the continuity ofthe matrix inversion and the last approximate equality follows from the K-FAC approximation (15).
To show Assertion (ii), we consider lim λ g →∞ and lim λ a →∞ independently, that is
[Equation]
and
[Equation]
This then implies
[Equation]
which concludes the proof.
D Additional Experiments
[Figure] [Figure] [Figure]