A Layer-Wise Natural Gradient Optimizer for Training Deep Neural Networks
Authors: Xiaolei Liu
Shaoshuai Li
Kaixin Gao
Binfeng Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: Natural gradient method
second-order optimization
deep neural networks
Kronecker factorization
Primary area: optimization_for_deep_networks

Abstract: Second-order optimization algorithms, such as the Newton method and the natural gradient descent (NGD) method exhibit excellent convergence properties for training deep neural networks, but the high computational cost limits its practical application. In this paper, we focus on the NGD method and propose a novel layer-wise natural gradient descent (LNGD) method to further reduce computational costs and accelerate the training process. Specifically, based on the block diagonal approximation of the Fisher information matrix, we first propose the layer-wise sample method to compute each block matrix without performing a complete back-propagation. Then, each block matrix is approximated as a Kronecker product of two smaller matrices, one of which is a diagonal matrix, while keeping the traces equal before and after approximation. By these two steps, we provide a new approximation for the Fisher information matrix, which can effectively reduce the computational cost while preserving the main information of each block matrix. Moreover, we propose a new adaptive layer-wise learning rate to further accelerate training. Based on these new approaches, we propose the LNGD optimizer. The global convergence analysis of LNGD is established under some assumptions. Experiments on image classification and machine translation tasks show that our method is quite competitive compared to the state-of-the-art methods.

A Layer- Natural Gradient Optimizer for Training Deep Neural Networks
Xiaolei Liu" Ant Group Hangzhou, China
Shaoshuai Li" Ant Group Hangzhou, China lishaoshuai. 1ss@mybank... cn
liuxiaolei. 1elllx1mmmaaa @eymankmankc.o
Kaixin Gao
Binfeng Wang Ant Group Hangzhou, China vangbinfeng wbf @mybank... cn
Ocean University China Qingdao, China gaokaixin060163.cm
Abstract
Second-order optimization algorithms, such the Newton method and the natu- ral gradient descent (NGD) method exhibit excellent convergence properties for training deep neural networks, but the high computational cost limits its practical application. In this paper, we focus on the NGD method and propose a novel layer- wise natural gradient descent (LNGD) method to further reduce computational costs and accelerate the training process. Specifically, based on the block diagonal approximation of the Fisher information matrix, we first propose the layer-wise sample method to compute each block matrix without performing complete back- propagation. Then, each block matrix approximated as Kronecker product of two smaller matrices, one of which is diagonal matrix, while keeping the traces equal before and after approximation. By these two steps, we provide a new approximation for the Fisher information matrix, which can effectively reduce the computational cost while preserving the main information of each block matrix. Moreover, propose new adaptive layer-wise learning rate to further accelerate training. Based on these new approaches, we propose the LNGD optimizer. The global convergence analysis of LNGD is established under some assumptions. Experiments on image classification and machine translation tasks show that our method is quite competitive compared to the state-of-the ffft--hatt art methods.
Introduction
With the rapid increase in the size of deep neural networks (DNNs) models in both areas of computer vision (CV) and natural language processing (NLP), there have been remarkable attentions given to optimizing algorithms. An effective optimizer can significantly improve the training speed of models while ensuring high prediction performance. First-order gradient descent methods are workhorses of training DNNs, which can be broadly divided into two categories: methods use same learning rate, such as stochastic gradient descent (SGD) 1 and accelerations [2 3], and methods use adaptive learning rate, such as AdaDelta 4. RMSProp [. ADAM 6 and Adabelief [7]] Although first-order gradient descent methods enjoy low computational cost and ease of implementation, they might suffer from sensitivity to hyperparameters and slow convergence. It is challenging to reduce the number of iterations and computational time these methods.
Joint first author, these authors contributed equally to this work. Conresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Some work has considered introducing curvature information when updating parameters of DNNs to improve the convergence speed and overcome the above shortcomings of the first-order methods. However, second-order optimization methods need to store and compute the inverse of curvature matrix, which brings expensive storage and computation costs and limits the application of second- order methods in training large-scale DNNs. Therefore, many approximate second-order methods have been proposed for training large-scale models. For example, Keskar and Berahas [8] proposed a
stochastic quasi-Newton algorithm for training recurrent neural networks. Yao et al. p approximated
the Hessian matrix as a diagonal operator, which is achieved by applying Hutchinson's 's method, and proposed the AdaHessian method. Goldfarb, Ren and Bahamou [10] developed Kronecker- factored block-diagonal BFGS and its limited memoor variants L-BFGS methods for training DNNs. Generalized Gauss-Newton Newton methods, such as the Hessian fesse method [I]] and the Krylov subspace method 12 also have been proposed to approximate the Hessian matrix.
The natural gradient descent (NGD) method [13. which preconditions the gradient by the Fisher information matrix instead of the Hessian matrix, also has shown effectiveness in training DNNs 14 15 16 17 NGD explores the steepest direction of the objective function when the parameter space has a Riemannian metric structure and has a faster convergence speed. In particular, NGD can also be seen as an approximation of the Netwon method when the objective function and the manifold metric are compatible [8]. However, it is still impossible to directly compute the inverse of the Fisher information matrix for DNNs with millions or even billions parameters. Quite a few approximate approaches have been proposed. Under some independency assumptions, Martens approximate curvature (KFAC) method, in which the Fisher information matrix is approximated as a block diagonal matrix and each block matrix is further approximated as the Kronecker product of two smaller matrices. Then, KFAC was extended to convolutional neural networks [19], recurrent neural networks [20] and variational Bayesian neural networks [2]] and showed significant speedup during training. In addition, George al. Kronecker factorization (EKFAC) method. Gao et al. [15 23
and Grosse [1] proposed the Kronecker-
[22
proposed the eigenvalue-corree
proposed the trace-restricted Kronecker-factored approximate (TKFAC) method. These approaches approximations of the Fisher information matrix. What's 's more, some works have also considered large-scale distributed computing using NGD for training DNNs and shows excellent experimental performance [6 251.
all focus on the Kronecker-factoon
In this paper, our main focus is the NGD method. Motivated by the effectiveness of diagonal approximations and the significance of diagonal elements in the curvature matrix, we prioritize the diagonal information and integrate it into our approximation and introduce novel method, namely Layer-wise Natural Gradient Descent (LNGD). Our contributions can be given as follows:
Based on the block diagonal approximation of the Fisher information matrix, we propose layer-wise sample method more efficiently compute each block matrix corresponding to each layer. By assuming that the predictive distribution of the output after the activation function for each layer follows a Gaussian distribution, each block matrix can be directly computed using the inputs and the outputs separately, without having to perform a complete back-propagation.
For each block matrix corresponding to each layer, we further approximate it as Kronecker product of two smaller matrices, one of which is a diagonal matrix, while keeping the traces equal before and after approximation. With this operation, we further reduce the cost of computing inverse matrices while still preserving the main information of each block matrix.
In order to further accelerate the training, we propose an adaptive layer-wise learning rate by optimizing a quadratic model, in which parameters in the same layer share the same adaptive learning rate. Moreover, a faster approach of computing the adaptive layer-wise learning rate is also provided, making speed up training while maintaining computationally efficient.
Based on the novel approximation mentioned above of the Fisher information matrix and the adaptive layer-wise learning rate, we propose the LNGD optimizer for training DNNs. The global convergence analysis are also established under some assumptions.
We perform experiments on image classification and machine translation tasks. Numerical results show that LNGD converges faster than SGD, ADAM and KFAC, and LNGD provides an significant improvement in computational time savings when achieves convergence.
The rest of this paper organized as follows. Section gives the notations and introduces the NGD method. In Section 3 we propose novel approximation of the Fisher information matrix and the
2

---

adaptive layer-wise learning rate. Furthermore, we give the framework of LNGD and establish the convergence analysis. Section presents the results of experiments on image classification and machine translation tasks. The conclusion is drawn in Section
2 Notations and Preliminaries
In this paper, for matrix A, Ai to denote its (i.j)th entry, tr(A) denote its trace and ||A||] to denote its Frobenius norm. We use o and & to denote the Hadamard and Kronecker product of two matrices. In the following, briefly introduce the NGD method for training DNNs. During the training process of neural networks, the purpose is to find the vector of parameters which minimizes the loss function h0). If the loss function h(0) is chosen the the cross-entropy loss function, (()) can be given as h)()) E[ - logp(y x, 0)]. where p(y) x, 0) is the density function of predictive distribution (0), and y.y are the training inputs and labels, respectively. Next, we give the definition of natural gradient, which gives the steepest direction of the objective function when the parameter space has a Riemannian metric structure. The natural gradient is defined as F (O), where F is the Fisher information matrix given by
F
Vo log ply x, 0)Vo logp(y/x,
x
In Eq. (T) the input x is independently sampled from a distribution Qx with density function being ((x) and the label y is sampled from the predictive distribution (0). the following pages, we abbreviate y x))) as E unless otherwise specified. Consider neural network with L layers, for each layer / E [L] with [L] = {1, 2,... L}, we denote a and Wi as the input (the activation from the previous layer) and the matrix of weights of this layer, respectively. What's more. 0, = vec(W,) and 0 = (f...... 01) = (vec(W,) ..... vec W L.) =)) where vec(-) indicates vectorization of matrix. For convenience, we denote the derivative of the loss function with respect to 0 as DD = - Vo log ply x, 0). Then the Fisher information matrix can be expressed as F D|D00
Due to the high computational and storage costs caused by the inverse operation of high-dimensional matrices, it is impractical to directly compute F in the training of DNNs. The family Kronecker- factored approximations provides effective approach for computing F-1 of parameters in high- dimensional space, which is usually achieved by two steps. In the first step, by assuming that the parameters between different layers are independent, these methods approximate the entire Fisher information matrix as a block diagonal matrix, i.e.,
F = diag F1. F2. F2),
(2)
where F, D0, ] for any [L]. By this way, the Fisher information matrix can be approximated by L block matrices. This step transforms the inverse of the entire Fisher information matrix into the inverse of series of small block matrices. In the second step, these methods further approximate each block matrix as the Kronecker product of some smaller factors. This approximation can transform the inverse of each block matrix into the inverse of some smaller factors combining the properties of the Kronecker product.
3 LNGD: A Layer-Wise Second-Order Optimizer
In this section, we first introduce the layer-wise sample approximation strategy. Then, we present the details of adaptive layer-wise learning rate mechanism and give the specific framework of LNGD. Finally, elaborate theoretical analysis of LNGD's convergence is also provided.
3.1 Layer-Wins Wise Sample Approximation
For NGD methods to train DNNs, the Fisher information matrix can be approximated by block diagonal one according to different layers as given by Eq. (2). this approximation can be found in
19 22 and references therein. We call such block diagonal approximate Fisher information matrix the layer Fisher information matrix, which is computed based on a distribution Qx and
14
15
predictive distribution (0) as given in Eq. (.. To obtain the layer Fisher information matrix,
3

---

100
100
100
200
200
200
300
800
300
200
200
200
20
20
30
30
30
20
20
Figure 1: Comparison of the exact Fisher information matrix F and our approximation FINGD- We LNGD train MNIST on fully-connected neural network, whose architecture 196-20-20-20- 20-10. We show the results of the Fisher information matrix of the first layer with 20 units in top, which is 400 x 400 matrix. The bottom portion displays partially enlarged parts of the top marked with red square, which is 40 x 40 matrix. Within both the top and bottom sections, on the left is the exact Fisher information matrix F, in the middle is our approximation LNGGD. and on the right the absolute error between them. The brightness levels correspond to the sizes of the absolute values.
need perform a complete back- propagation to sequentially compute Fz, FL-1,.... F1, which still consumes much computing time.
In this subsection, we propose a layer-wise sample approximation of the Fisher information matrix, in which each block matrix F, is computed based on the Ith layer's prediction distribution Pay ai (0)) with the input a E Rd. of this layer and the input a e R4+11 of the 1)th layer instead of using the same predictive distribution Py Py)((0) for all layers. Specifically, for a,. we assume that the predictive distribution P lay, (0:) follows Gaussian distribution, which is usual used as prior by variational auto-encoder (26). F, can be computed by sampling from normal distribution with expectation being a and variance being I. Similar assumption can also be found in (27), in which the normality is also supported by central limit theorem under the independence assumption. By this layer-wise wise sample approximation, we can compute the layer Fisher information matrix without having to perform a complete back propapropngttton and thus improve the computational efficiency.
Next, we can give the formula of each block F, in the layer Fisher information matrix as
Fi E.DODDoT |
(DW::
[
E[(a--aa (gig] }} E I maddddxm xmid;
where DØ, = Ve logp(ay aa: 01), gi = V logp(a/ a, E with ss = Wi and p(a//a;-1,00 0.) is the density function of the distribution Pa, a; (0t).
In practice, the dimension of each block matrix F, is often still too large to directly compute its inverse matrix. Therefore, additional approximation methods are required to handle this computational difficulty. Suppose that the predictive distribution of a, follows Gaussian distribution with expectation being and variance being I, and each element of activation output is independent and identically distributed random number, then each element of partial derivative gi also independent and identically distributed. It is easy to show that F; can be seen matrix with di x di block matrices, in which each block is an my x my matrix and the off off-diagonal elements are zero. Therefore, F, can be approximated as
Fi Rd (a @ diag(gig) }]
(3)

---

A B for any two invertible matrices A and B, we can significantly reduces the computational complexity. Thus, some approaches have considered
Combining the property that (A @ B) -
approximating the Fisher information matrix as the Kronecker product of two factors [14 15 19 20
22]. Inspired by these works, we also approximate F, as the Kronecker product of two factor matrices
4 IR and V, E IR IRdixdi To get factor matrices 1 and v, we first replace diag (gg in
Eq. 3 by its trace and obtain Pt. Then we compute V while keeping that tr(Fi) = tr (2 @ v) Specifically, 2 is given by
D (a L1) & tr(diag(g:g) })] (aaa-1a/_1) x gi gi],
(4)
On the other hand, 1 can be computed by
L [[[([]] 1) x
(5)
[[ ay (g) 8:)]
Based on Eq. and Eq. 5 we can show that tr(FF) ur 4 & 1)
Fig. presents the visualization results of the exact Fisher information matrix F, our approximation FINGD. and the absolute error between them. Brighter pixels indicate higher values. From the left column in the top row, we observe the elements in the principal diagonal exhibit quite higher values, indicating their significance with rich information. Similarly, F LNGD can also emphasize the importance of the diagonal elements. The error figure reveals that the errors of the diagonal elements are small, which indicates that FLNGD provides a good approximation effect for the diagonal elements. Furthermore, to achieve a clearer visualization, we show the results of the partially enlarged area marked with red square in the bottom row. Here, we can observe more clearly that F LNGD achieves a favorable approximation effect on the diagonal elements. What's more, FINGD can also provide an effective approximation of the elements the auxiliary diagonals. These visualizations demonstrate the effectiveness of our proposed approximation in capturing the main elements of the Fisher information matrix. Therefore, our proposed approximation FLNGD is efficient and FINGD can retain most of information.
3.2 Adaptive Layer- Wise Learning Rate
In this subsection, we propose an adaptive layer-wise learning rate to accelerate training DNNs. We first consider the cases that use the same learning rate for all elements and the adaptive element-wise learning rate. Then we present the adaptive layer- wise learning rate scheme.
Suppose that d is the update direction of the function /: R² - IR at the iteration point 0², We first recall the gradient descent methods for getting the minimization of h, in which the update rule be given as 0 = l a "d! where a* is the learning rate, which can be chosen according to the value of the quadratic model
h(0* a d h(00)) (a Voh(O*)) + (a+))² (a*) Voh(O+)a
Once the update direction chosen, the minimizer of a can be given by
<d't Voh(0*)) (d+) dd*)
(6)
if TVah)\Vhh(0+)) is nonzero. If (gh(O)) is positive definite and d = (Vah(0"))
then a = 1, which leads to the classical Netwon method. In gradient decent methods, the learning rate is often regarded as the most important hyperparameter that highly influences model training. A
fixed learning rate may lead to slow convergence or suboptimal performance in some cases. Therefore, many works have considered using adaptive learning rate in gradient decent methods 6 28). In the following, we consider giving an adaptive element-wise learning rate automatically scaled by the direction d². In this case, the update rule of parameters is given by 8 O² -- dt = = a where x E R is the learning rate, D E Raxn is a diagonal matrix with (D*)a = (d ); and (D*); = when a j for i, E [n]] and 5." denotes the element-wise product. The second Taylor expansion of h(O Da) at iteration & is
h(0² )) zh D Voh(O*)) + TV'a"))Thh(0''''s
5

---

Taking the derivative of h with respect to at and letting it equal to O, we get
2D 2D DD D"Veh(0") = =0
which yields that
a = (Vah
D and (Bh( are positive definite.
Note that in Eq. (7). it is impractical compute the inverse of Vih(0*)D* directly for large- scale models due to high computational and storage costs. For second-order optimization methods in deep learning, some methods have considered approximating the curvature matrix by block diagonal one according to different layers [10 14 observed that parameters in the same layer have gradients of similar magnitudes. Therefore, common learning rate can be efficiently shared by these parameters (29 30). Inspired by these works, we propose novel adaptive layer-wise learning method as follows. Suppose that d
15
19 24).. What's more, some works have
((df), (dg) ...... (d (d)) is the update direction of T. layers neural network at the iteration point 0 ((0)), (0)),,... ..... T)) the update rule of @ is given as gh+1 0 - - where
D = diag(df d
d d
(8)
is block diagonal matrix and at E R is the learning rate. The approximate second Taylor expansion of h(O Dā) at iteration l: is
1 (
h(ok D*aat ~2((0h) @* Veh(O*)) +
(9)
HH H......, H and H Vg h00 for l [L] and the Hessian matrix is approximated by the block diagonal matrix H Taking the derivative of h with respect to at and letting equal to we get (D&)*T Voh(O") ( ( H D which yields that
where H diag(H)
a diag( O1 O2 ........ ④2) ((d)) T h(0*), (dh) T (o.h( .. (d Vo,h(0")) (10)
if + is nonzero, where θ = (df) H (df)) for L E [L].
If same learning rate is used for all layers, the same way of computing the adaptive layer-wise learning rate, we can get
(d+) (oh(Ok)
(11)
(d H
Theorem 1. Let g(0) and gr (0) be the approximate second Taylor expansions of h0 ad) and
h(O (D) as given in (9). where D E _ R& and lR are given by (8) 10 and [1] respectively, then we have (0) I g(0).
Proof. The proof is given in the appendix.
By Theorem we know that the adaptive layer-wise learning rate may lead faster decline in terms of function values. In our proposed algorithm, we choose d = (F*) "VVoh("") where F is the Hessian matrix. Then, the Fisher
Fisher information matrix and can be seen as approximation information matrix is approximated by a block diagonal matrix each block matrix is approximated by the Kronecker product of two factor matrices. each layer, the update direction de is scaled by layer-wise damping learning rate of according to (10). which given by
(d()) Voh(O))) a (d/) F + +
(12)
where > 0 is parameter. Using this adaptive layer-wise learning rate can accelerate layers with smaller gradients. Moreover, this approach can also avoid computing the inverse matrix in element-win learning rate (13) and remain computationally efficient.

---

3.3 Algorithm Schema
To effectively apply LNGD in training DNNs, several certain techniques need to employed. In this section, we primarily focus on introducing the damping technique, which is commonly used in second-order methods. Meanwhile, a simple method can be used to compute the adaptive layer-wise learning rate according to Eq. 12 since the cost of computing (df) d is relatively expensive. Finally, we discuss the utilization of exponential moving averages to enhance the training process.
new damping technique: Damping plays a crucial role in second-order optimization methods. Large damping can weaken the effect of curvature matrix, while small damping may cause computa- tional difficulty and inaccuracy since most eigenvalues of the Fisher information matrix are close zero and only small number of eigenvalues take on large values. To make training stable, we propose the following damping for the /th layer: X min (aa((r(())))))).... (Ft)/d4. 21), 12), where 1/1 and V/2 are two constants to constrain the minimum and maximum of damping, and di the number of weight parameters. In our method, F; is approximated as the Kronecker product of two factors i and 1 so we add the damping to each factors by $ c + Xy and L T A,, where Pi)/n. n, v1), 1/2) and X min(max (t(( 1t)/n, v2).
X
Compute the learning rate the faster: In order compute the adaptive layer-wise learning rate given in Eq. 12 more quickly, we turn matrix computation into vector computation. Specifically,
(d F = (
E
[D0,D0] Jak
E (..y)((ax.))
[(d)) D0,D0] d
(2.grp(2.0))
(13)
1 (D.))))
E (#.y<p(E.))
~ [(dd)) D0,))
where F is the empirical Fisher information matrix and N the number of samples. The empir- ical version of Fisher information matrix with no need for sampling from the model's prediction distribution, making it more computationally efficient.
Exponential moving averages: In line with previous studies, we incorporate exponential moving averages into our approach. This involves updating the estimate by combining the previous estimate, weighted by e, with the estimate calculated from the new mini-batch, weighted by c. That is
$ k+1
cb. k+1 + (1 -)
e) and b+1 <- eii + - e) e
(14)
In summary, our proposed algorithm is shown in Algorithm]]
3.4 Convergence Analysis
In this subsection, we give the convergence analysis of LNGD. Following the model used in previous works about analysing the gradient descent B3 neural network activated by the ReLU function with neurons in the hidden layer as follows:
33 and NGD [34 we consider a two-layer
1 m a,,(00,xx) x))
f(0, a(),a
m Σ
where 01, O2..... ..... Pm E Rdd are the weight vectors of the first layer, Rd is the input, o, R is the
weight of unit in the second layer and up() is the ReLU activation function, i.c., ((z) = max (0, z).
Let v [f(0, a, x,), f(0, a, x2)....,, f(0, a, x,))] T. In the following, we only give the result of
convergence of Algorithm|] [] the specific proof, which uses some conclusions in [36 37 38 39, is given the appendix.
Theorem 2. (Convergence rate of LNGD) Under the Assumption 7 and the assumption that
rank X) = d. If we set the number of hidden units m = Ω
Xx we i.i.d initialize v
0, N(O. vI), or unif([-1,+ for any E [m], and we set the step size a S (1+c)² Then with probability at least 1
(1-2c)
E over the random initialization, we have for k = 0, 1,2,...
yy
113 S (1- a)* a1
lly - -
Proof. The proof is given in the appendix.

---

Algorithm
LNGD
Require:
learning rate a, learning rate parameter μ,, damping parameter X, damping constraints VI 11, momentum parameter т, exponential moving average parameter 6, Fisher information matrix and its inverse update intervals FIM and TINV-
1: & <- O, m <- O. Initialize $ and 1 for any / E [Z].
2: while convergence is not reached do
3: Select new mini-batch
4: for all / E [L] do
5: if k: = 0 (mod TFIM) then
6:
Update the factors e and 1 using Eq. 14
7:
end if
8:
if k = 0 (mod TINv) then
9:
Compute the inverses of o, and L
end if
11:
Compute ((,) using backpropagation
12:
Compute the approximated natural gradient Φ
13:
Compute the adaptive learning rate a using Eq.
12
14:
S +- -000 c 6 Vo, 1(0)
m 7m
(Update momentum)
16:
0: + 0; + m (Update parameters)
end for
18:
k + k +
end while
20: return θ
4
Experiments
In order to verify the effectiveness of the proposed optimizer, we apply the optimizer to both image classification and machine translation tasks. We first present the optimization performance of our optimizer by comparing with several baselines. Then, we pay attention to the contribution of different modules of our optimizer by conducting elaborate ablation analysis, which is given in the appendix. Unless otherwise stated, the batch size for all experiments in the following is set to 256. The initial learning rate hyperparameters for all optimizers
tuned using grid search with values a E (lc- 4, 3e - 4, ..,... 3) The damping parameter À in KFAC [4] are tuned using a grid search with values (1e 6, le 4, 3e 4, of damping parameters parameter and the momentum correlating with KFAC and LNGD Furthermore, a weight decay of 0.004 is applied in all optimizers. All experiments run on a single A100 GPU using TensorFlow. We average the results of 5 runs and the hyper- aapee--rramtter settings for these optimizers are the best values randomly searched for many times.
..... le 1, 3e 1}. The minimum and maximum and 12 in LNGD are set to le 5 and le - 2. The moving average to 0.95 and 0.9, respectively.
4.1 CIFAR-10 Training
We first report the optimizing performance benchmark optimization methods 66]
CIFAR-10 [40], which standard task used to
43 44) Following these previous works, the changes of testing accuracy and training loss versus time well as epoch are reported in Fig. 2 and detailed statistics are shown in Table6 From Fig. 2 it can be observed that LNGD exhibits the most rapid decline training loss during the initial epochs and seconds. This suggests that LNGD effective in quickly reducing the training loss and reaching convergence. All optimization methods convergent at around 200 epochs. However, is observed that second-order optimization methods, such as KFAC and LNGD, achieve lower training loss compared to first-order optimization methods like SGD and Adam. In terms of testing accuracy, as depicted in Fig. (b) and (d), LNGD achieves a top-1 accuracy of 91% at the fastest rate. It only requires 36 epochs and 189.69 seconds to achieve this accuracy level. In comparison, as presented in Table SGD and ADAM require least 100% and 30% more epochs and time, respectively, to achieve similar accuracy. Relative to KFAC, LNGD reduces the number of epochs and time by around 20% and 21%, respectively. Furthermore, as shown in Tabld LNGD gets the highest final testing accuracy after convergence.
41
42

---

DALK UNGO
LIN
Figure 2: Numerical performance on ResNet-18 with CIFAR-10.
Table 1: Detailed statistics CIFAR-10 when top-1
testing accuracy achieves 91%.
Epoch
Total Time
Time Per Epoch
Acceleration
Best Test Acc
SGD
79
268.67s
3,4s
29%
91.88%
ADAM
72
248.83s
3.77s
23%
92.62%
KFAC
45
241.86s
5.87s
21%
93.34%
LNGD
36
189.69s
5.08s
93.61%
2
175C
2000
2000
Figure 3: Numerical performance on ResNet-50 with ImageNet.
4.2 ImageNet Training
We extend examination of optimizer efficacy larger image classification dataset, ImageNet-1K [45] The changes of lesting accuracy and training loss versus time and epoch are reported in Fig 3 and Tablc2 The results show that the LNGD optimizer highly efficient in training large image datasets in terms of both speed and accuracy. LNGD, which requires only 36 epochs and 6,46 hours, is much faster in achieving the top-1 testing accuracy of 75.9% than other baselines. This is significant improvement over SGD, which takes 100% more epochs and 75% more time to reach the same accuracy level. As for Adam, exhibits a rapid decrease loss during training and reaches convergence at fast rate. However, the best achieved testing accuracy is only 74.05%, indicating that when training large-scale image tasks, a trade-off between efficiency and effectiveness needs to be considered. Compared to KFAC, although LNGD is better for only 3 epochs, it leads to 19% reduction in terms of the computing time. The training loss results further support the efficiency of LNGD, maintains the fastest rate of decline during the initial stages of training and ultimately yields the lowest training loss upon convergence. Overall, the results suggest that LNGD is a highly efficient optimizer for large-scale image classification tasks, providing faster convergence and better accuracy than other commonly used optimizers.
4.3 Transformer Training
In this experiment, apply LNGD the Transformer-Bi Big model [46] with 213.7M parameters. The training datasets WMT English-Germnn German machine translation corpus [46]. We use Bleu 47] as the evaluation metries, which is frequently used in machine translation tasks. The setting of learning rate updating strategy for SGD, Adam, KFAC and LNGD are the same as in ImageNet training.
In Fig 4 and Table3] we present the comparative evaluation of the performance of LNGD against SGD, Adam, and KFAC in terms of testing accuracy and training loss. ADAM demonstrates superior performance over SGD, evidenced by more rapid decrease in training loss and a lower converged loss value. This observation aligns with previous empirical findings that ADAM highly effective for transformer models. KFAC exhibits further enhancements in performance compared to Adam,

---

Table 2: Detailed statistics on ImageNet when top-1 testing aceuracy achieves 75.9%.
Epoch
Total Time
Time Per Epoch
Acceleration
Best Test Acc
SGD
78
11,28h
520.55s
43%
76.47%
ADAM
74.05%
KFAC
39
8.02h
739.93s
19%
76.58%
LNGD
36
6.46h
646.44s
76.73%
ARAM LICE
2TCC
Trandeccch
Tlaneleccch
Figure 4: Numerical performance on Transformer with WMT.
yet it does not surpass the efficacy of LNGD. LNGD outperforms its counterparts with the swiftest reduction in training loss and the highest convergence rates. In terms of testing accuracy, measured by the Blcu score, LNGD achieves a top-1 Bleu score 32% with remarkable efficiency, which is able to reduce the required steps by approximately 24% and computing time by 16% compared to Adam. When compared to KFAC, LNGD still shows significant improvements, reducing the steps by around 14% and computing time by 24%. As for SGD, it cannot reach the top-1 Bleu score of 32% and the best testing accuracy is only 31.8%, which indicates that SGD is not good choice for large language processing tasks. In summary, the results provide strong evidences for the effectiveness of LNGD as an optimization algorithm for transformer models and shed light for large practical NLP tasks where time and computational resources quite limited.
Table 3: Detailed statistics on WMT when Bleu achieves 32%.
Step
Total Time
Time Per 1K Acceleration
Best Test Bleu
SGD
31.87%
ADAM
25K
3.39h
488.16s
16%
33.05%
KFAC
22K
3.75h
613.63s
24%
33.45%
LNGD
19K
2.85h
540s
33.55%
5
Conclusion
In summary, propose novel NGD optimizer named LNGD for training DNNs, specifically targeting the computational inefficiencies that impede the practical application of conventional natural gradient techniques in large-scale neural networks. Our approach strategically computes Fisher information matrices for each individual layers using sample approximation and dynamically adjusts learning rates leveraging curvature information. This method facilitates a more refined representation of the optimization landscape at the layer level. Besides, we provide convergence analysis of LNGD. Experimental evaluations indicate its competitive performance relative existing state-of-the-ar optimizers. This work hold significant potential for enhancing the efficiency and scalability of training processes in deep learning frameworks.
References
[1] Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics, pages 400-407, 1951.
[2] Yu E Nesterov. A method for solving the convex programming problem with convergence rate of L. In Doklady Akademii Nauk SSSR, volume 269, pages 543-547, 1983.
10

---

[3] Ning Qian. On the momentum term in gradient descent learning algorithms, Neural networks, 12(1):145-151, 1999.
[4] Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprini 1212.5701 2012.
[5] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26- 31, 2012.
[6] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2014.
[7] Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 13:18795-18806, 2020.
[8] Nitish Shirish Keskar and Albert Berahas. ADAQN: An adaptive guasi-Newton algorithm for training mns. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 1-16, 2016.
[9] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney, AdaHessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAI conference artificial intelligence, volume 35, pages 10665-10677 2021.
[10] Donald Goldfarb, Ren, and Achraf Bahamou. Practical quasi-Newton methods for training deep neural networks. Advances in Neural Information Processing Systems, 33:2386-2396, 2020.
[11] James Martens, Deep learning via Hessian-free optimization. In International Conference on Learning Representations, volume 27, pages 735-742, 2010.
[12] Oriol Vinyals and Daniel Povey. Krylov subspace descent for deep learning. In Artificial Intelligence and Statistics, pages 1261-1268. PMLR, 2012.
[13] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2)2511 276, 1998.
[14] James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning, pages 2408-2417, 2015.
[15] Kaixin Gao, Xiaolei Liu, Zhenghai Huang, Min Wang, Zidong Wang, Dachuan Xu, and Fan Yu. A trace-restricted Kronecker-factoree approximation to natural gradient. In Proceedings the
AAAI Conference on Artificial Intelligence, volume 35, pages 7519-7527, 2021.
[16] Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Chuan- Foo, and Rio Yokota. Scalable and practical natural gradient for large-scale deep learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(1 14(1)):0444 415, 2020.
[17] Minghan Yang. Dong Xu, Qiwen Cui, Zaiwen Wen, and Pengxiang Xu. An efficient Fisher matrix approximation method for large-scale neural network optimization. JEEE Transactions on Pattern Analysis and Machine Intelligence, 45(5):5391-5403 2022.
[18] James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21:1-76, 2020.
[19] Roger Grosse and James Martens. A Kromecker-factored approximate Fisher matrix for convo- lution layers. In International Conference Machine Learning, pages 573-582, 2016.
[20] James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.

---

[21] Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International conference on machine learning, pages 5847-5856, 2018.
[22] Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a Kronecker factored eigenbasis. Advances in Neural Information Processing Systems, pages 9550-9560, 2018.
[23] Kaixin Gao, Zheng- Hai Huang, Xiaolei Liu, Min Wang, Shuangling Wang, Zidong Wang, Dachuan Xu, and Fan Yu. Eigenvaluuecorrected natural gradient based on a new approximation. Asia-Pacific Journal of Operational Research, 40(01):2340005, 2023.
[24] Mengyun Chen, Kaixin Gao, Xiaolei Liu, Zidong Wang, Ningxi Ni, Qian Zhang, Lei Chen, Chao Ding, Zhenghai Huang, Min Wang, et al. THOR, trace-based hardware-driven layer- oriented natural gradient descent computation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pages 7046-7054, 2021.
[25] Lin Zhang, Shaohuai Shi, Wei Wang. and Bo Li. Scalable K-FAC training for deep neural networks with distributed preconditioning. IEEE Transactions on Cloud Computing, 11(3):2365- 2378, 2023.
[26] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[27] Krishnakumar Balasubramanian, Pinar Donmez, and Guy Lebanon. Unsupervised supervised learning ii: Margin- based classification without labels. In Proceedings of the Fourteenth Inter- national Conference on Artificial Intelligence and Statistics, pages 137-145. JMLR Workshop and Conference Proceedings, 2011.
[28] Frédéric de Gournay and Alban Gossard. Adaptive scaling of the learning rate by second order automatic differentiation. arXiv preprint 10X1v:2210.14520, 2022.
(29] Bharat Singh, Soham De, Yangmuzi Zhang, Thomas Goldstein, and Gavin Taylor. Layer- specific adaptive learning rates for deep networks. In 2015 IEEE 14th International Conference on Machine Learning and Applications (ICMLA), pages 364-368. IEEE, 2015.
[30] Boris Ginsburg, Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Huyen Nguyen, Yang Zhang, and Jonathan M Cohen. Training deep networks with stochastic gradient normalized by layerwise adaptive second moments. 2019.
[31] Simon S. Du, Xiyu Zhai, Barnabás Póczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In 7th International Conference on Learning Representations, 2019.
[32] Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Proceedings of the 36th International Conference on Machine Learning, volume 97. pages 1675-1685, 2019.
[33] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterizee two-layer neural networks. In Proceed- ings of the 36th International Conference on Machine Learning, volume 97, pages 322-332, 2019.
[34] Guodong Zhang, James Martens, and Roger Grosse. Fast convergence of natural gradient descent for over-parane wver-parameterize neural networks. In Advances in Neural Information Processing Systems, volume 32, 2019.
[35] Ryo Karakida and Kazuki Osawa. Understanding approximate Fisher information for fast convergence of natural gradient descent in wide neural networks. In Advances in Neural Information Processing Systems, volume 33, pages 10891-10901, 2020.
[36] Alberto Bernacchia, Máté Lengyel and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. In Advances in Neural Information Processing Systems, pages 5945-5954, 2018.
12

---

[37] Joel A Tropp. An introduction to matrix concentration inequalities. arXiv:1501.01571, 2015.
arXiv preprint
[38] Jssai Schur. Bemerkungen zur theorie der beschränkten bilinearformen mit unendlich vielen veränderlichen. Journal für die reine und ungewandle Mathematik, 1911(140):1-23 1911.
[39] Shuangzhe Liu and Gotz Trenkler. Hadamard, Khatri-Rao, Kronecker and other matrix products. International Journal of Information and Systems Sciences, 4(1):160-177, 2008.
[40] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images.
2009.
[41] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangging Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 7706.02677 2017.
[42] Peng Xu, Fred Roosta, and Michael Mahoney. Second-order optimization for non-convex ma- chine learning: An empirical study. In Proceedings of the 2020 SIAM International Conference on Data Mining, pages 199-207. SIAM, 2020.
[43] Yuchao Li, Shaohui Lin, Jianzhuang Liu, Qixiang Ye, Mengdi Wang, Fei Chao, Fan Yang, Jincheng Ma. Qi Tian, and Rongrong Ji. Towards compact cnns via collaborative compression. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pultern Recognition, pages 6438 6447, 2021.
[44] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016
[45] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large- scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-255. leee, 2009.
[46] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[47] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 31 111311, 2002.
13

---

A Notations
In this paper, we denote [n] = (1, For a matrix A, we use A.j to denote its (i, entry, tr( (A) to denote its trace, (i) and max (A) ) to denote its smallest and largest singular value, | I to denote its Frobenius norm and | | to denote its spectral norm. A is positive semi-definite, Amin(A) and Amax(A) denote its smallest and largest eigenvalue, and define A Amax((), ///i
(A))/min (A), respectively. The identity matrix denoted as I. For vector a, ||al|2 denotes the Euclidean norm. We use (n, Σ) to denote the Gaussian distribution with mean l and covariance Σ. For two matrices
A |a a2, ................................ an] IRxx,
B= [b,, b.,..., bnl E RPx0
and @ denote the Hadamard and Kronecker product, respectively. The column-wise Khatri-Rao product * is defines as
A B laj b1, a2 2 ,
an @ bn] R x9.
Similarly, we can define the row-wise Khatri-Rao product * and we have (A B) (AT BT). Given event E, I{E) denotes its indicator function, i.e.,
KE} = {{
1, if E happens,
0, otherwise.
B
Comparisons and Explanations
B.1
Comparisons with Related Works
There have been some works to utilize the NGD or its approximations for training DNNs. One of the primary computational challenges lie in the storage and computing the inverse of the Fisher information matrix during NGD optimization. Recently, several studies have explored the adoption of the efficient cronecker-factoree approximation to the Fisher information matrix to address this computational challenge. The most related approaches to this work are the KFAC, EKFAC and [5F These works and LNGD all start with block-diagonal approximation of the Fisher information matrix. The differences among them are the approximations of the block matrix F,, By approximating the expectation of the Kronecker product as the Kronecker product of expectations, KFAC approximates F; as the Kronecker product of A = ([(aa-aa__)) and B = (((g)) )] with eigenbasis, EKFAC performs eigenvalue decomposition of the Fisher information matrix and re-scales the eigenvalues by S* to achieve better approximation, where S* is diagonal matrix defined by S = [ o (0)))))) and UA, Un are eigenvectors of A, B. TKFAC upproximates F; Kronecker product of two factors P and Q scaled by a coefficient& ō and keep the traces of each block equal. In this paper, we propose the LNGD, which approximates F, as a Kronecker product of matrix 4 and a diagonal matrix L which is computed by sampling from each layer. We summarize the above approximations of F, in Table
g = se log p(y By tracking the diagonal variance in the Kronecker-factoret
Table 4: Summary of some NGD optimizers
Optimizer
F
KFAC
A
B
EKFAC
(UA &
SB)
(U
&
T
TKFAC
8P 8 Q
LNGD
Φ x V
When these methods have the same update frequency, KFAC needs to compute two factor matrices A and B, and then invert them. However, B can only be computed after completely performing back- propagation. On the other hand, EKFAC modifies KFAC by incorporating eigenvalue decomposition to scale the eigenvalue during the inversion process. TKFAC, another variant KFAC, maintains the equality of traces of matrices before and after approximation. Both EKFAC and TKFAC involve increased computational requirements compared to KFAC. Our proposed L.NGD also requires computation and inversion of two factors & and V. However, the advantage of LNGD is that the
14

---

(a)
(b)
(c)
(e)
(f)
Figure 5: Comparison of the exact Fisher information matrix and the approximated Fisher information matrix of KFAC and LNGD. On the left is the exact Fisher information matrix, the middle the approximated Fisher information matrix, and on the right is the absolute error of these. The first row shows the result of KFAC, and the second row shows the results of LNGD.
matrix V is diagonal, and due to the use of hierarchical sampling. v can be directly inverse during the forward propagation phase, without the need to wait for the completion of the back bockk-ropagation process. This feature significantly reduces computational time.
In addition, THOR proposed in [24 also provides an optimizer using NGD for training DNNs. THOR mainly considers reducing the computational cost of NGD through two aspects. On the one hand, THOR gradually increases the updating frequency the inverse matrix Fisher information matrix and proposes trace- rrace-based updating rule for the Fisher information matrix of each layer. On the other hand, THOR approximates the approximated Fisher information matrix obtained by KFAC as some smaller matrices by splitting matrix dimensions. Our proposed LNGD first gives a layer-wise sample method to more efficiently compute each block matrix corresponding to each layer and proposes a novel approximate scheme of the Fisher information matrix. Furthermore, LNGD also adopts adaptive layer-wise learning rate to speed up training. The contributions and ideas of our proposed LNGD are different from THOR.
B.2
Comparisons Between KFAC and LNGD
Fig. 5 shows the visualization results of KFAC and LNGD. From Fig. (b) and (e), we can see that KFAC and LNGD all emphasize the importance of the diagonal elements in the exact Fisher information matrix. In addition, can also be seen clearly that KFAC still retains some elements the main diagonal, while LNGD does not, which also reflects that LNGD provides efficient approximation of Fisher information matrix with less computational cost in comparison with KFAC.
B.3 Illustration of the Gaussian Distribution Assumption
In Subsection 3.1, we assume that the predictive distribution Pa (0,) follows Gaussian distri- bution. To illustrate the validity of the Gaussian distribution assumption. we collect the output of two layers of the ResNet-18 network on CIFAR-10 and show the results in Fig. 6 Fig. (a) and (b) show the distributions of sample representation vectors' values in some dimension. Since we use the ReLU activation function, the obtained distributions are accord with the Gaussian distribution in
the positive quadrant. Fig. (c) and (d) show the distributions of values of sample representation vectors' Fuclidean norm, from which we can see that the two distributions can also approximated Gaussian distributions.
15

---

(a)
(b)
(d)
Figure 6: Illustration of Gaussian distribution.
C Proof of Theorem 1
Proof. Since g(0) and (0) the approximate second Taylor expansions of h(O ad) and h(0 Dāx), have
1 Σ
9L ( = h(0) - ad, Voh(O)) + 2 a Hd - h(()) +
(d Vo,h(0)
1=1
(o))) (d (d vo)o)(()) d H,d:
Σ add Hd Hid:
0.
Hdd Hd
This completes the proof.
D Convergence of LNGD
In this section, give the convergence analysis of LNGD. Following the model used in previous
works about analysing the gradient descent [1 32 33 and NGD [34] 35), we consider a two-layer
neural network activated by the ReLU function with m neurons in the hidden layer as follows:
f(0, a,)
(,,
where 01. .2..... On E IRd are the weight vectors of the first layer, x E Rd the input, a, E R the weight of unit in the second layer and (p.) is the ReLU activation function, i.e., (o(x) = max(0, For convenience, we define O [0 0 ..... IT E Imnd We first initialize the parameters randomly by
0, O0 II), a, unif (- Vr E [m].
where 0 < controls the magnitude of initialization.
Given the training dataset S = {(x,, yi) containing (input, target) examples (x,, W:). Following 31 3 34 we make the following assumption for the data.
Assumption 1. For all i, ||x;1|2 and ly:l O()). For any # x, H Xj.
In this subsection, we mainly focus on the mean squared error loss (MSE) function
((
Σ (f(0, a, x.) - yi)²
2n
m
(0, x) y1)
m
fix the weights of second layer and only optimize the weights of first layer. Then the update rule of NGD can be written as
Following 31 3 34
0 0² -a ((F)) V )
16

---

As shown in [18], if the network's s predictive distribution is in the exponential family, the Fisher information matrix equivalent to the generalized Fauss-Newtor matrix, which is defined by
(.,, J HcJ,],
where H is the Hessian matrix of the loss function ((O) with respect to the prediction f(O, x;) and J; is the Jacobian matrix of f(0, a, x;) with respect the parameters O. Under our setting that (()) is the MSE loss function, the Hessian matrix He is the identity matrix 1. J; can be computed by
J: ((,
Vo.. s0, a, x.)x;)))
where
Or 1 x; 0)x;, Vr [m].
V0,,(0,a,x.) (0,
(15)
vm
Let J (J1, J2
J.) Raxdd then the Fisher information matrix can be written
F = (x. wx.. J HcJ,) = mss J JJ J.
J.
As discussed in [34]. when m the Fisher information matrix is singular. So in this case, use the generalized inverse given in [36
F nJ (JJ (JJ
(16)
and the update rule of NGD can be written as
g 8
F (J*) ((* (v* - y),
where y
W1. 12
V/ xn)
and
U1, U2....y Un]
[S(0, a, f(), ((), f(0, in this subsection, since we fixed the weights in second layer, and the Fisher information matrix of this model is approximated by
Consider the two-layer neural network described
P v
For simplicity, we ignore the index of layer. Define
X X1,
Xn1 gadd
and
G = (X01), 6 (X02)..... (X0m)) E =x"
where X is the input matrix formed by input vectors and G is the pre-activation derivatives matrix whose entry is given by
1
G
1xx 0, 20.
Vie [n], E [m].
(17)
Vm
Then we have and
Vef a, X G xa
JJ
X T
(X G) R mndxmdd
n
n
The formulas of + and W can be given as
X
(GG I)X Raxd
(18)
G o IIG
V
(19)
TI tn
XX ool)o (GG I))
Under the Assumption 1 we have XXX )a = E [d]. Therefore, Eq. simplified to
18 and Eq.
can be
w = G
Φ = Zx.G.
nn
where 1 = T o I) (GGT o I)) and Z X (GGT 1)X. Finally, the update rule of
LNGD can be given by
g++1 = g+ 0+ (Z & ((G*) T (v* y).
To analyze the global convergence of LNGD, we need the following two conditions given in [34]
17

---

Condition 1. The matrix GT)) is positive definite.
To verify this condition, we need the following two lemmas.
Lemma Define k Vo X; 2xz20 0,0%x, 2 0}}, i,J E [n]] then we have the
matrix K is strictly positive definite.
The result of this lemma has been given and discussed matrix K whose entry is given by
34 We define Ak ( 2 0 and
1 Σ T(O, x; 2 0, 0) (GG )o Vi,j E (m).
Kij
Then we can show Condition Tholds by the following lemma.
Lemma 2. If m = Q log )).
we have with probability at least 1 that Amin (K(0)) 2 2 }>K-
Proof.
Note that (0) can be written as the sum of random symmetric matrices
((0)= Σ K(0,), K,,(0,) _ T(O 0, 0,0, 0)-
Furthermore, we have KO0,) are positive semi-definite and ||((0,)l22 = tr(K((,)) S m. Thus, by the matrix Chernoff bound (37], we can obtain
(- A
P Amia (K(0)) - (1 A
S exp
42
Let & n exp (-4 Akm) we have m 52 log ) Proof complete.
This proof is similar to the Lemma 6 33], the difference is the definition of K. For completeness of the proof, also give the detailed proof here. By this lemma, we have Amin (O)
Amin (K)) 2 Ax > 0, which implies that Condition holds. Next, we give the other condition.
Condition 2. For all parameters 0 that satisfy 110
0 S
21y-vv Kzx.a. there exists 0 S <
Vxx/2
such that
J J- 22c VAK
KZx
To show this condition holds, need the following lemma.
Lemma 3. 134) For all weight vectors O that satisfy 0 00 < R. we have probability at least
1 that
2n R2/3
| - 1
/2/332/33m1
64 / n 1| 1/2/3 2/3 "2x
By taking R = 21ly-voll vll
KZx.c we have |J
J
/ Therefore, if we let
/ Avve
f|y
0 1333xe
12
the Condition holds. What's more, have probability at least 1
that ||y
v
O which has been given in 31. Thus write the condition of as
m Ω
Before giving the main result, we first give some necessary lemmas.
n1 n'r2x Lemma 4. =f
) for all parameters O that satisfy ||0 - -0012 S 21y--v Riz o
vvo/2
have probability at least that Amin (K)
VxK/2.
18

---

Proof. Combine Eq.
15). Eq. 17 and Assumption we have || G JJ J Let
ntm x
and H 2 volla
m = Ω
VAk/2 K c by Condition and Lemma 3 we have
|G G - || G(O))l
S J J 1
AK
8 Z
2
8
Therefore,
mmin (G) 2 Fmin (G())) - G G(()]l2
3
V2
AK
VλK:
Note that the large m is, the smaller c is. Therefore, we can choose a slight larger m satisfying this inequality. So we have
Ak 2
Amin (K) = -miin (GG 2
Proof complete.
Lemma 5. 38
Let A and B be two positive define matrices, we have
Amax (A o B) mmax A
(inxx().. Amin (A B) min Asi Amin Amin.B).
Lemma 6. (39 Let @ denote the Kronecker product and * denote the column-wise Khatri-Rao product, we have
(A 2 B) D) D)= AC BD,
(A B) B) A
Now, we give the convergence analysis of LNGD.
Theorem 3. (Convergence rate of LNGD) Under the Assumption 7 and the assumption that
rank(X) (X) = d. If we set the number of hidden units = - n. K2x we i.i.d initialize D
0, NO vI), Or unif[[-1, +1}] for any r E [m]] and we set the step size a S (1-2c) Then with probability at least 1 e over the random initialization, we have for k = 0, 1, 1,2,...
(1+c)2
yy 11 < (1 (1=a)* ||y - v
Proof. Consider the predictive error at the (A: 1)-th iteration, we have
yy
k+ 11 = lyy * + v -
(20)
* 11 2 v²)T(p (v - + || v
v²)
Next, we need to estimate the bound of prediction
k+1
* v(0 (FFFt) (J (ar))(v++ -v-- +
aJ'( (F (v+ - -y
0=
aJ (J( T (v" - y)ę
0=0
a(Jh - (P ((J+) (v* y)
5=0
aJ* (Ft) (J (v+ -y)
Term
(f( (J+ - -5){ (F+)) (v y);
Term
19

---

where J
dv( 80
and 0 = +
(1-
0h (JK)
(vk y).
We first analyse Term
We omit the index in J. G and for simplicity.
Term 1 = -aJFFFF- -aJJ--JTTy" y)
=a * G) @ (G * GT - vn)
=a(1XZx, X GG G) - -
=a(1XZx,xX o oI(yyvv)) -
The second equation follows the update rule of LNGD. The third equation is obtained according to the properties of Kronecker, Hadamard and Khatri-Rao products given in Lemma The last equation uses the definition of generalized inverse as given by Eq. (16). By Lemma 5 we have
Amax XZx X o I) < max X X.G G,
)ii)max())
(21)
I Amax (x,a) max(XX )==
Amin (Zx.G)
Therefore, bound Term by
[eerm 1|2 ||a(nXZxx x, X
o II v/
S aml x X I|l2lly v²112 5
07)
v²
(22)
Amiu (Zx,c)
Based on the Condition
2 we have the following inequality
(J - JE)G
| J J6 = || -J
E=0
2=0
1 J0ll2 + || JJ
(23)
/2c VTK < c
Amin(GGT]
2
KZx.a
KZx
Next, we bound Term 2. By Eq.
(23p. we have
Term 2112
a
(J - J5){ -
(F (J -y)
E-0
00
Amin DGn ([ ) - -
ZZx.
ca
- Amiu (GG((G)))((7 (GT G) G -
KZx.G
0a
Amiu (GG (GGT)llIZxox (GT - v=112
Kix,c
c0
Amiu (GG (GGT)|lnZZl X *G (GG GG ||2ly v*112
KZx.G
coO
Amin (GG T (GG 1
(24)
KZx.G
Define A Z -
XT G (GGTT) then we have
|||| max (A)
Amax A ATA) A)
(25)
Xmax XZx I
Ix XT
XT
(GGT)-1).
Similar to Eq. 21 by Lemma
we can prove
Al]2 S
(26)
Xmii (Zx,G) Amiul GG
20

---

Combine Eq.
24
and Eq.
26
we have
[Termal
S
CO7)
Aunin (GG)
1 Amia (Zx,c) Amin
1
by - v²
KZx.c
(27)
6007
122
Amax
ZZx
Combine Eq.
20
Eq.
27 and Eq.
22 we can obtain
1l
2 = yy
12 2(y v²) (o -²) + // - v0112 112
= Fly v 112 20 - v²) J*) T - v²)
+2a - v²) T ( (J+ - -6) (F (J+)) T(y (y v²)
{=0
+ +l|vet1 v 12
20m
200m
+ X (ZX,G)
Amax (Zx,a)
Amax (Zx,c)
20a2112
+ Amax (Zx Amin (ZX.G)
+
I1y - vv
N
(ZX,G)
In the last second inequality, we use the fact that Amiul (XZx
11))
2 Xmax
Let
o01)
2com
Q217
Amax ZZ
Amax
(Zx.G)
Xa (Zx,G)
20a21²
c'alp²
+ Amax (Zx, Amin (ZX,G)
S 0,
max (Zx,
we have
(1 20) Amax (Zx,c)
(1 2c)
c)+
(1 + c)2
and
lyy v S aa
-
This completes the proof.
So far, we have already proved Theorem S under the an assumption that the parameters stay close to the initialization point. We now verify this assumption by the following lemma.
Lemma 7. If Conditions and hold, then as long as Amin (GGT) { we have
| 00 y
voll l|
Kzx.a
V //K/2 2
Proof. By the update rule of LNGD, have k
10
-00 =
Σ (J' - v²)
S
I Z T (G v'll2
1=0
k
< 0
Amin (Zx,G) vAK/2 lly- v'll2 v2212
1=0
2/AK
(1 (1-X a
1/2
S
lly - vo112
1=1 Amiin (Zx,c)
Amnve (Zx,G)
2
v01/2 Amax (Zx,c)
21ly -v'12
KZx.a
/ /xx/2
Amii (ZX
Vxk/2
This completes the proof.
21

---

E Experiments
E.1 Sctup of CIFAR-10
The training of ResNet- 18 44 on the CIFAR-I0 10 40 dataset serves as a fundamental experiment in the field of image classification. In this subsection, we present comparison of LNGD with several established baselines including SGD with momentum (referred to as SGD), ADAM (6], and KFAC. We follow the standard experimental settings and employ a commonly used data augmentation scheme involving random crop and horizontal flip. The initial learning rate is multiplied by 0.1 every 40 epochs. The update intervals for the curvature matrix and inverse matrix correlating with KFAČ and LNGD are set to be 100. All experimental runs are conducted over duration of 200 epochs.
E.2 Setup of ImageNet
The implementation of ResNe150 [4 follows the TensorFlow version which can be found in the website 31
We use the linear warmup strategy [1 in the first 5 epochs for SGD, ADAM and KFAC. The update intervals for the curvature matrix and inverse matrix correlating with KFAC and LNGD are set to be 500. For SGD and Adam, the max epoch is set be 80, while for KFAC and LNGD, the max epoch is set to be 50. SGD uses the cosine learning rate updating strategy and is set to be a = 0.001 + 0.5 (ag - 0.001) * (1 + c2 * 0.47 + * t/max_epochi) where is the number of epochs. For Adam, KFAC and LNGD, the learning rate uses the exponential updating strategy ag = au * (1 - t/ _eaa______)___ where is decay rate E f2,3, 4, 5, 6}. ag is the initial learning rate tuned using grid search with values E {le - 4, le .......1 1).
E.3
Ablation Analysis
E.3.1
Setup
In this subsection, to further elucidate the contributions of distinct components within the LNGD, a series of ablation studies are performed. The ablation experiments aim to isolate the effects of adaptive learning rate and sampling optimization the LNGD's performance. The variant denoted as LNGD-Ir corresponds to the iteration of the algorithm that employs an adaptive learning rate, but does not incorporate sampling optimization. Conversely, LNGD-sample represents the iteration that utilizes sampling optimization, but does not implement an adaptive learning rate. These ablation studies are executed on the ImageNet-1 IK dataset, All hyperparameters are maintained consistent with those outlined in the ImageNet training section.
E.3.2 Results
The results of the ablation experiments, as shown in 7 and Table5] reveal some interesting findings. Specifically, analyzing the training loss and testing accuracy versus epoch, we observe that LNGD-Ir achieves the fastest decrease in training loss and the most rapid initial increase in testing accuracy within the initial few epochs. This can be attributed the fact that LNGD-Ir computes the exact Fisher information matrix at each epoch without using any approximation sampling strategy. However, this advantage comes at the cost of increased computational complexity, leading to a 15% increase in the time required to reach a top-1 testing accuracy of 75.9% compared to LNGD, which employs both the sampling approximation strategy and the adaptive learning rate strategy at each layer. Moreover, LNGD-Ir also takes 5% more time compared to LNGD -amplle which only utilizes the sampling approximation strategy. Notably, LNGGD-sampl exhibits the slowest decrease in training loss and increase in testing accuracy during the initial epochs due to its approximation sampling of the Fisher information matrix at each step. Nevertheless, when considering the time dimension, LNGD-sample still achieves a faster speed compared with INGD-Ir in reaching final testing accuracy of 75.9% due to the significant reduction in the computation of the exact Fisher information matrix. In contrast to LNGD, LNGD- sample takes 9% more time to reach testing accuracy of 75.9% due to the absence of automatic scaling learning rate. In conclusion, considering the constraints of limited computational resources and time, LNGD demonstrates superior optimizing performance.
22

---

Table 5: Detailed statistics abalation study when top-1 testing aceuracy achieves 75.9%.
Epoch
Total Time
Time Per Epoch Acceleration Best Test Acc
LNGD-Ir
35
7,43h
764.39s
13%
76.50%
LNGD-ssmple
41
7.06h
619.86s
9%
76.57%
LNGD
36
6,46h
646.44s
76.73%
5.0
4.
70
LNGD comple
Thintag Log
3.5
3.0
2.5
Tod
LN30
2.0
1NGO
1.0
20
50
20
10
20
30
Epoch
Epoch
LNNED
4.5
INCD
4.0
Thtning Lese
F0
3.0
2.5
a2
LNGD LNGD lr LNGD smmlle
2.0
15
75.9
1.0
10000
20000
30000
10000
20000
30000
Time/Seords
TennsSeconds
Figure 7: The optimization performance of variants of LNGD.
E.4 Results of More Comparisons
Table 6: Detailed statistics on CIFAR-10 10 when top-1 testing accuracy achieves %.
Epoch
Total Time
Time Per Epoch
Acceleration
SGD
79
268.67s
3.4s
29%
ADAM
248.83s
3.77s
23%
KFAC
45
241.86s
5,87s
21%
EKFAC
41
247.64 447.64s
6.04s
23%
TKFAC
39
239.20s
5.98s
20%
NG+
204.45s
5.11s
7%
LNGD
36
189.69s
5.08s
In order to further validate the effectiveness of LNGD, we conduct additional experiments on the CIFAR-10 dataset, which three methods including EKFAC (22). TKFAC [15], and NG+ [17] are added for comparison. The detailed statistics are presented in Table 6 From this table, we observe that LNGD achieves a testing accuracy of 91% with the fewest epochs and the shortest total time. Furthermore, LNGD exhibits the smallest computational time epoch. Additionally. due to the efficient Fisher information matrix approximation strategy adopted by NG+, can significantly reduce the computational time compared to EKFAC and TKFAC.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Due to space constraints in the main text, we have not created a separate "Limitations" section. However, for each theory presented, we explicitly outline any strong assumptions and provide justification for the validity of these theories. Additionally, thorough comparison with related work is furnished in the appendix.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: For each theoretical result, we provide the full set of assumptions and a complete (and correct) proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide every detail of the the methodology and experimental setup to ensure the reproducibility of the main experimental results presented in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
25

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [No]
Justification: Due to the involvement of proprietary code resources, the disclosure of such materials must adhere to the company' relevant disclosure processes. If necessary, data and code can be made available upon request.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips.co tttpcc//pbbb. muullc//guddee
wttp///ooeeesssiiiooco wommissssoooiccy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specify all the training and test details necessary to understand the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: For each experiment, we average the results of runs and the hyper-paramett settings are the best values randomly searched for many times.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
26

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide sufficient information on the computer resources needed to repro- duce the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpu://neurips.. http://p....cpulllccccccccccaaaa
Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
27

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: For each relevant work, whether it be code or data, we have provided appropri- ate citations.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
28

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, batasets.papersstt paperswi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We present optimization method, which does not involve the release of new
assets.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
29

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "NcmNFN8yjo", "reviewer": "Reviewer_HWh2", "summary": "The paper introduces a new optimization algorithm called LNGD (Layer-wise Natural Gradient Descent). This optimizer aims to enhance the training efficiency of deep neural networks by approximating the Fisher information matrix in a computationally efficient manner and introducing adaptive layer-wise learning rates.", "questions": "See Weaknesses.", "limitations": "See Weaknesses.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- **Research Significance**: The paper addresses a highly valuable problem in the field of deep learning training. Improving the efficiency and effectiveness of training deep neural networks with second-order methods.\n- **Practical Utility**: The proposed LNGD method shows practical promise. By reducing computational costs and introducing adaptive learning rates, the method can potentially be applied to many real-world training scenarios rather than KFAC.", "weaknesses": "1. **Assumption of Gaussian Distribution**: The assumption that $P_{a_l \\mid a_{l-1}}$ follows a standard Gaussian distribution lacks strong motivation and empirical evidence. Even for a single layer in a random linear neural network, this assumption is not convincingly justified. Further explanation or practical case demonstrations where $P_{a_l \\mid a_{l-1}}$  follows or approximates a Gaussian distribution are necessary.\n2. **Convergence Analysis**: The convergence analysis results are unusual. Typically, new optimization algorithms share the relationship between convergence speed and iteration steps, assuming a certain smoothness. The approach result, similar to NTK, is unconventional and may not highlight the superiority of NGD since different optimization algorithms tend to have similar rates with a large number of hidden units.\n3. **Complexity of Adaptive Learning Rate**: The method for calculating the adaptive learning rate is overly complicated. It needs to be shown whether the proposed Adaptive Layer-Wise Learning Rate has practical advantages over optimizers like Adam or LAMB. Ablation studies comparing this approach with existing LR schedules (e.g., linear warmup + cosine decay) should be conducted.\n4. **Computational Overhead**: The computation of sample-wise metrics in Eq. (13) poses significant memory and computational overhead. It is necessary to verify if the final equation holds true, specifically whether the squared term's batch average is equivalent to the mean of the squared terms.\n5. **Comparison with Latest Methods**: The experiments lack comparisons with some of the latest KFAC-variant methods. Additionally, current results omit several key metrics, such as peak memory usage, and selected baselines appear weak. For example, advanced training recipes for smaller models (like ResNet-50) on ImageNet typically achieve over 80% accuracy, making comparisons with sub-75% settings less meaningful.\n\nMinor Comments\n- **Layer Outputs**: Clarify the relationship between the outputs of the previous layer and the inputs of the next layer, as this affects the solution for $g_l$ in Eq. (3).\n- **Notation Consistency**: The notation $F_ {LNGD}$ is not defined.\n- **Citations**: The citations appear to be added post hoc, with critical parts lacking references. The related work section is also absent, suggesting a possible unfamiliarity with the latest advancements in the field, as evidenced by the absence of citations from 2023 onward.\n- **Typos**: There are typographical errors in the paper, such as inconsistent matrix multiplication dimensions between lines 195-197 and missing add sign  between Lines 183-184 and\n- **Algorithm Description**: The description of Algorithm 1 should not be placed in the appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1:  Assumption of Gaussian distribution.\n\nA1: Thanks for your valuable suggestion. We have added Figure 1 to illustrate the validity of the Gaussian distribution assumption. Please refer to the submitted pdf file. We collect the output of two layers of the ResNet-18 network on CIFAR-10. Figure 1 (a) and (b) show the distributions of sample representation vectors' values in some dimension. Since we use the ReLU activation function, the obtained distributions are in accord with the Gaussian distribution in the positive quadrant. Figure 1 (c) and (d) show the distributions of  values of sample representation vectors' Euclidean norm, from which we can see that the two distributions can also be approximated as Gaussian distributions. We will also consider this problem further and add relevant explanations in the revised version.\n\nQ2: Convergence analysis.\n\nA2: Thank you for your insightful feedback. The convergence analysis presented in our manuscript is predominantly grounded in the framework utilized for the convergence analysis of natural gradient descent in related works. As you correctly pointed out, this result does have limitations, as different methods may exhibit similar rates when applied to a large number of hidden units. We have recognized this issue and have endeavored to provide a convergence analysis in terms of convergence rate and iteration steps within the framework of stochastic optimization with rigorous and appropriate assumptions. This work is currently ongoing and we intend to present such results in the following work in order to comprehensively illustrate the advantages of natural gradient descent in comparison with other methods.\n\nQ3: Complexity of adaptive learning rate.\n\nA3: We apologize for any ambiguity in the previous description. The adaptive learning rate is presented in Eq. (12). Furthermore, to enhance computational efficiency, we transform matrix computations into vector form as delineated in Eq. (13).  Notably, both ${\\bf d}^k_l$ and ${\\bf F}_l^k$  are common components in Natural Gradient Descent (NGD) approaches, ensuring that our method does not introduce additional variables but instead incorporates vector computations. In Theorem 1, we have demonstrated that the adaptive layer-wise learning rate contributes to a more rapid convergence of learnable parameters. Additionally, we conducted an Ablation Analysis in Section F.3 of the appendix, wherein confirm that the adaptive layer-wise learning rate effectively accelerates the training procedures. \n\n\nIn addition, LNGD serves as a second-order optimization algorithm. In practical applications, it operates similarly to other second-order optimizers (e.g., KFAC, EKFAC and  TKFAC) and first-order optimization algorithms (e.g., Adam), which target to efficiently modify the gradient descending magnitude of training parameters and often necessitate the integration of various learning rate adjustment schedules, including the linear warmup and cosine decay strategies as you mentioned. In our manuscript, we reference other second-order optimization algorithms and incorporate an exponential updating strategy for the learning rate.\n\nQ4: Computational overhead. \n\nA4:   The aim of Eq. (13) is to expedite the computation of the adaptive layer-wise learning rate provided in Eq. (12) by converting the matrix computations to vector form. It is noteworthy that the terms ${\\bf d}_l^k$ and $\\mathcal{D}\\theta_l$ in Eq. (13) are common variables in the NGD method, therefore we have not introduced additional parameters, but have indeed incorporated additional vector multiplication calculations to yield the adaptive layer-wise learning rate. Nonetheless, through Theorem.1 and Ablation Analysis in section F.3 in the appendix of our manuscript, the additional calculations introduced for computing the adaptive layer-wise learning rate are acceptable and can significantly speed up the convergence of training procedures.\n\n\nRegarding the validity of the final equation in Eq. (13), indeed, the strict equality does not hold. We are very grateful for your careful reading and correction. In reality, an approximate equality should be used here, where we utilize the batch average of the squared term to estimate the mean of the squared terms in computation.  We will amend this in the revised version. Thank you once again.\n\nQ5: Comparison with latest methods.\n\nA5: We acknowledge and appreciate the valuable suggestion provided. Our objective is to achieve a satisfactory balance between model training speed and accuracy. Building upon the existing NDG approaches, we propose a layer-wise sample method to efficiently compute each block matrix corresponding to each layer and introduce an adaptive layer-wise learning rate to expedite training. In Appendix B of our manuscript, we present comparisons of LNGD and KFAC along with its recent variants (EKFAC and TKFAC). To further validate the effectiveness of LNGD, we conducted additional comparative experiments, which affirm LNGD's superior performance at achieving the same level of accuracy. The detailed experimental results can be found in the response to all reviewers.  As our emphasis is not on whether LNGD can ultimately achieve the highest accuracy, but rather on the performance of various optimizations in achieving acceptable accuracy, we may have overlooked metrics such as peak memory usage. If time permits, we will evaluate the final accuracy level achievable by LNGD in the final version. Once again, we express our gratitude for the important and valuable suggestions provided for the further improvement of LNGD.\n\nQ6: Minor comments.\n\nA6: Thank you for your insightful suggestions. We will diligently proofread the manuscript and address the notations and typographical errors. Additionally, we will incorporate a comprehensive overview of recent pertinent literature published in 2023 and 2024 in the revised version. The details of Algorithm 1 will also be presented in subsection 3.3.", "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": "", "extra_scores": {}}, {"review_id": "l2SvRtka9p", "reviewer": "Reviewer_k9Mw", "summary": "The authors propose a computationally feasible second-order method for training neural nets, layer-wise natural GD, which includes an Adaptive Layer-Wise Learning Rate scheme. The method eliminates the backprop pass by using a layer-local sampling approach to approximate the Fisher information matrix; thus providing a novel approach to *local learning*.  The authors provide a coherent and reasonably thourough theoretical analysis and motivation for their method, which is backed up by a *somewhat modest amount pf experimental results*.", "questions": "* Have you considered how your method might affect the well-known implicit regularization of GD training? Specifically, the implicit bias of depth? It seems to me that skipping the backward pass might significantly affect this.\n* Please explicitly define the vector, **v**, of Thm 2 (near line 255). It is only defined in the appendix.\n* In Algo 1 you use both T_FIM and T_Fisherinformationmatrix. Please be consistent.", "limitations": "* What is your justifcation for calling your three experiments ``extensive''? Am I being too grumpy here? ;-)", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper is generally well-written with only few grammatical errors/typos. The authors did a good job of presenting the math and theoretical foundation for the method (including proofs in the appendix). Clearly, a good amount of work was done. While the authors frame their work as a computationally feasible second order method, I find that it is indeed an interesting addition to the lit on local learning.", "weaknesses": "While the authors claim (in the abstract) to provide ``extensive experiments'', I find the work lacking wrt. empirical validation. Only three results are provided, and those are for old architectures (far from SOTA), and without error bars. The proposed method could easily have been tested on a much wider variety of (even smaller) models and benchmarks --- providing the reader with a much more complete picture of how the method behaves in priactice. This leaves a good amount of doubt wrt. the robustness of the approach.\n\nGiven the known issues with other local learning methods, I suspect that this method too might have a strong tendency to overfit. Three experiments are not sufficient to convince me otherwise.  \n\n## Speaking of local learning, the authors did not touch on the topic at all. Thus, several relevant references were missed, such as:\n* Nøkland, A. (2016). Direct Feedback Alignment Provides Learning in Deep Neural Networks. ArXiv Preprint ArXiv:1609.01596. http://arxiv.org/abs/1609.01596\n* Nøkland, A., & Eidnes, L. H. (2019). Training Neural Networks with Local Error Signals. Proceedings of the 36th International Conference on Machine Learning, 4839–4850. https://proceedings.mlr.press/v97/nokland19a.html\n* Belilovsky, E., Eickenberg, M., & Oyallon, E. (2019). Greedy Layerwise Learning Can Scale To ImageNet. Proceedings of the 36th International Conference on Machine Learning, 583–593. https://proceedings.mlr.press/v97/belilovsky19a.html\n* Belilovsky, E., Eickenberg, M., & Oyallon, E. (2020). Decoupled Greedy Learning of CNNs. Proceedings of the 37th International Conference on Machine Learning, 736–745. https://proceedings.mlr.press/v119/belilovsky20a.html\n* Ren, M., Kornblith, S., Liao, R., & Hinton, G. (2022). Scaling Forward Gradient With Local Los\n* Xiong, Y., Ren, M., & Urtasun, R. (2020). LoCo: Local Contrastive Representation Learning. Advances in Neural Information Processing Systems, 33, 11142–11153.ses. https://arxiv.org/abs/2210.03310v3\n* Wang, Y., Ni, Z., Song, S., Yang, L., & Huang, G. (2021). Revisiting Locally Supervised Learning: an Alternative to End-to-end Training. ICLR 2021 - 9th International Conference on Learning Representations. https://arxiv.org/abs/2101.10832v1\n\nMoreover, I really think that the authors should have included their code in the submission. Their method substantially differs from the standard SGD variants, and the code would help a lot wrt. reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1: Weaknesses.\n\nA1: We would like to express our gratitude for the insightful suggestions provided. Our primary aim is to establish an optimal equilibrium between model training speed and accuracy. Building upon the existing approaches in NGD, we propose a layer-wise sampling methodology to efficiently compute the block matrix corresponding to each layer, as well as an adaptive layer-wise learning rate to enhance training efficiency. In Appendix B of our manuscript, we offer comparisons of LNGD and KFAC alongside its recent variants (EKFAC and TKFAC). To further validate the effectiveness of LNGD, we conducted additional comparative experiments on CIFAR-10 and ImageNet using recent NDG methods, consistently confirming LNGD's superior performance in achieving the same level of accuracy. Detailed experimental results are available in the response to all reviewers. Throughout the experimental process, extensive and multiple rounds of hyperparameter search were conducted, and the best results were chosen for comparison, therefore no error bars were included. As for the code, we are currently undergoing the approval process for the open-source workflow in compliance with the company's regulations, and it will be open-sourced shortly.\n\nQ2:  Have you considered how your method might affect the well-known implicit regularization of GD training? Specifically, the implicit bias of depth? It seems to me that skipping the backward pass might significantly affect this.\n\nA2: Thanks for your question. In LNGD, the key in parameter updating strategy is to calculate the updating direction ${\\bf d}^k_l=({\\bf F}^k_l)^{-1}\\nabla_{\\theta}h(\\theta^k_l)$. To save computational costs, we propose the layer-wise sample approximation approach to compute ${\\bf F}^k_l$ without having to perform a complete back-propagation. However, the calculation of the first-order gradient $\\nabla_{\\theta}h(\\theta^k_l)$ still requires performing a complete back-propagation. This may not have a significant impact on the implicit bias of depth. Since we have not focused on the implicit regularization in this paper, we may not provide a complete and detailed explanation. These contents will be further considered in our following work. Thank you again for providing a valuable direction for our further study.\n\nQ3:  Please explicitly define the vector, v, of Thm 2 (near line 255). It is only defined in the appendix.  \nIn Algo 1 you use both T_FIM and T_Fisherinformationmatrix. Please be consistent.\n\nA3: Thanks for your questions. We will add the definition of $\\bf{v}$ before Theorem 2 and correct $T_{Fisherinformationmatrix}$ as $T_{FIM}$ in the revised version.", "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": "", "extra_scores": {}}, {"review_id": "5u3LrzRIh8", "reviewer": "Reviewer_gNTY", "summary": "The authors propose a novel way to approximate the Fisher information matrix. They do this by starting with the block diagonal approximation of the Fisher information matrix which they compute using a novel layer-wise sampling method without performing a complete backpropagation. Then they approximate each block as a Kronecker product of two smaller matrices, where one is diagonal. They also keep the traces equal before and after the approximation. They call this new method the Layer-wise Natural Gradient Descent (LNGD). They further propose a new adaptive layer-wise learning rate that accelerates training. The authors also provide global convergence analysis of LNGD. At the end, they show extensive experiments on image classification and machine translation tasks, demonstrating their method is competitive.", "questions": "- I like Figure 1. But it would be interesting to also compare this with KFAC.\n\n- In Section 3.2, doesn't having adaptive learning rates contradict the spirit of the natural gradient?\n    - In the equation between line 172-173, isn't this computing the Taylor expansion, and thus also the Hessian? But I thought with the natural gradient and tricks to compute it, there should be no need to compute the Hessian?\n\n- Can you also apply your adaptive learning rate strategy to KFAC?", "limitations": "The authors provide limitations for each theory.", "rating": 4, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The idea of factoring the Fisher information matrix as in Section 3.1 seems novel and interesting.\n\n- The authors provide a convergence result in Theorem 2.\n\n- The authors have a great experimental section on complicated datasets. Their method outperforms all compared methods in both epoch and wallclock time.", "weaknesses": "- The biggest and most glaring weakness is that the paper doesn't compare with recent algorithms, some of which they even list in their own related works sections. The most recent algorithm that was compared is KFAC, which was published in 2015 -- almost a decade ago. If the authors provide comparisons to more recent algorithms (EKFAC, TKFAC, relevant extensions to convolutional neural networks, and NLP-related architectures, etc.), that might push it over to acceptance.\n\n- There is the main idea (Section 3.1) of factoring the Fisher information matrix which seems novel. But then there are various techniques (Section 3.2 and 3.3) that are employed to make the method more effective. But it's hard to disentangle whether the main contribution to the numerical performance is due to the novel factoring that allows using the natural gradient, or the techniques such as adaptive learning rates based on the Hessian. As asked in the Questions section of this review, can you apply the adaptive learning rate techniques to KFAC? How does it compare?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1:  The first part of weaknesses.\n\nA1: We are sincerely thankful for the valuable suggestions. Our primary objective is to propose and optimizer and achieve an optimal balance between the speed and accuracy of model training. Expanding upon established methodologies in the field of NGD, we propose a layer-wise sampling approach for efficiently computing the block matrix corresponding to each layer, as well as an adaptive layer-wise learning rate to improve training efficiency. In Appendix B of our manuscript, we provide comparisons of LNGD and KFAC as well as its recent variants (EKFAC and TKFAC). Additionally, in order to further substantiate the effectiveness of LNGD, we conducted additional experiments on CIFAR-10 and ImageNet using recent NGD methods, consistently demonstrating LNGD's superior performance in achieving comparable accuracy levels. Detailed experimental findings are available in our response to all reviewers, and we intend to include these pertinent results in the final version of our paper.\n\nQ2: The second part of weaknesses. \n\nA2: We appreciate the insightful suggestion. In order to provide a deeper understanding of the individual contributions of different components within the LNGD, a set of ablation studies have been conducted and are detailed in Appendix F.3 of our manuscript. These experiments are aimed at isolating the impact of adaptive learning rate and sampling optimization on the performance of the LNGD. Due to limitations in the rebuttal space, we will not elaborate further on this matter here and appreciate your understanding.\n\nQ3: I like Figure 1. But it would be interesting to also compare this with KFAC.\n\nA3: We have added the visualization results of KFAC in Figure 2, in which we only show partially enlarged figures due to page limitation.  From Figure (b) and (e), we can see that KFAC and LNGD can all emphasize the importance of the diagonal elements in the exact Fisher information matrix (FIM). In addition, it can also be seen clearly that KFAC still retains some elements near the main diagonal, while LNGD does not, which also reflects that LNGD provides an efficient approximation of FIM with less computational cost in comparison with KFAC.\n\nQ4: In Section 3.2, doesn't having adaptive learning rates contradict the spirit of the natural gradient?\n\nA4: Apologies for the lack of clarity in our presentation. We extend the current natural gradient methods by introducing a layer-wise sampling technique and an adaptive layer-wise learning rate to speed up the training process. In Section 3.2, we first discuss the necessity of an adaptive layer-wise learning rate, and then derive the formula for calculating it using a second-order Taylor expansion approach (i.e., Eq. (11)). Adhering to the NG principle, based on Eq. (11), we necessitate the computation of the Hessian matrix. However, across various loss functions such as cross-entropy or mean squared loss, the Fisher information matrix and the Hessian matrix are, in fact, equivalent. Hence, we proceed to compute the Fisher information matrix and utilize Eq. (12) to obtain the adaptive learning rate in LNGD.\n\nQ5: Can you also apply your adaptive learning rate strategy to KFAC?\n\nA5: Thanks for the valuable question. Our proposed adaptive learning rate strategy can also be applied to KFAC. As given in Appendix B, the main difference between KFAC and LNGD is the approximations of the block matrix $F_l$. KFAC approximates  $F_l$ as the Kronecker product of $A=E[a_{l-1}a_{l-1}^\\top]$ and $ B=E[g_lg_l^\\top]$. Our proposed LNGD \napproximates  $F_l$ as the Kronecker product of a matrix ${\\bf\\Phi}_l$ and a diagonal matrix ${\\bf\\Psi}_l$, which is computed by sampling from each layer, details are given in Section 3. The adaptive layer-wise learning rate is computed by Eq. (12), which would be the same for KFAC and LNGD. Hence, it is feasible to integrate the adaptive layer-wise learning rate into KFAC with minimal costs.", "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": "", "extra_scores": {}}, {"review_id": "8dyzXbyMPe", "reviewer": "Reviewer_Mvn4", "summary": "The proposed approach uses the block diagonal form of the Fischer Information Matrix (FIM) which is achieved by computing this FIM with every layer in the network sequentially. In addition, the space complexity is limited by consider the diagonal form of this FIM matrix for each layer and its respective gradient outer products.", "questions": "1. What is $\\mathbf{s}_l$ in line 136 ?\n2. This reveiwer would like to know what's the difference between the proposed approach and https://ojs.aaai.org/index.php/AAAI/article/view/16867 ?", "limitations": "This is not really a limitation - Given that this paper improves upon fundamental results significantly, an open source implementation of this approach would a be really good to have.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors find a clever way of approximating the Fischer information matrix by only consider the block diagonal form of the FIM. In addition, due to i.i.d. assumptions, the block diagonal reduces to a diagonal matrix for \n$\\Psi_l \\otimes \\Phi_l$\n2. The adaptive learning rate for each layer further makes it an interesting way of optimizing the weights.", "weaknesses": "The reviewer has actually found it difficult to find any weaknesses with the paper. It is really well written, with sufficient proofs and sufficient results to support their claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1: What is $s_l$ in line 136 ?\n\nA1: Sorry for the unclear description, $s_l= W_l a_{l-1}$. We will correct it in the final version.\n\nQ2: This reveiwer would like to know what's the difference between the proposed approach and https://ojs.aaai.org/index.php/AAAI/article/view/16867? \n\nA2: Thanks for your question. The THOR, which is proposed in https://ojs.aaai.org/index.php/AAAI/article/view/16867, mainly considers reducing the computational cost of natural gradient descent through two aspects. On the one hand, THOR gradually increasess the updating frequency of the inverse matrix of Fisher information matrix (FIM) and proposes a trace-based updating rule for FIM for each layer. On the other hand, THOR approximates the approximated FIM obtained by KFAC as some smaller matrices by splitting matrix dimensions. Our proposed LNGD first gives a layer-wise sample method to more efficiently compute each block matrix corresponding to each layer and proposes a novel approximate scheme of the FIM. Furthermore, LNGD also adopts an adaptive layer-wise learning rate to speed up training. The contributions and ideas of our proposed LNGD are different from THOR.", "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": "", "extra_scores": {}}]

Author rebuttal:
We are very grateful to the four reviewers for their constructive comments and valuable suggestions on our manuscript.

The tables and figures mentioned in the reply are given in the the submitted pdf file. Please see it for details. 

In our manuscript, we aim to propose an optimizer that can achieve a good balance between the training speed and accuracy of the model. To this end, based on the existing NDG approaches, we introduce LNGD. In Appendix B of our manuscript, we provide comparisons of the primary variances between LNGD, KFAC, and its recent variants (EKFAC and TKFAC). In order to further validate the effectiveness of LNGD, we conducted additional experiments on the CIFAR-10 dataset, in which three methods including EKFAC [1], TKFAC [2], and NG+ [3] are added for comparison. The detailed statistics are presented in Table 1. From the table, we observe that LNGD achieves a testing accuracy of 91\% with the fewest epochs and the shortest total time. Furthermore, LNGD exhibits the smallest computational time per epoch. Additionally, due to the efficient FIM approximation strategy adopted by NG+, it can significantly reduce the computational time and number of epochs compared to EKFAC and TKFAC. Thus, considering the time constraints, we prioritize testing the performance of NG+ on ImageNet, and the results are presented in Table 2. It is evident that LNGD consistently demonstrates a noteworthy reduction in computational time compared to NG+.

References:

[1] Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a Kronecker factored eigenbasis. In Advances in Neural Information Processing Systems, pages 9550–9560, 2018.

[2] Kaixin Gao, Xiaolei Liu, Zhenghai Huang, Min Wang, Zidong Wang, Dachuan Xu, and Fan Yu. A trace-restricted Kronecker-factored approximation to natural gradient. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 7519–7527, 2021.

[3] Minghan Yang, Dong Xu, Qiwen Cui, Zaiwen Wen, and Pengxiang Xu. An efficient Fisher matrix approximation method for large-scale neural network optimization. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(5):5391–5403, 2023.

Meta-review: The paper proposes a new second-order method for training neural networks using layer-wise natural gradient descent. The key design in their method is an Adaptive Layer-Wise Learning Rate scheme, where a layer-local sampling approach is used to approximate the Fisher information matrix to improve local learning. This approximation method is novel and efficient. The proposed methods have been proved effective in both theory and numerical experiments on challenging datasets.
Decision: Accept (poster)