A Layer-Wise Natural Gradient Optimizer for Training Deep Neural Networks

Authors: xiaolei  liu
shaoshuai  li
kaixin  gao
binfeng  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 04:14:07

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.

Main paper: 
 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 a same learning rate, such as stochastic gradient descent (SGD) [1] and its accelerations [2, 3] , and methods use adaptive learning rate, such as AdaDelta [4] , RMSProp [5] , 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 of these methods. 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 secondorder 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. [9] approximated the Hessian matrix as a diagonal operator, which is achieved by applying Hutchinson's method, and proposed the AdaHessian method. Goldfarb, Ren and Bahamou [10] developed Kroneckerfactored block-diagonal BFGS and its limited-memory variants L-BFGS methods for training DNNs. Generalized Gauss-Newton methods, such as the Hessian-free method [11] 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 [18] . 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 and Grosse [14] proposed the Kronecker-factored 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 [21] and showed significant speedup during training. In addition, George et al. [22] proposed the eigenvalue-corrected Kronecker factorization (EKFAC) method. Gao et al. [15, 23] proposed the trace-restricted Kronecker-factored approximate (TKFAC) method. These approaches all focus on the Kronecker-factored approximations of the Fisher information matrix. What's more, some works have also considered large-scale distributed computing using NGD for training DNNs and shows excellent experimental performance [16, 24, 25] .In this paper, our main focus is on 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 a 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 a layer-wise sample method to 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 a 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 it 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 is organized as follows. Section 2 gives the notations and introduces the NGD method. In Section 3, we propose a novel approximation of the Fisher information matrix and the adaptive layer-wise learning rate. Furthermore, we give the framework of LNGD and establish the convergence analysis. Section 4 presents the results of experiments on image classification and machine translation tasks. The conclusion is drawn in Section 5.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 or 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
notations and preliminaries: 
 In this paper, for a matrix A, we use A ij to denote its (i, j)th entry, tr(A) to denote its trace and ∥A∥ F to denote its Frobenius norm. We use • and ⊗ to denote the Hadamard and Kronecker product of two matrices. In the following, we 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 h(θ). If the loss function h(θ) is chosen as the the cross-entropy loss function, h(θ) can be given as h(θ) = E[log p(y|x, θ)], where p(y|x, θ) is the density function of a predictive distribution P y|x (θ), and x, 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 asF -1 ∇ θ h(θ),where F is the Fisher information matrix given byEQUATIONIn Eq. ( 1), the input x is independently sampled from a distribution Q x with density function being q(x) and the label y is sampled from the predictive distribution P y|x (θ). In the following pages, we abbreviate E x∼q(x),y∼p(y|x,θ) as E unless otherwise specified. Consider a neural network with L layers, for each layer l ∈ [L] with [L] = {1, 2, . . . , L}, we denote a l-1 and W l as the input (the activation from the previous layer) and the matrix of weights of this layer, respectively. What's more,θ l = vec(W l ) and θ = (θ 1 , . . . , θ L ) ⊤ = (vec(W 1 ) ⊤ , . . . , vec(W L ) ⊤ ) ⊤ , where vec(•)indicates vectorization of a matrix. For convenience, we denote the derivative of the loss function with respect to θ as Dθ = -∇ θ log p(y|x, θ). Then the Fisher information matrix can be expressed asF = E[DθDθ ⊤ ].Due to the high computational and storage costs caused by the inverse operation of high-dimensional matrices, it is impractical to directly compute F -1 in the training of DNNs. The family of Kroneckerfactored approximations provides an effective approach for computing F -1 of parameters in highdimensional 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.,EQUATIONwhereF l = E[Dθ l Dθ ⊤ l ] for any l ∈ [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 a 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.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, a thorough comparison with related work is furnished in the appendix.Guidelines:• The answer NA means that the paper has no 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 a 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
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.For NGD methods to train DNNs, the Fisher information matrix can be approximated by a block diagonal one according to different layers as given by Eq. ( 2), this approximation can be found in [14, 15, 19, 22] and references therein. We call such a block diagonal approximate Fisher information matrix the layer Fisher information matrix, which is computed based on a distribution Q x and a predictive distribution P y|x (θ) as given in Eq. (1) . To obtain the layer Fisher information matrix, we need perform a complete back-propagation to sequentially compute F L , F L-1 , . . . , F 1 , 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 l is computed based on the lth layer's prediction distribution P a l |a l-1 (θ l ) with the input a l-1 ∈ R d l of this layer and the input a l ∈ R d l+1 of the (l + 1)th layer instead of using the same predictive distribution P y|x (θ) for all layers. Specifically, for a l , we assume that the predictive distribution P a l |a l-1 (θ l ) follows Gaussian distribution, which is usual used as prior by variational auto-encoder [26] , so F l can be computed by sampling from a normal distribution with expectation being a l and variance being I. Similar assumption can also be found in [27] , in which the normality is also supported by a central limit theorem under the independence assumption. By this layer-wise sample approximation, we can compute the layer Fisher information matrix without having to perform a complete back-propagation and thus improve the computational efficiency.Next, we can give the formula of each block F l in the layer Fisher information matrix asF l = E[ Dθ l Dθ ⊤ l ] = E[vec( DW l )vec( DW l ) ⊤ ] = E[vec(g l a ⊤ l-1 )vec(g l a ⊤ l-1 ) ⊤ ] = E[(a l-1 a ⊤ l-1 ) ⊗ (g l g ⊤ l )] ∈ R m l d l ×m l d l , where Dθ l = -∇ θ log p(a l |a l-1 , θ l ), g l = -∇ s l log p(a l |a l-1 , θ l ) ∈ R m l with s l = W l a l-1, and p(a l |a l-1 , θ l ) is the density function of the distribution P a l |a l-1 (θ l ). In practice, the dimension of each block matrix F l 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 l follows Gaussian distribution with expectation being a l and variance being I, and each element of activation output a l is independent and identically distributed random number, then each element of partial derivative g l is also independent and identically distributed. It is easy to show that F l can be seen as a matrix with d l × d l block matrices, in which each block is an m l × m l matrix and the off-diagonal elements are zero. Therefore, F l can be approximated asEQUATIONCombining the property that (A ⊗ B) -1 = A -1 ⊗ B -1 for any two invertible matrices A and B, we can significantly reduces the computational complexity. Thus, some approaches have considered 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 l as the Kronecker product of two factor matrices Φ l ∈ R d l ×d l and Ψ l ∈ R d l ×d l . To get factor matrices Φ l and Ψ l , we first replace diag(g l g ⊤ l ) in Eq. ( 3) by its trace and obtain Φ l . Then we compute Ψ l while keeping that tr(F l ) = tr(Φ l ⊗ Ψ l ). Specifically, Φ l is given byEQUATIONOn the other hand, Ψ l can be computed byΨ l = E[(a ⊤ l-1 a l-1 ) × diag(g l g ⊤ l )] E[(a ⊤ l-1 a l-1 )(g ⊤ l g l )].Based on Eq. ( 4) and Eq. ( 5), we can show that tr(F l ) = tr(Φ l ⊗ Ψ l ).Fig. 1 presents the visualization results of the exact Fisher information matrix F, our approximation F LNGD , 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 F LNGD 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, F LNGD can also provide an effective approximation of the elements in 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 F LNGD is efficient and F LNGD can retain most of information.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 k is the update direction of the function h : R n → R at the iteration point θ k . We first recall the gradient descent methods for getting the minimization of h, in which the update rule can be given as θ k+1 = θ kα k d k , where α k is the learning rate, which can be chosen according to the value of the quadratic modelh(θ k -α k d k ) ≈ h(θ k ) -α k ⟨d k , ∇ θ h(θ k )⟩ + (α k ) 2 2 (d k ) ⊤ ∇ 2 θ h(θ k )d k .Once the update direction is chosen, the minimizer of α k can be given byα k = ⟨d k , ∇ θ h(θ k )⟩ (d k ) ⊤ ∇ 2 θ h(θ k )d k (6) if (d k ) ⊤ ∇ 2 θ h(θ k )d k is nonzero. If ∇ 2 θ h(θ k ) is positive definite and d k = (∇ 2 θ h(θ k )) -1 ∇ θ h(θ k ), then α k = 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 [5, 6, 28] . In the following, we consider giving an adaptive element-wise learning rate automatically scaled by the direction d k . In this case, the update rule of parameters is given by θk+1 = θ k -α k •d k = θ k -D k α k , where α k ∈ R n is the learning rate, D k ∈ R n×n is a diagonal matrix with (D k ) ii = (d k ) i and (D k ) ij = 0 when i ̸ = j for i, j ∈ [n]and "•" denotes the element-wise product. The second Taylor expansion of h(θ -Dα) at iteration k ish(θ k -D k α k ) ≈h(θ k ) -⟨D k α k , ∇ θ h(θ k )⟩ + 1 2 (D k α k ) ⊤ ∇ 2 θ h(θ k )D k α k .Taking the derivative of h with respect to α k and letting it equal to 0, we getEQUATIONif D k and ∇ 2 θ h(θ k ) are positive definite. Note that in Eq. (7) , it is impractical to compute the inverse of ∇ 2 θ h(θ k )D k directly for largescale 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 a block diagonal one according to different layers [10, 14, 15, 19, 24] . What's more, some works have observed that parameters in the same layer have gradients of similar magnitudes. Therefore, a common learning rate can be efficiently shared by these parameters [29, 30] . Inspired by these works, we propose a novel adaptive layer-wise learning method as follows. Suppose thatEQUATIONis a block diagonal matrix and αk ∈ R L is the learning rate. The approximate second Taylor expansion of h(θEQUATIONwhereH k = diag(H k 1 , H k 2 , . . . , H k L ) and H k l = ∇ 2 θ l h(θ k ) for l ∈ [L]and the Hessian matrix is approximated by the block diagonal matrix H k . Taking the derivative of h with respect to αk and letting it equal to 0, we get ( Dk ) ⊤ ∇ θ h(θ k ) = ( Dk ) ⊤ H ⊤ Dk αk , which yields thatα k = diag(Θ 1 , Θ 2 , . . . , Θ L ) -1 ((d k 1 ) ⊤ ∇ θ1 h(θ k ), (d k 2 ) ⊤ ∇ θ2 h(θ k ), . . . (d k L ) ⊤ ∇ θ L h(θ k )) ⊤ (10) if Θ l is nonzero, where Θ l = (d k 1 ) ⊤ H k 1 (d k 1 ) for l ∈ [L]. If a same learning rate is used for all layers, as the same way of computing the adaptive layer-wise learning rate, we can getEQUATION)Theorem 1. Let g(θ) and g L (θ) be the approximate second Taylor expansions of h(θαd) and h(θ -D α) as given in (9) , where D ∈ R n×L , α ∈ R L and α ∈ R are given by ( 8), (10) and (11) respectively, then we have g L (θ) ≤ g(θ).Proof. The proof is given in the appendix.By Theorem 1, we know that the adaptive layer-wise learning rate may lead to a faster decline in terms of function values. In our proposed algorithm, we choosed k = (F k ) -1 ∇ θ h(θ k ), where F k is the Fisher information matrix and can be seen as a approximation of the Hessian matrix. Then, the Fisher information matrix is approximated by a block diagonal matrix each block matrix is approximated by the Kronecker product of two factor matrices. In each layer, the update direction d k l is scaled by a layer-wise damping learning rate α k l according to (10) , which is given byEQUATIONwhere µ > 0 is a 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-wise learning rate (13) and remain computationally efficient.To effectively apply LNGD in training DNNs, several certain techniques need to be employed. In this section, we primarily focus on introducing the damping technique, which is a 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 (d k l ) ⊤ F k l d k l is relatively expensive. Finally, we discuss the utilization of exponential moving averages to enhance the training process.A 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 computational difficulty and inaccuracy since most eigenvalues of the Fisher information matrix are close to zero and only a small number of eigenvalues take on large values. To make training stable, we propose the following damping for the lth layer: λ l = min(max(tr(F l )/d l , ν 1 ), ν 2 ), where ν 1 and ν 2 are two constants to constrain the minimum and maximum of damping, and d l is the number of weight parameters. In our method, F l is approximated as the Kronecker product of two factors Φ l and Ψ l , so we add the damping to each factors by Φl = Φ l + λ Φ l and Ψl = Ψ l + λ Ψ l , where λ Φ l = min(max(tr(Φ l )/n, ν 1 ), ν 2 ) and λ Ψ l = min(max(tr(Ψ l )/n, ν 1 ), ν 2 ). Compute the learning rate the faster: In order to compute the adaptive layer-wise learning rate given in Eq. ( 12) more quickly, we turn matrix computation into vector computation. Specifically,EQUATIONwhere F k l is the empirical Fisher information matrix and N is the number of samples. The empirical version of Fisher information matrix with no need for sampling from the model's prediction distribution, making it more computationally efficient.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 ϵ, with the estimate calculated from the new mini-batch, weighted by 1ϵ. That isEQUATIONIn summary, our proposed algorithm is shown in Algorithm1.In this subsection, we give the convergence analysis of LNGD. Following the model used in previous works about analysing the gradient descent [31, 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 (θ, a, x) = 1 √ m m r=1 a r φ(θ ⊤ r x),where θ 1 , θ 2 , . . . , θ m ∈ R d are the weight vectors of the first layer, x ∈ R d is the input, a r ∈ R is the weight of unit r in the second layer and φ(•) is the ReLU activation function, i.e., φ(x) = max{0, x}.Let v = [f (θ, a, x i ), f (θ, a, x 2 ), . . . , f (θ, a, x n )] ⊤ .In the following, we only give the result of convergence of Algorithm 1, the specific proof, which uses some conclusions in [36, 37, 38, 39] , is given in the appendix. Theorem 2. (Convergence rate of LNGD) Under the Assumption 1 and the assumption that rank(X) = d. If we set the number of hidden units m = Ωn 4 κ 8 Z X,G ν 2 ε 3 λ 4 G , we i.i.d initialize θ r ∼ N (0, νI), a r ∼ unif[{-1, +1}] for any r ∈ [m],and we set the step size α ≤ (1-2c) (1+c) 2 . Then with probability at least 1ε over the random initialization, we have for k = 0, 1, 2, . . .∥y -v k ∥ 2 2 ≤ (1 -α) k ∥y -v 0 ∥ 2 2 .Proof. The proof is given in the appendix.Require: learning rate α, learning rate parameter µ, damping parameter λ, damping constraints ν 1 ,ν 2 , momentum parameter τ , exponential moving average parameter ϵ, Fisher information matrix and its inverse update intervals T FIM and T INV . 1: k ← 0, m ← 0. Initialize Φl and Ψl for any l ∈ [L]. 2: while convergence is not reached do if k ≡ 0 (mod T FIM ) then 6:Update the factors Φl and Ψl using Eq. ( 14)7: end if 8: if k ≡ 0 (mod T INV ) then 9:Compute the inverses of Φl and Ψl Compute ∇ θ l h(θ) using backpropagation 12:Compute the approximated natural gradient ( Φ-1l ⊗ Ψ-1 l )∇ θ l h(θ) 13:Compute the adaptive learning rate α l using Eq. ( 12)14:ζ ← -αα l ( Φ-1 l ⊗ Ψ-1 l )∇ θ l h(θ) 15: m ← τ m + ζ (Update momentum) 16: θ l ← θ l + m (Update parameters) 17:end for 18:k ← k + 1 19: end while 20: return θ 
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 are tuned using a grid search with values α ∈ {1e -4, 3e -4, . . . , 1, 3}. The damping parameter λ in KFAC [14] are tuned using a grid search with values λ ∈ {1e -6, 1e -4, 3e -4, 1e -3, . . . , 1e -1, 3e -1}. The minimum and maximum of damping parameters ν 1 and ν 2 in LNGD are set to 1e -5 and 1e -2. The moving average parameter and the momentum correlating with KFAC and LNGD are set to 0.95 and 0.9, respectively. 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-parameter settings for these optimizers are the best values randomly searched for many times.We first report the optimizing performance on CIFAR-10 [40] , which is a standard task used to benchmark optimization methods [6, 41, 42, 43, 44] . Following these previous works, the changes of testing accuracy and training loss versus time as 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 in training loss during the initial epochs and seconds. This suggests that LNGD is effective in quickly reducing the training loss and reaching convergence. All optimization methods convergent at around 200 epochs. However, it is observed that second-order optimization methods, such as KFAC and LNGD, achieve a lower training loss compared to first-order optimization methods like SGD and Adam. In terms of testing accuracy, as depicted in Fig. 2 (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 Table6, SGD and ADAM require at 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 Table6, LNGD gets the highest final testing accuracy after convergence. neural network activated by the ReLU function with m neurons in the hidden layer as follows:f (θ, a, x) = 1 √ m m r=1 a r φ(θ ⊤ r x),where θ 1 , θ 2 , . . . , θ m ∈ R d are the weight vectors of the first layer, x ∈ R d is the input, a r ∈ R is the weight of unit r in the second layer and φ(•) is the ReLU activation function, i.e., φ(x) = max{0, x}.Let v = [f (θ, a, x i ), f (θ, a, x 2 ), . . . , f (θ, a, x n )] ⊤ .In the following, we only give the result of convergence of Algorithm 2, the specific proof is given in the appendix. Theorem 2. (Convergence rate of LNGD) Under the Assumption 1 and the assumption that rank(X) = d. If we set the number of hidden units m = Ωn 4 κ 8 Z X,G ν 2 ε 3 λ 4 G , we i.i.d initialize θ r ∼ N (0, νI), a r ∼ unif[{-1, +1}] for any r ∈ [m],and we set the step size α ≤ (1-2c) (1+c) 2 . Then with probability at least 1ε over the random initialization, we have for k = 0, 1, 2, . . .∥y -v k ∥ 2 2 ≤ (1 -α) k ∥y -v 0 ∥ 2 2 .Proof. The proof is given in the appendix.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 are tuned using a grid search with values α ∈ {1e -4, 3e -4, . . . , 1, 3}. The damping parameter λ in KFAC [14] are tuned using a grid search with values λ ∈ {1e -6, 1e -4, 3e -4, 1e -3, . . . , 1e -1, 3e -1}. The minimum and maximum of damping parameters ν 1 and ν 2 in LNGD are set to 1e -5 and 1e -2. The moving average parameter and the momentum correlating with KFAC and LNGD are set to 0.95 and 0.9, respectively. 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-parameter settings for these optimizers are the best values randomly searched for many times.We first report the optimizing performance on CIFAR-10, which is a standard task used to benchmark optimization methods [6, 36, 37, 38] . Following these previous works, the changes of testing accuracy and training loss versus time as well as epoch are reported in Fig. 2 , and detailed statistics are shown in Table .1. From Fig. 2 , it can be observed that LNGD exhibits the most rapid decline in training loss during the initial epochs and seconds. This suggests that LNGD is effective in quickly reducing the training loss and reaching convergence. All optimization methods convergent at around 200 epochs. However, it is observed that second-order optimization methods, such as KFAC and LNGD, achieve a lower training loss compared to first-order optimization methods like SGD and Adam. In terms of testing accuracy, as depicted in Fig. 2 We extend our examination of optimizer efficacy to a larger image classification dataset, ImageNet-1K [45] . The changes of testing accuracy and training loss versus time and epoch are reported in Fig. 3 and Table2. The results show that the LNGD optimizer is 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 a significant improvement over SGD, which takes 100% more epochs and 75% more time to reach the same accuracy level. As for Adam, it exhibits a rapid decrease in loss during training and reaches convergence at a 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, as it 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.In this experiment, we apply LNGD to the Transformer-Big model [46] with 213.7M parameters.The training datasets is WMT English-German machine translation corpus [46] . We use Bleu [47] as the evaluation metrics, 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, as evidenced by a more rapid decrease in training loss and a lower converged loss value. This observation aligns with previous empirical findings that ADAM is highly effective for transformer models. KFAC exhibits further enhancements in performance compared to Adam, 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, as it 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.In this experiment, we apply LNGD to the Transformer-Big model [40] with 213.7M parameters.The training datasets is WMT English-German machine translation corpus [40] . We use Bleu [41] as the evaluation metrics, 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. 5 and Table .3, 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, as evidenced by a more rapid decrease in training loss and a lower converged loss value. This observation aligns with previous empirical findings that ADAM is highly effective for transformer models. KFAC exhibits further enhancements in performance compared to Adam, 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 Bleu score, LNGD achieves a top-1 Bleu score of 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 a 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 are quite limited.  
conclusion: 
 In summary, we propose a novel NGD optimizer named as 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. 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 Bleu score, LNGD achieves a top-1 Bleu score of 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 a 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 are quite limited. In summary, we propose a novel NGD optimizer named as 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 to existing state-of-the-art optimizers. This work hold significant potential for enhancing the efficiency and scalability of training processes in deep learning frameworks.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: Due to the involvement of proprietary code resources, the disclosure of such materials must adhere to the company's 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: For each experiment, we average the results of 5 runs and the hyper-parameter 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• 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. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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.• 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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 appropriate citations.Guidelines:• The answer NA means that the paper does not use 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 is used and, if 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • 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 are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: We present an optimization method, which does not involve the release of new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model 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.  
 
 Appendix: In this paper, we denote [n] = {1, 2, . . . , n}. For a matrix A, we use A ij to denote its (i, j)-th entry, tr(A) to denote its trace, σ min (A) and σ max (A) to denote its smallest and largest singular value, ∥A∥ F to denote its Frobenius norm and ∥A∥ 2 to denote its spectral norm. If A is positive semi-definite, λ min (A) and λ max (A) denote its smallest and largest eigenvalue, and define κ A = λ max (A)/λ min (A), respectively. The identity matrix is denoted as I. For a vector a, ∥a∥ 2 denotes the Euclidean norm. We use N (µ, Σ) to denote the Gaussian distribution with mean µ and covariance Σ. For two matricesA = [a 1 , a 2 , . . . , a n ] ∈ R p×q , B = [b 1 , b 2 , . . . , b n ] ∈ R p×q ,we use • and ⊗ to denote the Hadamard and Kronecker product, respectively. The column-wise Khatri-Rao product * is defines asA * B = [a 1 ⊗ b 1 , a 2 ⊗ b 2 , . . . , a n ⊗ b n ] ∈ R p 2 ×q .Similarly, we can define the row-wise Khatri-Rao product ⋆ and we have(A ⋆ B) ⊤ = (A ⊤ * B ⊤ ).Given an event E, I{E} denotes its indicator function, i.e., I{E} = 1, if E happens, 0, otherwise.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 Kronecker-factored approximation to the Fisher information matrix to address this computational challenge. The most related approaches to this work are the KFAC, EKFAC and TKFAC [15] . These works and LNGD all start with a block-diagonal approximation of the Fisher information matrix. The differences among them are the approximations of the block matrix F l . By approximating the expectation of the Kronecker product as the Kronecker product of expectations, KFAC approximates F l as the Kronecker product of A = E[(a l-1 a ⊤ l-1 )] and B = E[(ĝ l ĝ⊤ l )] with ĝl = -∇ s l log p(y|x). By tracking the diagonal variance in the Kronecker-factored eigenbasis, EKFAC performs eigenvalue decomposition of the Fisher information matrix and re-scales the eigenvalues by S * to achieve a better approximation, where S * is a diagonal matrix defined byS * ii = E[(U ⊤ B ∇ θ h(θ) 2 ) i ],and U A , U B are eigenvectors of A, B. TKFAC approximates F l as a 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 l as a Kronecker product of a matrix Φ l and a diagonal matrix Ψ l , which is computed by sampling from each layer. We summarize the above approximations of F l in Table 4 .F l KFAC A ⊗ B EKFAC (U A ⊗ U B )S * (U A ⊗ U B ) ⊤ TKFAC δP ⊗ Q LNGD Φ ⊗ Ψ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 backpropagation. On the other hand, EKFAC modifies KFAC by incorporating eigenvalue decomposition to scale the eigenvalue during the inversion process. TKFAC, another variant of 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 LNGD also requires computation and inversion of two factors Φ and Ψ. However, the advantage of LNGD is that the matrix Ψ is diagonal, and due to the use of hierarchical sampling, Ψ can be directly inverse during the forward propagation phase, without the need to wait for the completion of the back-propagation 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 of the inverse matrix of Fisher information matrix and proposes a trace-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 an adaptive layer-wise learning rate to speed up training. The contributions and ideas of our proposed LNGD are different from THOR.Fig. 5 shows the visualization results of KFAC and LNGD. From Fig. 5 (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. 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 Fisher information matrix with less computational cost in comparison with KFAC.In Subsection 3.1, we assume that the predictive distribution P a l |a l-1 (θ l ) follows Gaussian distribution. 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 . Proof. Since g(θ) and g L (θ) are the approximate second Taylor expansions of h(θαd) and h(θ -D α), we haveg L (θ) -g(θ) = h(θ) -α⟨d, ∇ θ h(θ)⟩ + 1 2 α 2 d ⊤ Hd -h(θ) + L l=1 αl ⟨d l , ∇ θ l h(θ)⟩ - L l=1 1 2 α2 l d ⊤ l H l d l = 1 2 L l=1 (d ⊤ l ∇ θ l h(θ)) 2 d ⊤ l H l d l - (d ⊤ ∇ θ h(θ)) 2 d ⊤ Hd ≤ 0.This completes the proof.In this section, we give the convergence analysis of LNGD. Following the model used in previous works about analysing the gradient descent [31, 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 (θ, a, x) = 1 √ m m r=1 a r φ(θ ⊤ r x),where θ 1 , θ 2 , . . . , θ m ∈ R d are the weight vectors of the first layer, x ∈ R d is the input, a r ∈ R is the weight of unit r in the second layer and φ(•) is the ReLU activation function, i.e., φ(x) = max{0, x}.For convenience, we define θ= [θ ⊤ 1 , θ ⊤ 2 , . . . , θ ⊤ m ] ⊤ ∈ R md .We first initialize the parameters randomly by θ r ∼ N (0,ν 2 I), a r ∼ unif[{-1, +1}], ∀r ∈ [m], where 0 < ν ≤ 1 controls the magnitude of initialization.S = {(x i , y i )} n i=1 containing (input, target) examples (x i , y i ).Following [31, 33, 34] , we make the following assumption for the data. Assumption 1. For all i, ∥x i ∥ 2 2 = 1 and |y i | = O(1). For any i ̸ = j, x i ∦ x j . In this subsection, we mainly focus on the mean squared error loss (MSE) functionL(θ) = 1 2n n i=1 (f (θ, a, x i ) -y i ) 2 = 1 2n n i=1 ( 1 √ m m r=1 a r φ(θ ⊤ r x) -y i ) 2 .Following [31, 33, 34] , we fix the weights of second layer and only optimize the weights of first layer.Then the update rule of NGD can be written asθ k+1 = θ k -α(F k ) -1 ∇ θ L(θ k ).As shown in [18] , if the network's predictive distribution is in the exponential family, the Fisher information matrix is equivalent to the generalized Gauss-Newton matrix, which is defined byE (xi,yi)∈S [J ⊤ i H L J i ], where H L is the Hessian matrix of the loss function L(θ) with respect to the prediction f (θ, a, x i ) and J i is the Jacobian matrix of f (θ, a, x i ) with respect to the parameters θ. Under our setting that L(θ) is the MSE loss function, the Hessian matrix H L is the identity matrix I. J i can be computed byEQUATIONLet J = (J 1 , J 2 , . . . , J n ) ⊤ ∈ R n×md , then the Fisher information matrix can be written asF = E (xi,yi)∈S [J ⊤ i H L J i ] = E (xi,yi)∈S [J ⊤ i J i ] = 1 n J ⊤ J.As discussed in [34] , when m > n, the Fisher information matrix is singular. So in this case, we use the generalized inverse given in [36 ]F † = nJ ⊤ (JJ ⊤ ) -1 (JJ ⊤ ) -1 J (16)and the update rule of NGD can be written asθ k+1 = θ k - α n (F k ) † (J k ) ⊤ (v k -y), where y = [y 1 , y 2 , . . . , y n ] ⊤ and v = [v 1 , v 2 , . . . , v n ] ⊤ = [f (θ, a, x i ), f (θ, a, x 2 ), . . . , f (θ, a, x n )] ⊤ .Consider the two-layer neural network described in this subsection, since we fixed the weights in second layer, and the Fisher information matrix of this model is approximated by F ≈ Φ ⊗ Ψ. For simplicity, we ignore the index of layer. DefineX = [x 1 , x 2 , . . . , x n ] ⊤ ∈ R n×d and G = [φ ′ (Xθ 1 ), φ ′ (Xθ 2 ), . . . , φ ′ (Xθ m )] ∈ R n×m, where X is the input matrix formed by n input vectors and G is the pre-activation derivatives matrix whose entry is given byEQUATIONThen we haveJ = ∇ θ f (θ, a, x i ) = X ⋆ G ∈ R n×md and F = 1 n J ⊤ J = 1 n (X ⋆ G) ⊤ (X ⋆ G) ∈ R md×md .The formulas of Φ and Ψ can be given asEQUATIONEQUATION)Under the Assumption 1, we have(XX ⊤ ) ii = 1, i ∈ [d]. Therefore, Eq. ( 18) and Eq. ( 19) can be simplified toΦ = Z X,G , Ψ = G ⊤ G nη ,where η = tr((XX⊤ • I) • (GG ⊤ • I)) and Z X,G = X ⊤ (GG ⊤ • I)X.Finally, the update rule of LNGD can be given byθ k+1 = θ k -α[η(Z k X,G ) -1 ⊗ ((G k ) ⊤ G k ) -1 ](J k ) ⊤ (v k -y).To analyze the global convergence of LNGD, we need the following two conditions as given in [34] .Proof. Combine Eq. ( 15), Eq. ( 17) and Assumption 1, we have ∥G -G(0)∥ 2 F ≤ ∥J -J 0 ∥ 2 F . Let m = Ω n 4 κ 8 Z X,G ν 2 ε 3 λ 4 K and R = 2∥y-v 0 ∥2 √ λ K /2κ Z X,G , by Condition 2 and Lemma 3 we have∥G -G(0)∥ 2 2 ≤ ∥G -G(0)∥ 2 F ≤ ∥J -J 0 ∥ 2 F ≤ c 8 λ K κ 2 Z X,G ≤ c 8 λ K . Therefore, σ min (G) ≥ σ min (G(0)) -∥G -G(0)∥ 2 ≥ 3 4 λ K - c 8 λ K ≥ √ 2 2 λ 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λ min (K) = λ min (GG ⊤ ) ≥ λ K 2 .Proof complete.Lemma 5. [38] Let A and B be two positive define matrices, we haveλ max (A • B) ≤ max i A ii λ max (B), λ min (A • B) ≥ min iA ii λ min (B). Lemma 6. [39] Let ⊗ denote the Kronecker product and * denote the column-wise Khatri-Rao product, we have (1+c) 2 . Then with probability at least 1ε over the random initialization, we have for k = 0, 1, 2, . . .(A ⊗ B)(C * D) = AC * BD, (A * B) ⊤ (A * B) = A ⊤ A • B ⊤ B.∥y -v k ∥ 2 2 ≤ (1 -α) k ∥y -v 0 ∥ 2 2 .Proof. Consider the predictive error at the (k + 1)-th iteration, we haveEQUATIONNext, we need to estimate the bound of predictionv k+1 -v k = v(θ k -α(F k ) -1 (J k ) ⊤ (v k -y)) -v k = - 1 ξ=0 αJ ξ (F k ) -1 (J k ) ⊤ (v k -y)ξ . = - 1 ξ=0 αJ k (F k ) -1 (J k ) ⊤ (v k -y)ξ . + 1 ξ=0 α(J k -J ξ )(F k ) -1 (J k ) ⊤ (v k -y)ξ . = -αJ k (F k ) -1 (J k ) ⊤ (v k -y) Term 1 + α 1 ξ=0 (J k -J ξ )ξ . (F k ) -1 (J k ) ⊤ (v k -y) Term 2, where J ξ = ∂v(θ ξ ) ∂θ ξ , andθ ξ = ξθ k + (1 -ξ)θ k+1 = θ k -ξα(F k ) -1 (J k ) ⊤ (v k -y).We first analyse Term 1. We omit the index k in J, G and F for simplicity.Term 1 = -αJF -1 J ⊤ (v k -y) =α(X ⋆ G)[ηZ -1 X,G ⊗ (G ⊤ G) -1 ](X ⊤ * G ⊤ )(y -v k ) =α(ηXZ -1 X,G X ⊤ • G(G ⊤ G) -1 G ⊤ )(y -v k ) =α(ηXZ -1 X,G X ⊤ • I)(y -v k ),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 6. The last equation uses the definition of generalized inverse as given by Eq. ( 16). By Lemma 5, we haveEQUATIONTherefore, we can bound Term 1 byEQUATIONBased on the Condition 2, we have the following inequalityEQUATIONNext, we bound Term 2. By Eq. ( 23), we haveEQUATIONEQUATIONSimilar to Eq. ( 21), by Lemma (5) we can proveEQUATIONCombine Eq. ( 24) and Eq. ( 26), we haveEQUATIONCombine Eq. ( 20), Eq. ( 27) and Eq. ( 22), we can obtain∥y -v k+1 ∥ 2 2 =∥y -v k ∥ 2 2 -2(y -v k ) ⊤ (v k+1 -v k ) + ∥v k+1 -v k ∥ 2 2 =∥y -v k ∥ 2 2 -2α(y -v k ) ⊤ J k (F k ) -1 (J k ) ⊤ (y -v k ) +2α(y -v k ) ⊤ 1 ξ=0 (J k -J ξ )ξ . (F k ) -1 (J k ) ⊤ (y -v k ) +∥v k+1 -v k ∥ 2 2 ≤ 1 - 2αη λ max (Z X,G ) + 2cαη λ max (Z X,G ) + α 2 η 2 λ 2 min (Z X,G ) + 2cα 2 η 2 λ max (Z X,G )λ min (Z X,G ) + c 2 α 2 η 2 λ 2 max (Z X,G ) ∥y -v k ∥ 2 2 .In the last second inequality, we use the fact thatλ min (XZ -1 X,G X ⊤ • I) ≥ 1 λmax(Z X,G ) . Let - αη λ max (Z X,G ) + 2cαη λ max (Z X,G ) + α 2 η 2 λ 2 min (Z X,G ) + 2cα 2 η 2 λ max (Z X,G )λ min (Z X,G ) + c 2 α 2 η 2 λ 2 max (Z X,G ) ≤ 0, we have α ≤ (1 -2c)λ max (Z X,G ) (1 + c) 2 η ≤ (1 -2c) (1 + c) 2 , and ∥y -v k ∥ 2 2 ≤ (1 -α) k ∥y -v 0 ∥ 2 2 .This completes the proof.So far, we have already proved Theorem 3 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 1 and 2 hold, then as long as λ min (GG ⊤ ) ≥ 1 2 λ K , we have∥θ k+1 -θ 0 ∥ 2 ≤ 2∥y -v 0 ∥ 2 λ K /2 κ Z X,G .Proof. By the update rule of LNGD, we have∥θ k+1 -θ 0 ∥ 2 = k t=0 α(F t ) -1 (J t ) ⊤ (y -v t ) 2 ≤ α k t=0 ∥∥Z -1 X,G X ⊤ * (G t ) ⊤ (G t (G t ) ⊤ ) -1 ∥ 2 ∥y -v t ∥ 2 ≤ α k t=0 1 λ min (Z X,G ) 1 λ K /2 ∥y -v t ∥ 2 ≤ α k t=0 2/λ K λ min (Z X,G ) 1 - α λ max (Z X,G ) t/2 ∥y -v 0 ∥ 2 ≤ 2∥y -v 0 ∥ 2 λ K /2 λ max (Z X,G ) λ min (Z X,G ) = 2∥y -v 0 ∥ 2 λ K /2 κ Z X,G .This completes the proof.The training of ResNet-18 [44] on the CIFAR-10 [40] dataset serves as a fundamental experiment in the field of image classification. In this subsection, we present a 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 KFAC and LNGD are set to be 100. All experimental runs are conducted over a duration of 200 epochs.The implementation of ResNet50 [44] follows the TensorFlow version which can be found in the website 3 . We use the linear warmup strategy [41] 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 to 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 α t = 0.001 + 0.5 * (α 0 -0.001) * (1 + cos(2 * 0.47 * π * t/max_epoch)), where t is the number of epochs. For Adam, KFAC and LNGD, the learning rate uses the exponential updating strategyα t = α 0 * (1 -t/max_epoch) E, where E is decay rate ∈ {2, 3, 4, 5, 6}. α 0 is the initial learning rate tuned using a grid search with values α ∈ {1e -4, 1e -3, . . . , 1}.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 on the LNGD's performance. The variant denoted as LNGD-lr 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-1K dataset, All hyperparameters are maintained consistent with those outlined in the ImageNet training section.The results of the ablation experiments, as shown in Fig. 7 and Table5, reveal some interesting findings. Specifically, analyzing the training loss and testing accuracy versus epoch, we observe that LNGD-lr 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 to the fact that LNGD-lr 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-lr also takes 5% more time compared to LNGD-sample, which only utilizes the sampling approximation strategy. Notably, LNGD-sample 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 LNGD-lr in reaching a 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. In order to further validate the effectiveness of LNGD, we conduct additional experiments on the CIFAR-10 dataset, in 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 per epoch. Additionally, due to the efficient Fisher information matrix approximation strategy adopted by NG+, it can significantly reduce the computational time compared to EKFAC and TKFAC.NeurIPS Paper Checklisthttps://github.com/google-deepmind/dm-haiku/tree/main/examples/imagenet 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. **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.
3. **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.
4. **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.
5. **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.

Minor Comments
- **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).
- **Notation Consistency**: The notation $F_ {LNGD}$ is not defined.
- **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.
- **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
- **Algorithm Description**: The description of Algorithm 1 should not be placed in the appendix.
questions: questions: See Weaknesses.
limitations: limitations: See Weaknesses.
paper_summary: paper_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*.
strength_weakness: strength_weakness: 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.

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

## Speaking of local learning, the authors did not touch on the topic at all. Thus, several relevant references were missed, such as:
* 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ø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
* 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
* 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
* Ren, M., Kornblith, S., Liao, R., & Hinton, G. (2022). Scaling Forward Gradient With Local Los
* 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
* 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

Moreover, 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.
questions: 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.
* Please explicitly define the vector, **v**, of Thm 2 (near line 255). It is only defined in the appendix.
* In Algo 1 you use both T_FIM and T_Fisherinformationmatrix. Please be consistent.
limitations: limitations: * What is your justifcation for calling your three experiments ``extensive''? Am I being too grumpy here? ;-)
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

- 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?
questions: questions: - I like Figure 1. But it would be interesting to also compare this with KFAC.

- In Section 3.2, doesn't having adaptive learning rates contradict the spirit of the natural gradient?
    - 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?

- Can you also apply your adaptive learning rate strategy to KFAC?
limitations: limitations: The authors provide limitations for each theory.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: 1. What is $\mathbf{s}_l$ in line 136 ?
2. 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: 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.

Comments: title: Thank you for feedback
comment: Dear reviewer:

We would like to express our gratitude for all the feedback provided. Additionally, we hope that the additional responses above address your concerns and contribute to an enhancement in your overall rating of this manuscript.
Dear reviewer:

We sincerely appreciate your taking the time to review our response. We hope that the above clarifications address your concerns and contribute positively to your overall assessment of the paper.
Q 1. About the reference.

The related work is presented as follows：

[15] 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.

[16] Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Chuan-Sheng Foo, and Rio Yokota. Scalable and practical natural gradient for large-scale deep learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(1):404-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. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(5):5391-5403, 2022.

[22] 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.

[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 on Artificial Intelligence, volume 35, pages 7046-7054, 2021.

Q 2. About the integration of the newly proposed strategies into existing NGD methods.


Thank you for your response. As previously mentioned, layer-wise sampling and adaptive layer-wise learning rates can indeed be integrated with KFAC and its variants, such as EKFAC and TKFAC. However, it is important to note that the incorporation of these techniques into KFAC, EKFAC, and TKFAC results in new variants of these methodologies. Consequently, it would not be appropriate to conduct comparative experiments using these resulting new variants against LNGD. Actually, to facilitate a comprehensive understanding of the distinct contributions of various components within the LNGD framework, we have conducted a series of ablation studies, detailed in Appendix F.3 of our manuscript. These experiments aim to isolate the effects of adaptive learning rates and sampling optimization on LNGD performance. 
Furthermore, as you suggested, for a more thorough analysis of the impact of layer-wise sampling and adaptive layer-wise learning rates as independent components on the performance enhancement of existing methods such as EKFAC and TKFAC, it would be beneficial to incorporate these techniques into EKFAC and TKFAC and compare their performance against the baseline.
Thank you for your clarifying response, it is much appreciated. I just had a few more questions:

- For the 91% and 75.9% benchmarks, can you provide references to the literature where these numbers are utilized? I'm so sorry, but I'm having trouble find it.

- I'm still a bit uncomfortable with the "bundling" of many methods, and not providing similar "boosts" to the other methods in the literature. It does seem like multiple new methods are being introduced in the paper. Thus, I think it would be more scientifically enlightening to have a study where these new methods also apply to previous methods in the literature, where applicable. This is so the contribution of each part can be examined, and thus the causes of the performance increases can be made more clear.

Thank you again, and I do find the research direction and the method quite interesting. And please correct me if I missed something.
Regarding the open-sourcing of the code, as we are affiliated with an internet company, we initiated the process for code release upon receiving the first round of review comments. Given the multiple stages of the internal workflow, we are currently at the final stage of this process. We assure the reviewers and the area chair that the code will be made publicly available immediately upon completion of this final step. 

In regard to the error bars, we conducted multiple rounds of hyperparameter searches throughout the experimental process and selected the best results for final comparison; consequently, no error bars were included. If quite necessary, we can incorporate error bars in the final version.
Thank you for your valuable suggestion. Regarding the open-sourcing of the code, as we are affiliated with an internet company, we initiated the process for code release upon receiving the first round of review comments. Given the multiple stages of the internal workflow, we are currently at the final stage of this process. We assure the reviewers and the area chair that the code will be made publicly available immediately upon completion of this final step. 

In regard to ambiguous descriptions such as "extensive experiments," we will revise and refine these aspects in the final version. We sincerely apologize for any confusion caused.
Dear Authors,

Thanks for your careful response.

I still *highly recommend* sharing your code. Also, please do not say "extensive experiments" in your paper; you are still far from that in my opinion ;-)

I maintain my original rating of "Accept".
Dear Reviewer k9Mw,

Could you please respond with how you think about the authors' response? Please at least indicate that you have read their responses.

Thank you,
Area chair
Dear Authors,

Since Reviewer k9Mw has asked, could you please explain why the codes were not included in the submission? Also why weren't the error bars created in the plots?

Thank you,
AC
Dear Reviewer HWh2,

Could you please respond with how you think about the authors' response? Please at least indicate that you have read their responses.

Thank you, Area chair
Thank you for your new valuable comments and suggestions.

We will add explanations in the final version to make a clear distinction between THOR and LNGD. We will also try to perform experiments to compare their performance.

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.

We will add figures and explanations in the final version to illustrate the difference between the exact FIM and the approximated FIM in LNGD along the diagonal.
Thank you for your new valuable comments and questions.

Q 1. The new methods and experiments. 

A 1. Thank you for your question. We are sorry for the lack of detailed descriptions of the new methods’ setting in the reply. The code of EKFAC and NG+ are obtained from GitHub (links will be added in the final version since they are not permitted in our comment), and the code of TKFAC is obtained from its authors. The hyperparameters, such as the learning rate, the update of the curvature matrix strategy, the damping parameter, etc., are used according to the experimental settings in their papers and codes. We will add the detailed descriptions in the final version. Regarding the rationale for selecting accuracy benchmarks of 91% and 75.9%, we predominantly adhere to the settings established in the literature on natural gradient descent, specifically referencing several variants of KFAC.

Q 2. KFAC with the adaptive learning rates and sampling technique.

A 2. Thank you for your insightful question. We would like to take this opportunity to clarify the main contributions of our work. In this paper, we introduce a novel optimizer, LNGD, which builds upon existing natural gradient methodologies. Central to our approach are the novel layer-wise sampling strategy and the adaptive layer-wise learning rate strategy, which are essential components for executing the natural gradient optimization process within LNGD. As mentioned in our previous response, layer-wise sampling and adaptive layer-wise learning rates can also be integrated with KFAC and its variants, including EKFAC and TKFAC. However, once these techniques are incorporated into KFAC, EKFAC, and TKFAC, the resulting methods may be regarded as new variants of these methods rather than the original ones since some new techniques are incorporated.

Certainly. To illustrate the effects of incorporating sampling strategy and adaptive learning rate into KFAC, we conducted experiments on the CIFAR-10 dataset. We denote L-KFAC as a model variant where the sampling strategy, the adaptive learning rate and the new damping have been integrated into KFAC. On the CIFAR-10 dataset, when the top-1 testing accuracy reached 91%, both L-KFAC and LNGD utilized 36 epochs; however, L-KFAC incurred a time cost of 5.62 seconds per epoch, in contrast to LNGD, which took only 5.08 seconds. This highlights the effectiveness of our new factorization scheme. We will explore the application of the layer-wise sampling and learning rate adjustments to other natural gradient descent methods and hope to incorporate pertinent content in the final version.

Q 3. KFAC is more accurate than LNGD.

A 3. Regarding the Fisher information matrix (FIM), KFAC can indeed provide a more accurate approximation compared to LNGD. However, the purpose of proposing the LNGD optimizer is to achieve an optimal balance between the training speed and accuracy of one model, and to make every effort to reduce computational cost while preserving the main information of the FIM as much as possible. Therefore, the approximation accuracy of LNGD to FIM may not be the highest, but LNGD achieves a "close-enough" approximation much faster.

Q 4. Compute the Hessian matrix.

A 4. In the context of LNGD, computation of the Hessian matrix is not required. When calculating the adaptive learning rate within LNGD, we approximate the Hessian matrix using the Fisher Information Matrix (FIM). This is justified by the equivalence of the FIM and the Hessian matrix under various loss functions, such as cross-entropy and mean squared loss. Accordingly, we compute the FIM and employ Eq. (12) to derive the adaptive learning rate for LNGD. However, if our proposed adaptive layer-wise learning rate is applied  to other methods without FIM, the Hessian matrix or its approximation is need to be computed.
This reviewer thanks the authors for their feedback.

This reviewer upholds his/her decision on the paper. There is solid contribution 

However, here are some minor changes the reviewer would suggest
1. Make the clear distinction between THOR and the proposed approach. At a first glance, the methods seem to have a lot of overlap, albeit the contributions might be different. It would be interesting to see a comparative performance of THOR against LNGD as well (if the code is reproducible and available to use).
2. Reiterating on reproducibility - It would be interesting to release a reproducible version of this code to the community.
3. Absolute difference between the exact FIM and the proposed approach FIM $\textbf{along the diagonal}$ - It would be interesting to see how different the diagonal elements of the FIM are from the exact version. It would be great if that information is provided in the final version of the paper as well.
Thank you for your clarifications. I had a few more questions.

Reply to Q.1 reply: For the first part of the weakness, yes I already saw Appendix B, but I was looking for performance comparisons. I see now that the authors have done this in the reply to all reviewers, thank you very much, I really appreciate it. But now can the authors expand on the experimental procedure with the new methods? For example, did the authors take the code from public code repositories and run it on the datasets? What hyperparameters were used? Also, apologies if I'm ignorant, but I just realized the oddity of using 91% and 75.9% as accuracy benchmarks. Why were these specific numbers chosen?

Reply to Q.2 reply: For the second part of the weakness in my original review, I did see Section F.3 (which wasn't referenced in the main text, but perhaps should be referenced around line 260), but perhaps I wasn't being clear, my apologies. I was hoping to see the adaptive learning rates and sampling technique also applied to KFAC if applicable, and in light of the new experiments but me being sympathetic to the large amount of work involved, perhaps just to NG+. This is so we can examine if it's the novel factoring that shows impressive performance (which seems to me the main contribution of the paper), or the adaptive learning rate/sampling, or the new damping technique. I even see throughout the paper (Line 175-178 with adaptive learning rates, and Line 225-228 with damping) that adaptive learning rates and damping techniques are required to achieve optimal performance, so I would have hoped these advantages would be conferred onto KFAC as well, and now NG+ if they were also conferred to the proposed method, LNGD. I do realize this is a lot of work in a limited time, but I also believe this would make for a fairer comparison.

Reply to Q.3 reply: From my eyes, and perhaps I'm looking at it incorrectly, but it seems KFAC is more accurate than LNGD? But are the authors arguing that LNGD achieves a "close-enough" approximation much faster?

Reply to Q.4 reply: Thank you for the answer, it was enlightening. Apologies, but just to reclarify for my own ignorance: so to compute the adaptive learning rate, the Hessian matrix must be computed?

Reply to Q.5 reply: See my answers to "Reply to Q.1" and "Reply to Q.2" above.