Title: A Layer-Wise Natural Gradient Optimizer for Training Deep Neural Networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: Natural gradient method

second-order optimization

deep neural networks

Kronecker factorization
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-Wise Natural Gradient Optimizer for Training Deep Neural 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 layerwise 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 backpropagation. 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.

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

2 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 as F -1 ∇ θ h ( θ ) , where F is the Fisher information matrix given by

In 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 as F = 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.,

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

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-Wise Sample Approximation

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

0

200

0

100

200

300

0

200

0

100

200

300

0

200

0

100

200

300

0

20

0

10

20

30

0

20

0

10

20

30

0

20

0

10

20

30

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 l th 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 as

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

Combining 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 by

On the other hand, Ψ l can be computed by

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.

Based on Eq. (4) and Eq. (5), we can show that tr ( F l ) = tr ( Φ l ⊗ Ψ l ) .

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

Once the update direction is chosen, the minimizer of α k can be given by

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 is

̸

Taking the derivative of h with respect to α k and letting it equal to 0 , we get

if 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 that d k = (( d k 1 ) ⊤ , ( d k 2 ) ⊤ , . . . , ( d k L ) ⊤ ) ⊤ is the update direction of a L layers neural network at the iteration point θ k = (( θ k 1 ) ⊤ , ( θ k 2 ) ⊤ , . . . , ( θ k L ) ⊤ ) ⊤ , the update rule of θ k is given as θ k +1 = θ k -˜ D k ˜ α k , where

is a block diagonal matrix and ˜ α k ∈ R L is the learning rate. The approximate second Taylor expansion of h ( θ -˜ D ˜ α ) at iteration k is

where H 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 ( ˜ D k ) ⊤ ∇ θ h ( θ k ) = ( ˜ D k ) ⊤ H ⊤ ˜ D k ˜ α k , which yields that

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 get

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 choose d 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 by

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

which yields that

3.3 Algorithm Schema

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 l th 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,

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

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 ϵ , with the estimate calculated from the new mini-batch, weighted by 1 -ϵ . That is

In summary, our proposed algorithm is shown in Algorithm1.

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 [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:

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 -2 c ) (1+ c ) 2 . Then with probability at least 1 -ε over the random initialization, we have for k = 0 , 1 , 2 , . . .

Proof. The proof is given in the appendix.

Algorithm 1 LNGD

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 .

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 are tuned using a grid search with values α ∈ { 1 e -4 , 3 e -4 , . . . , 1 , 3 } . The damping parameter λ in KFAC[14] are tuned using a grid search with values λ ∈ { 1 e -6 , 1 e -4 , 3 e -4 , 1 e -3 , . . . , 1 e -1 , 3 e -1 } . The minimum and maximum of damping parameters ν 1 and ν 2 in LNGD are set to 1 e -5 and 1 e -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.

4.1 CIFAR-10 Training

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.

0

50

100

150

200

0.0

0.5

1.0

1.5

Training Loss

Epoch

SGD

ADAM

KFAC

LNGD

0

100

200

300

400

500

600

0.0

0.5

1.0

1.5

Training Loss

Time/Seconds

SGD

ADAM

KFAC

LNGD

0

50

100

150

200

50

60

70

80

90

Testing Accuracy

Epoch

SGD

ADAM

KFAC

LNGD

91

0

100

200

300

400

500

600

50

60

70

80

90

Testing Accuracy

Time/Seconds

SGD

ADAM

KFAC

LNGD

91

0.5

1.0

1.5

Training Loss

SGD

ADAM

KFAC

LNGD

0.5

1.0

1.5

Training Loss

SGD

ADAM

KFAC

LNGD

0.5

1.0

1.5

Training Loss

SGD

ADAM

KFAC

LNGD

0.5

1.0

1.5

Training Loss

SGD

ADAM

KFAC

LNGD

0.0

0.0

100

Epoch

Figure 3: Comparison of the exact Fisher information matrix

F

and our approximation

F

150

200

LNGD. We

20-10. We show the results of the Fisher information matrix of the first layer with 20 units in top,

which is a exact Fisher information matrix

with red square, which is a where

θ

1

Let

weight of unit

v

= [

f

(

·

convergence of Algorithm 2, the specific proof is given in the appendix.

Theorem 2.

rank

θ

r

Proof.

the fastest rate.

4

, . . . , f

{

}

θ

, a,

x

i

)

, f

(

θ

, a,

x

2

)

(

θ

, a,

x

n

)]

/latticetop

.

In the following, we only give the result of

(Convergence rate of LNGD) Under the Assumption 1 and the assumption that

(

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

∼

N

(0

, ν

I

)

,

a

r

∼

unif[

{-

1

,

+1

}

]

for any

r

∈

[

m

]

, and we set the step size

α

(1

-

2

c

)

(1+

c

)

2

.

Then with probability at least

1

-

ε

over the random initialization, we have for

k

= 0

,

1

,

2

, . . .

0

10

20

30

40

50

60

70

80

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

Training Loss

Epoch

SGD

ADAM

KFAC

LNGD

0

10000

20000

30000

40000

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

Training Loss

Time/Seconds

SGD

ADAM

KFAC

LNGD

0

10

20

30

40

50

60

70

80

20

30

40

50

60

70

80

SGD

ADAM

KFAC

LNGD

75.9

Testing Accuracy

Epoch

0

10000

20000

30000

40000

20

30

40

50

60

70

80

Testing Accuracy

Time/Seconds

SGD

ADAM

KFAC

LNGD

75.9

2

k

2

The proof is given in the appendix.

Experiments comparison, as presented in Table.1, SGD and ADAM require at least 100 4.2 ImageNet Training

y

v

0

2

.

It only requires 36 epochs and 189.69 seconds to achieve this accuracy level.

%

and 30

%

more epochs

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 α ∈ { 1 e -4 , 3 e -4 , . . . , 1 , 3 } . The damping parameter λ in KFAC[14] are tuned using a grid search with values λ ∈ { 1 e -6 , 1 e -4 , 3 e -4 , 1 e -3 , . . . , 1 e -1 , 3 e -1 } . The minimum and maximum of damping parameters ν 1 and ν 2 in LNGD are set to 1 e -5 and 1 e -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. 4.1 CIFAR-10 Training epochs and time by around 20 % and 21 % , respectively. Furthermore, as shown in Table.1, LNGD gets the highest fincal testing accuracy after convergence. Table 1: Detailed statistics on 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% 4.2 ImageNet Training Table 2: Detailed statistics on ImageNet when top-1 testing accuracy achieves 75.9%. Epoch Total Time Time Per Epoch Acceleration Best Test Acc 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 and time, respectively, to achieve similar accuracy. Relative to KFAC, LNGD reduces the number of

11.28h

520.55s

43%

76.47%

optimization methods [6, 36, 37, 38]. Following these previous works, the changes of testing accuracy

We first report the optimizing performance on CIFAR-10, which is a standard task used to benchmark 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 KFAC 39 8.02h 739.93s 19% 76.58% LNGD 36 6.46h 646.44s 76.73% We extend our examination of optimizer efficacy to a larger image classification dataset, ImageNet-1K [39]. The changes of testing accuracy and training loss versus time and epoch are reported in Fig.4 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.

SGD 78 ADAM -4.3 Transformer Training

-

-

-

74.05%

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 8 and Table.2. 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 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,

9

k

0

50

150

.

5

In summary, we propose a novel NGD optimizer named as LNGD for training DNNs, specifically

4

0

20000

40000

60000

80000

100000

1

2

3

Training Loss

Step

SGD

ADAM

KFAC

LNGD

0

10000

20000

30000

40000

50000

1

2

3

Training Loss

Time/Seconds

SGD

ADAM

KFAC

LNGD

0

20000

40000

60000

80000

100000

20

25

30

Step

SGD

ADAM

KFAC

LNGD

32

0

10000

20000

30000

40000

50000

20

25

30

Time/Seconds

SGD

ADAM

KFAC

LNGD

32

4

35

35

that when training large-scale image tasks, a trade-off between efficiency and effectiveness needs

The training datasets is WMT English-German machine translation corpus [40]. We use Bleu [41] as 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. 4.3 Transformer Training In this experiment, we apply LNGD to the Transformer-Big model [40] with 213.7M parameters. 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.

the evaluation metrics, which is frequently used in machine translation tasks. The setting of learning

for transformer models. KFAC exhibits further enhancements in performance compared to Adam, reduction in training loss and the highest convergence rates. In terms of testing accuracy, measured

yet it does not surpass the efficacy of LNGD. LNGD outperforms its counterparts with the swiftest

%

by the Bleu score, LNGD achieves a top-1 Bleu score of 32 5 Conclusion

with remarkable efficiency, which is

%

and computing time by 16

-

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. Table 3: Detailed statistics on WMT when Bleu achieves 32%. Step Total Time Time Per 1K Acceleration Best Test Bleu 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.

able to reduce the required steps by approximately 24

SGD

-

22K

ADAM KFAC References

-

3.75h

-

613.63s

31.87%

24%

33.45%

25K

3.39h

488.16s

16%

33.05%

LNGD 19K 2.85h 540s 33.55% [1] Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathematical Statistics , pages 400-407, 1951.

Conclusion [2] Yu E Nesterov. A method for solving the convex programming problem with convergence rate o ( 1 k 2 ) . In Doklady Akademii Nauk SSSR , volume 269, pages 543-547, 1983.

targeting the computational inefficiencies that impede the practical application of conventional natural information matrices for each individual layers using sample approximation and dynamically adjusts

gradient techniques in large-scale neural networks.

Our approach strategically computes Fisher 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.

compared to

%

Ning Qian. On the momentum term in gradient descent learning algorithms. Neural networks , 12(1):145-151, 1999.

Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701 , 2012.

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):2631, 2012.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2014.

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 , 33:18795-18806, 2020.

Nitish Shirish Keskar and Albert S Berahas. ADAQN: An adaptive quasi-Newton algorithm for training rnns. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases , pages 1-16, 2016.

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 on artificial intelligence , volume 35, pages 10665-10673, 2021.

Donald Goldfarb, Yi Ren, and Achraf Bahamou. Practical quasi-Newton methods for training deep neural networks. Advances in Neural Information Processing Systems , 33:2386-2396, 2020.

James Martens. Deep learning via Hessian-free optimization. In International Conference on Learning Representations , volume 27, pages 735-742, 2010.

Oriol Vinyals and Daniel Povey. Krylov subspace descent for deep learning. In Artificial Intelligence and Statistics , pages 1261-1268. PMLR, 2012.

Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation , 10(2):251276, 1998.

James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International conference on machine learning , pages 2408-2417, 2015.

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.

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.

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.

James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research , 21:1-76, 2020.

Roger Grosse and James Martens. A Kronecker-factored approximate Fisher matrix for convolution layers. In International Conference on Machine Learning , pages 573-582, 2016.

James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations , 2018.

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.

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.

Kaixin Gao, Zheng-Hai Huang, Xiaolei Liu, Min Wang, Shuangling Wang, Zidong Wang, Dachuan Xu, and Fan Yu. Eigenvalue-corrected natural gradient based on a new approximation. Asia-Pacific Journal of Operational Research , 40(01):2340005, 2023.

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 layeroriented natural gradient descent computation. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 7046-7054, 2021.

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):23652378, 2023.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Krishnakumar Balasubramanian, Pinar Donmez, and Guy Lebanon. Unsupervised supervised learning ii: Margin-based classification without labels. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics , pages 137-145. JMLR Workshop and Conference Proceedings, 2011.

Frédéric de Gournay and Alban Gossard. Adaptive scaling of the learning rate by second order automatic differentiation. arXiv preprint arXiv:2210.14520 , 2022.

Bharat Singh, Soham De, Yangmuzi Zhang, Thomas Goldstein, and Gavin Taylor. Layerspecific adaptive learning rates for deep networks. In 2015 IEEE 14th International Conference on Machine Learning and Applications (ICMLA) , pages 364-368. IEEE, 2015.

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.

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.

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.

Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Proceedings of the 36th International Conference on Machine Learning , volume 97, pages 322-332, 2019.

Guodong Zhang, James Martens, and Roger Grosse. Fast convergence of natural gradient descent for over-parameterized neural networks. In Advances in Neural Information Processing Systems , volume 32, 2019.

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.

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.

Joel A Tropp. An introduction to matrix concentration inequalities. arXiv preprint arXiv:1501.01571 , 2015.

Jssai Schur. Bemerkungen zur theorie der beschränkten bilinearformen mit unendlich vielen veränderlichen. Journal für die reine und angewandte Mathematik , 1911(140):1-28, 1911.

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.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Peng Xu, Fred Roosta, and Michael W Mahoney. Second-order optimization for non-convex machine learning: An empirical study. In Proceedings of the 2020 SIAM International Conference on Data Mining , pages 199-207. SIAM, 2020.

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 Pattern Recognition , pages 6438-6447, 2021.

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.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics , pages 311-318, 2002.

A Notations

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 matrices

we use ◦ and ⊗ to denote the Hadamard and Kronecker product, respectively. The column-wise Khatri-Rao product ∗ is defines as

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

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 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 [(ˆ g l ˆ g ⊤ l )] with ˆ g 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 by S ∗ 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.

⊗

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

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

(a)

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

(b)

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

(c)

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

(d)

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

(e)

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

(f)

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.

B.2 Comparisons Between KFAC and LNGD

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.

B.3 Illustration of the Gaussian Distribution Assumption

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. Fig. 6 (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. Fig. 6 (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.

values

count

(a)

values

count

(b)

values

count

(c)

values

count

(d)

C Proof of Theorem 1

Proof. Since g ( θ ) and g L ( θ ) are the approximate second Taylor expansions of h ( θ -α d ) and h ( θ -˜ D ˜ α ) , we have

This completes the proof.

D Convergence of LNGD

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:

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 where 0 < ν ≤ 1 controls the magnitude of initialization.

Given the training dataset 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.

̸

In this subsection, we mainly focus on the mean squared error loss (MSE) function

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

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 by

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 by where

Let J = ( J 1 , J 2 , . . . , J n ) ⊤ ∈ R n × md , then the Fisher information matrix can be written as

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]

and the update rule of NGD can be written as

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

For simplicity, we ignore the index of layer. Define and

where X is the input matrix formed by n input vectors and G is the pre-activation derivatives matrix whose entry is given by

Then we have and

The formulas of Φ and Ψ can be given as

Under the Assumption 1, we have ( XX ⊤ ) ii = 1 , i ∈ [ d ] . Therefore, Eq. (18) and Eq. (19) can be simplified to

where η = tr (( XX ⊤ ◦ I ) ◦ ( GG ⊤ ◦ I )) and Z X , G = X ⊤ ( GG ⊤ ◦ I ) X . Finally, the update rule of LNGD can be given by

To analyze the global convergence of LNGD, we need the following two conditions as given in [34].

Condition 1. The matrix G 0 ( G 0 ) ⊤ is positive definite.

To verify this condition, we need the following two lemmas.

Lemma 1. Define K ∞ i,j = E θ ∼N ( 0 ,ν 2 I ) [ I { θ ⊤ x i ≥ 0 , θ ⊤ x j ≥ 0 } ] , i, j ∈ [ n ] , then we have the matrix K ∞ is strictly positive definite.

The result of this lemma has been given and discussed in [34]. We define λ K = λ min ( K ∞ ) ≥ 0 and matrix K whose entry is given by

Lemma 2. If m = Ω ( n λ K log n ε ) , we have with probability at least 1 -ε that λ min ( K (0)) ≥ 3 4 λ K .

Then we can show Condition 1 holds by the following lemma.

Proof. Note that K (0) can be written as the sum of random symmetric matrices

Furthermore, we have K ( θ r ) are positive semi-definite and ∥ K ( θ r ) ∥ 2 ≤ tr ( K ( θ r )) ≤ n m . Thus, by the matrix Chernoff bound [37], we can obtain

Condition 2. For all parameters θ that satisfy ∥ θ -θ 0 ∥ 2 ≤ 2 ∥ y -v 0 ∥ 2 √ λ K / 2 κ Z X , G , there exists 0 ≤ c < 1 2 such that

This proof is similar to the Lemma 6 in [34], the difference is the definition of K . For completeness of the proof, we also give the detailed proof here. By this lemma, we have λ min ( G (0) G (0) ⊤ ) = λ min ( K (0)) ≥ 3 4 λ K > 0 , which implies that Condition 1 holds. Next, we give the other condition.

To show this condition holds, we need the following lemma.

Lemma 3. [34] For all weight vectors θ that satisfy ∥ θ -θ 0 ∥ 2 ≤ R , we have probability at least 1 -ε that

By taking R = 2 ∥ y -v 0 ∥ 2 √ λ K / 2 κ Z X , G , we have ∥ J -J 0 ∥ 2 2 ≤ 64 1 / 3 n ∥ y -v 0 ∥ 2 / 3 2 κ 2 / 3 Z X , G ν 2 / 3 ε 2 / 3 m 1 / 3 λ 1 / 3 K . Therefore, if we let m = Ω ( n 3 ∥ y -v 0 ∥ 2 2 κ 8 Z X , G ν 2 ε 2 λ 4 K ) , the Condition 2 holds. What's more, we have probability at least 1 -ε that ∥ y -v 0 ∥ 2 2 = O ( n ε ) , which has been given in [31]. Thus we can write the condition of m as m = Ω ( n 4 κ 8 Z X , G ν 2 ε 3 λ 4 K ) .

Lemma 4. If m = Ω ( n 4 κ 8 Z X , G ν 2 ε 3 λ 4 K ) , for all parameters θ that satisfy ∥ θ -θ 0 ∥ 2 ≤ 2 ∥ y -v 0 ∥ 2 √ λ G / 2 κ Z X , G , we have probability at least 1 -ε that λ min ( K ) ≥ √ λ K / 2 .

Before giving the main result, we first give some necessary lemmas.

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

Therefore,

Note that the large m is, the smaller c is. Therefore, we can choose a slight larger m satisfying this inequality. So we have

Proof complete.

Lemma 5. [38] Let A and B be two positive define matrices, we have

Lemma 6. [39] Let ⊗ denote the Kronecker product and ∗ denote the column-wise Khatri-Rao product, we have

Now, we give the convergence analysis of LNGD.

Theorem 3. (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 -2 c ) (1+ c ) 2 . Then with probability at least 1 -ε over the random initialization, we have for k = 0 , 1 , 2 , . . .

Proof. Consider the predictive error at the ( k +1) -th iteration, we have

Next, we need to estimate the bound of prediction

∂ θ ---We first analyse Term 1. We omit the index k in J , G and F for simplicity.

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 have

Therefore, we can bound Term 1 by

Based on the Condition 2, we have the following inequality

Next, we bound Term 2. By Eq. (23), we have

Define ∆ = Z -1 X , G X ⊤ ∗ G ⊤ ( GG ⊤ ) -1 , then we have

Similar to Eq. (21), by Lemma (5) we can prove

Combine Eq. (24) and Eq. (26), we have

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

Proof. By the update rule of LNGD, we have

This completes the proof.

E Experiments

E.1 Setup of CIFAR-10

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.

E.2 Setup of ImageNet

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 α ∈ { 1 e -4 , 1 e -3 , . . . , 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 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.

E.3.2 Results

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.

5.0

Epoch

Total Time

Time Per Epoch

80

Acceleration

Best Test Acc

0

10

20

30

40

50

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

10

20

30

40

50

20

30

40

50

60

70

0

10000

20000

30000

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

5.0

10000

20000

30000

20

30

40

50

60

70

80

Training Loss

Epoch

LNGD

LNDG_lr

LNGD_sample

a

Test Accuracy

Epoch

LNGD

LNGD_lr

LNGD_sample

75.9

b

Training Loss

Time/Seconds

LNGD

LNGD_lr

LNGD_sample

c

Test Accuracy

Time/Seconds

LNGD

LNGD_lr

LNGD_sample

75.9

d

E.4 Results of More Comparisons

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

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

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 proofs in the paper should be numbered and crossreferenced.

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 to provide a 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 the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of 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.

If the paper includes experiments, a No answer to this question will not be 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 a 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 is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or 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 to replicate the results, access to a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

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.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

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

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

It 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 a 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 of range (e.g. negative error 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 in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide sufficient information on the computer resources needed to reproduce the experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or 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 as well as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

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

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

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing 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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 not involve crowdsourcing nor research with human subjects.

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

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
comments: 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.

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 
$\Psi_l \otimes \Phi_l$
2. 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.

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:
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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors provide limitations for each theory.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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*.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\n* What is your justifcation for calling your three experiments ``extensive''? Am I being too grumpy here? ;-)", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- **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.\n\nWEAKNESSES:\n1. **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.\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nSee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}