Tensor Normal Training for Deep Learning Models

Abstract:
Despite the predominant use of first-order methods for training deep learning models, second-order methods, and in particular, natural gradient methods, remain of interest because of their potential for accelerating training through the use of curvature information. Several methods with non-diagonal preconditioning matrices, including KFAC [34], Shampoo [18], and K-BFGS [15], have been proposed and shown to be effective. Based on the so-called tensor normal (TN) distribution [31], we propose and analyze a brand new approximate natural gradient method, Tensor Normal Training (TNT), which like Shampoo, only requires knowledge of the shape of the training parameters. By approximating the probabilistically based Fisher matrix, as opposed to the empirical Fisher matrix, our method uses the block-wise covariance of the sampling based gradient as the pre-conditioning matrix. Moreover, the assumption that the sampling-based (tensor) gradient follows a TN distribution, ensures that its covariance has a Kronecker separable structure, which leads to a tractable approximation to the Fisher matrix. Consequently, TNT’s memory requirements and per-iteration computational costs are only slightly higher than those for first-order methods. In our experiments, TNT exhibited superior optimization performance to state-of-the-art first-order methods, and comparable optimization performance to the state-of-the-art second-order methods KFAC and Shampoo. Moreover, TNT demonstrated its ability to generalize as well as first-order methods, while using fewer epochs.

heading: 1 Introduction
text: First-order methods are currently by far the most popular and successful optimization methods for training deep learning models. Stochastic gradient descent (SGD) [39] uses the (stochastic) gradient direction to guide its update at every iteration. Adaptive learning rate methods, including AdaGrad [11], RMSprop [20], and Adam [23], scale each element of the gradient direction (possibly modified to incorporate momentum) by the square root of the second moment of each element of the gradient. These first-order methods use little curvature information to "pre-condition" the gradient direction; SGD uses an identity pre-conditioning matrix, whereas the others use a diagonal matrix.
On the other hand, second-order methods attempt to greatly accelerate the optimization process by exploring the rich curvature information of the problem. Traditional second-order methods such as Newton’s method, BFGS [6, 13, 14, 41], and limited-memory BFGS (L-BFGS) [28], without modification, are not practical in a deep learning setting, because these methods require enormous amounts of memory and computational effort per iteration due to the huge number of parameters such models have. Some second-order methods have been proposed to deal with the non-convexity and stochasiticity of objective functions arising in machine learning (see e.g. [36, 7, 16, 44]), but directly using these methods to train deep learning models still requires large amounts of memory and computing resources.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Recently, there has been considerable advancement in the development of second-order methods that are suitable for deep learning models with huge numbers of parameters. These methods usually approach pre-conditioning of the gradient in a modular way, resulting in block-diagonal pre-conditioning matrices, where each block corresponds to a layer or a set of trainable parameters in the model. Inspired by the idea of the natural gradient (NG) method [1], [34] proposed KFAC, an NG method that uses a Kronecker-factored approximation to the Fisher matrix as its pre-conditioning matrix that can be applied to multilayer perceptrons, and which has subsequently been extended to other architectures, such as convolutional neural networks [17] and recurrent neural networks [35]. Kronecker-factored preconditioners [15, 38] based on the structure of the Hessian and quasi-Newton methods have also been developed. Despite the great success of these efficient and effective second-order methods, developing such methods requires careful examination of the structure of the preconditioning matrix to design appropriate approximations for each type of layer in a model.
Another well-recognized second-order method, Shampoo [18, 3], extends the adaptive learning rate method AdaGrad, so that the gradient is pre-conditioned along every dimension of the underlying tensor of parameters in the model, essentially replacing the diagonal pre-conditioning matrix of the adaptive learning rate methods by a block diagonal Kronecker-factored matrix which can be viewed as an approximation to a fractional power of the empirical Fisher (EF) matrix. However, while estimating the Fisher matrix, in a deep learning setting, by the EF matrix saves some computational effort, it usually does not capture as much valuable curvature information as the Fisher matrix [26].
Variants of the normal distribution, i.e. the matrix-normal distribution [9] and the tensor-normal distribution [31], have been proposed to estimate the covariance of matrix and higher-order tensor observations, respectively. By imposing a Kronecker structure on the covariance matrix, the resulting covariance estimate requires a vastly reduced amount of memory, while still capturing the interactions between the various dimensions of the respective matrix or tensor. Iterative MLE methods for estimating the parameters of matrix-normal and tensor-normal distributions have been examined in e.g. [12, 31], and various ways to identify the unique representation of the distribution parameters have been proposed in [43, 10]. However, to the best of our knowledge, this advanced statistical methodology has not been used to develop optimization methods for deep learning. In this paper, we describe a first attempt to do this and demonstrate its great potential.
Our Contributions. In this paper, we propose a brand new approximate natural gradient (NG) method, Tensor-Normal Training (TNT), that makes use of the tensor normal distribution to approximate the Fisher matrix. Significantly, the TNT method can be applied to any model whose training parameters are a collection of tensors without knowing the exact structure of the model.
To achieve this, we first propose a new way, that is suitable for optimization, to identify the covariance parameters of tensor normal (TN) distributions, in which the average eigenvalues of the covariance matrices corresponding to each of the tensor dimensions are required to be the same (see Section 3).
By using the Kronecker product structure of the TN covariance, TNT only introduces mild memory and per-iteration computational overhead compared with first-order methods. Also, TNT’s memory usage is the same as Shampoo’s and no greater than KFAC’s, while its per-iteration computational needs are no greater than Shampoo’s and KFAC’s (see Section 5).
The effectiveness of TNT is demonstrated on deep learning models. Specifically, on standard autoencoder problems, when optimization performance is compared, TNT converges faster than the benchmark first-order methods and roughly the same rate as the benchmark second-order methods. Moreover, on standard CNN models, when generalization is concerned, TNT is able to achieve roughly the same level of validation accuracy as the first-order methods, but using far fewer epochs (see Section 6).
We also prove that, if the statistics used in TNT can be estimated ideally, it converges to a stationary point under mild assumptions (see Section 4).

heading: 2 Preliminaries
text: Supervised Learning. Throughout this paper, we consider the classic supervised learning setting where we learn the parameters θ of a model, by minimizing Lpθq “ 1N řN i“1 lpyi, fθpxiqq, where tpxi, yiquNi“1 denotes a given dataset (xi being the input to the model and yi being the target), fθpxiq denotes the output of the model when xi is provided as the input, and l denotes a loss function
(e.g. least-squares loss for regression and cross entropy loss for classification) that measures the discrepancy between the model output fθpxiq and the target yi. Natural Gradient Method and the Fisher Matrix. In a first-order method, say SGD, the updating direction is always derived from an estimate to the gradient direction ∇θLpθq. In a natural gradient (NG) method [1], however, the Fisher matrix is used as a pre-conditioning matrix that is applied to the gradient direction. As shown in [34], the Fisher matrix is defined as
F “ Ex„Qx,y„pp¨|x,θq ” ∇θ log ppy|x, θq p∇θ log ppy|x, θqqJ ı , (1)
where Qx is the data distribution of x and pp¨|x, θq is the density function of the conditional distribution defined by the model with a given input x.
In many cases, such as when p is associated with a Gaussian distribution and the loss function l measures least-squares loss, or when p is associated with a multinomial distribution and l is crossentropy loss, log p is equivalent to l (see e.g. [33, 34]). Hence, if Dθ denotes the gradient of l w.r.t. θ for a given x and y, we have that F “ Ex„Qx,y„prDθDθJs. Consequently, one can sample x from Qx and perform a forward pass of the model, then sample y from pp¨|x, θq, and perform a backward pass to compute Dθ, and then use DθDθJ to estimate F . We call Dθ a sampling-based gradient, as opposed to the empirical gradient ∇θlpyi, fθpxiqq where pxi, yiq is one instance from the dataset. It is worth noting that the first moment of Dθ is zero. To see this, note that, with given x,
Ey„pr∇θ log ppy|x, θqs “ ż ∇θ log ppy|x, θqppy|x, θqdy “ ż ∇θppy|x, θqdy
“∇θ ˆ ż ppy|x, θqdy ˙ “ ∇θ1 “ 0.
Hence, Ex„Qx,y„prDθs “ Ex„Qx tEy„pr∇θ log ppy|x, θqs | xu “ 0. Thus, the Fisher matrix F can be viewed as the covariance matrix of Dθ. Note that the empirical Fisher CANNOT be viewed as the covariance of the empirical gradient, because the first moment of the latter is, in general, NOT zero.
Tensor-Normal Distribution. The development of our new method makes use the so-called tensornormal distribution [31, 10]:
Definition 1. An arbitrary tensor G P Rd1ˆ¨¨¨ˆdk is said to follow a tensor normal (TN) distribution with mean parameter M P Rd1ˆ¨¨¨ˆdk and covariance parameters U1 P Rd1ˆd1 , ..., Uk P Rdkˆdk if and only if vecpGq „ NormalpvecpMq, U1 b ¨ ¨ ¨ b Ukq.
In the above definition, the vec operation refers to the vectorization of a tensor, whose formal definition can be found in Sec A in the Appendix. Note that matrix-normal distribution can be viewed as a special case of TN distribution, where k “ 2. Compared with a regular normal distribution, whose covariance matrix has
śk i“1 d 2 i elements, the covariance of a k-way tensor-normal distribution
is stored in k smaller matrices with a total number of elements equal to řk i“1 d 2 i .
To estimate the covariance submatrices U1, . . . , Uk, the following property (e.g., see [10]) is used:
ErGpiqs “ Ui ¨ ź
j‰i trpUjq, (2)
where Gpiq :“ matipGqmatipGqJ P Rdiˆdi denotes the contraction of G with itself along all but the ith dimension and mati refers to matricization of a tensor (see Section A for the formal definitions). By (2), we can sample G to obtain estimates of the Gpiq’s, and hence, estimates of the Ui’s. The complexity of computing Gpiq is di śk j“1 dj , which is also far less than the complexity of computing vecpGqvecpGqJ needed to estimate the covariance of a regular normal distribution.

heading: 3 Tensor-Normal Training
text: In this section, we propose Tensor-Normal Training (TNT), a brand new variant of the natural gradient (NG) method that makes use of the tensor-normal distribution.

heading: 3.1 Block Diagonal Approximation
text: In this paper, we consider the case where the parameters of the model θ consists of multiple tensor variables W1, ...,WL, i.e. θ “ pvecpW1qJ, ..., vecpWLqJqJ. This setting is applicable to most common models in deep learning such as multi-layer perceptrons, convolutional neural networks, recurrent neural networks, etc. In these models, the trainable parameter Wl (l “ 1, . . . , L) come from the weights or biases of a layer, whether it be a feed-forward, convolutional, recurrent, or batch normalization layer, etc. Note that the index l of Wl refers to the index of a tensor variable, as opposed to a layer.
To obtain a practical NG method, we assume, as in KFAC and Shampoo, that the pre-conditioning Fisher matrix is block diagonal. To be more specific, we assume that each block corresponds to the covariance of a tensor variable in the model. Hence, the approximate Fisher matrix is:
F « diagLl“1 Ex„Qx,y„p “ vecpDWlqpvecpDWlqqJ ‰( “ diagLl“1 tVarpvecpDWlqqu .
The remaining question is how should one approximate VarpvecpDWlqq for l “ 1, ..., L.

heading: 3.2 Computing the Approximate Natural Gradient Direction by TNT
text: We consider a tensor variable W P Rd1ˆ¨¨¨ˆdk in the model and assume that G :“ DW P Rd1ˆ¨¨¨ˆdk follows a TN distribution with zero mean and covariance parameters U1, ..., Uk where Ui P Rdiˆdi . Thus, the Fisher matrix corresponding to W is FW “ Ex„Qx,y„prVarpvecpGqqs “ U1 b ¨ ¨ ¨ b Uk. Loosely speaking, the idea of relating the Fisher matrix to the covariance matrix of some normal distribution has some connections to Bayesian learning methods and interpretations of NG methods (see e.g., [22]). Let ∇WL P Rd1ˆ¨¨¨ˆdk denote the gradient of L w.r.t. W . The approximate NG updating direction for W is computed as
F´1W vecp∇WLq “ pU ´1 1 b ¨ ¨ ¨ b U´1k qvecp∇WLq “ vec ` ∇WLˆ1 U´11 ˆ2 ¨ ¨ ¨ ˆk U´1k ˘ , (3)
where ˆi (i “ 1, ..., k) denotes a mode-i product (see Section A in the Appendix). Note that the last equality of (3) makes use of the following proposition, which also appears in [18] (see Sec A in the Appendix for a proof):
Proposition 1. Let G P Rd1ˆ¨¨¨ˆdk and Ui P Rdiˆdi for i “ 1, ..., k. Then, we have `
bki“1Ui ˘ vecpGq “ vecpGˆ1 U1 ˆ2 U2 ¨ ¨ ¨ ˆk Ukq. (4)
To summarize, the generic Tensor-Normal Training algorithm is:
Algorithm 1 Generic Tensor-Normal Training (TNT) Require: Given batch size m, and learning rate α
1: for t “ 1, 2, . . . do 2: Sample mini-batch Mt of size m 3: Perform a forward-backward pass over Mt to compute the mini-batch gradient 4: Perform another backward pass over Mt with y sampled from the predictive distribution to compute Gl “ DWl (l “ 1, ..., L) averaged across Mt 5: for l “ 1, ...L do 6: Estimate ErGpiql s (i “ 1, ..., kl) from Gl 7: Determine U plq1 , ..., U plq kl from ErGp1ql s, ...,ErG pklq l s 8: Compute the inverses of U plq1 , ..., U plq kl
9: Compute the updating direction pl by (3) 10: Wl “Wl ´ α ¨ pl. 11: end for 12: end for

heading: 3.3 Identifying the Covariance Parameters of the Tensor Normal Distribution
text: By (2), Ui can be inferred from ErGpiqs up to a constant multiplier. However, different sets of multipliers can generate the same F , i.e. the same distribution. This is less of a problem if one
only cares about F . However, we need F´1 to compute the approximate natural gradient. That is, we first must choose a representation of F “ cpŨ1 b ¨ ¨ ¨ b Ũkq (see below), and then compute F´1 “ c´1ppŨ1 ` Iq´1 b ¨ ¨ ¨ b pŨk ` Iq´1q with a proper choice of ą 0, where I plays a damping role in the preconditioning matrix. In this case, different representations of F will lead to different F´1.
The statistics community has proposed various representations for Ũi’s. For example, [43] imposed that c “ 1 and the first element of Ũi to be one for i “ 1, ..., k ´ 1, whereas [10] imposed that trpŨiq “ 1 for i “ 1, ..., k. Although these representations have nice statistical properties, they are not ideal from the perspective of inverting the covariance for use in a NG method in optimization.
We now describe one way to determine Ũ1, ..., Ũk, and c from ErGp1qs, ...,ErGpkqs. In particular, we first set c “ 1, so that F´1 has a constant upper bound ´kI . We then require that trpŨiqdi is constant w.r.t i. In other words, the average of the eigenvalues of each of the Ũi’s is the same. This helps the Ũi’s have similar overall "magnitude" so that a suitable can be found that works for all dimensions. Moreover, this shares some similarity with how KFAC splits the overall damping term between KFAC matrices, although KFAC adjusts the damping values, whereas TNT adjusts the matrices. A bit of algebra gives the formula
Ũi “ ErGpiqs
ck´10 ś j‰i dj , (5)
where c0 “ ´ trpErGpiqsq ś
j dj
¯1{k .

heading: 3.4 Comparison with Shampoo and KFAC
text: Shampoo, proposed in [18], and later modified and extended in [3], is closely related to TNT. Both methods use a block-diagonal Kronecker-factored preconditioner based on second-order statistics of the gradient and are able to handle all sorts of tensors, and hence, can be applied to all sorts of deep neural network models, easily and seamlessly. The major differences between them are:
(i) The TN distribution cannot be directly applied to EF, which is used in Shampoo, because the empirical gradient does not have a zero mean; hence its covariance and second moment are different. It is also believed that EF does not capture as much valuable curvature information as Fisher [26]. (ii) Using the statistics ErGpiqs’s, TNT approximates the Fisher matrix as the covariance of the block-wise sampling-based gradients assuming that they are TN distributed. On the other hand, Shampoo computes 1{2k-th power of the statistics of each direction of the tensor-structured empirical gradient and forms a preconditioning matrix from the Kronecker product of them. It is unclear to us how to interpret statistically such a matrix other than by its connection to EF. We further note that Shampoo was developed as a Kronecker-factored approximation to the full-matrix version of AdaGrad [11], whereas TNT was developed as a NG method using a TN-distributed approximation to the Fisher matrix. (iii) TNT computes the updating direction using the inverse (i.e. power of ´1) of the Kronecker factors of the approximate Fisher matrix, whereas Shampoo uses the ´1{2k-th power1 of the Kronecker factors of the EF matrix.
Another method related to TNT is KFAC [34, 17], which, like TNT, uses Fisher as its preconditioning matrix. Their major differences are:
(i) KFAC develops its approximation based on the structure of the gradient and Fisher matrix for each type of layer. Admittedly, this could lead to better approximations. But it is relatively hard to implement (e.g. one need to store some intermediate variables to construct the KFAC matrices). Also, if new types of layers with different structures are considered, one needs to develop suitable Kronecker factorizations, i.e., KFAC matrices. On the contrary, TNT, like Shampoo, is a modelagnostic method, in the sense that, TNT can be directly applied as long as the shape of the tensor variables are specified.
1In [3], for autoencoder problems involving tensors of order 2, the power was set to ´α 2 , where α P r0, 1s was treated as a hyper-parameter which required tuning, and was set to α “ 1 after tuning.
(ii) Each block of TNT corresponds to a tensor variable whose shape needs to be specified, whereas each block of KFAC corresponds to all variables in a layer. For example, for a linear or convolutional layer, the KFAC block would correspond to the Fisher of both its weights and bias (and their correlation), whereas TNT would produce two blocks corresponding to the weights and bias, respectively.
In order to gain more insight into how well TNT approximates the Fisher matrix compared with other methods, we computed the cosine similarity between the direction produced by each method and that by a block Fisher method, where each block corresponded to one layer’s full Fisher matrix in the model (see Figure 1). For all methods shown in Figure 1, we always followed the trajectory produced by the block Fisher method. In our implementation of the block Fisher method, both the gradient and the block-Fisher matrices were estimated with a moving-average scheme, with the decay factors being 0.9. In all of the other methods compared to the block Fisher method, moving averages were also used, with the decay factors being 0.9, as described in Section D in the Appendix, to compute the relevant gradients and approximate block-Fisher matrices used by them, based on values computed at points generated by the block-Fisher method.
As shown in Figure 1, the cosine similarity for TNT is always around 0.7 to 0.8, which is similar to (and sometimes higher) than the structure-aware method KFAC, and always better than Shampoo. To provide more information, we also include SGD with momentum and Adam, whose similarity to the block Fisher direction is usually lower that of the second-order methods.

heading: 4 Convergence
text: In this section, we present results on the convergence of an idealized version of TNT that uses the actual covariance of Dθ, rather than a statistical estimate of it (see Algorithm 2 in the Appendix). In particular, our results show that Algorithm 2, with constant batch size and decreasing step size, converges to a stationary point under some mild assumptions. For simplicity, we assume that the model only contains one tensor variableW . However, our results can be easily extended to the case of multiple tensor variables. To start with, our proofs, which are delayed to Section B in the Appendix, require the following assumptions: Assumption 1. L : Rn Ñ R is continuously differentiable. Lpθq is lower bounded by a real number Llow for any θ P Rn. ∇L is globally Lipschitz continuous with Lipschitz constant L; namely for any θ, θ1 P Rn, }∇Lpθq ´∇Lpθ1q} ď L}θ ´ θ1}. Assumption 2. For any iteration t, we have
aq Eξt r∇lpθt, ξtqs “ ∇Lpθtq bq Eξt ” }∇lpθt, ξtq ´∇Lpθtq}2 ı ď σ2
where σ ą 0 is the noise level of the gradient estimation, and ξt, t “ 1, 2, . . . , are independent samples, and for a given t the random variable ξt is independent of tθjutj“1 Assumption 3. Let G :“ Dθ. For any θ P Rn, the norm of the Fisher matrix F “ Ex„Qx,y„prvecpGqvecpGqJs is bounded above.
Since F represents the curvature of the KL divergence of the model’s predictive distribution, Assumption 3 controls the change of predictive distribution when the model’s parameters change; hence,
it is a mild assumption for reasonable deep learning models. Essentially, we would like to prove that, if the Fisher matrix is upper bounded, our approximated Fisher (by TNT) is also upper bounded.
We now present two lemmas and our main theorem; see Section B in the Appendix for proofs.
Lemma 1. }Ex„Qx,y„prGpiqs} ď ´ 1 di śk i1“1 di1 ¯ }Ex„Qx,y„prvecpGqvecpGqJs}, @ i “ 1, . . . , k.
Lemma 2. Suppose Assumption 3 holds. Let FTNT :“ pŨ1 ` Iq b ¨ ¨ ¨ b pŨk ` Iq where Ũi’s are defined in (5). Then, the norm of FTNT is bounded both above and below. Theorem 1. Suppose that Assumptions 1, 2, and 3 hold for tθtu generated by Algorithm 2 with batch size mt “ m for all t. If we choose αt “ κLκ̄2 t ´β , with β P p0.5, 1q, then
1
N
N ÿ t“1 Etξju8j“1 ” }∇Lpθtq}2 ı ď 2L
` ML ´ Llow ˘ κ̄2
κ2 Nβ´1 ` σ
2
p1´ βqm ` N´β ´N´1 ˘ ,
where N denotes the iteration number and the constant ML ą 0 depends only on L. Moreover, for a given δ P p0, 1q, to guarantee that 1N řN t“1 Etξju8j“1 ” }∇Lpθtq}2 ı ă δ, N needs to be at most O ´ δ´ 1 1´β ¯ .

heading: 5 Implementation Details of TNT and Comparison on Complexity
text: Implementation Details of TNT. In practice, we compute G “ DW averaged over a minibatch of data at every iteration, and use the value of G piq to update a moving average estimate yGpiq of ErGpiqs. The extra work for these computations (as well as for updating the inverses of Ũi) compared with a stochastic gradient descent method is amortized by only performing them every T1 (and T2) iterations, which is also the approach used in KFAC and Shampoo, and does not seems to degrade the overall performance of the TNT algorithm. Moreover, we compute ErGpiqs using the whole dataset at the initialization point as a warm start, which is also done in our implementations of Shampoo and KFAC. See Algorithm 3 in the Appendix for the detailed implementation of TNT.
A Comparison on Memory and Per-iteration Time Complexity. To compare the memory requirements and per-iteration time complexities of different methods, we consider the case where we optimize one tensor variable of size d1 ˆ ¨ ¨ ¨ ˆ dk using minibatches of size m at every iteration. A plain SGD method requires Op
śk i“1 diq to store the model parameters and the gradient, whereas
its per-iteration time complexity is Opm śk i“1 diq. Table 1 lists the memory requirements and per-iteration time complexities in excess of that required by SGD for different methods.
Compared with a classic Newton-like method (e.g. BFGS), TNT (as well as Shampoo) reduces the memory requirement from Op
śk i“1 d 2 i q to Op řk i“1 d 2 i q, which is comparable to that of Adam-like
adaptive gradient methods. In fact, if the di’s are all equal to d and 3 ď k ăă d, the Kroneckerfactored TNT pre-conditioning matrix requires kd2 storage, which is less than that required by the diagonal pre-conditioners used by Adam-like methods. On the other hand, in terms of per-iteration time complexity, TNT (as well as Shampoo) only introduces a mild overhead for estimating the statistics ErGpiqs’s, inverting the pre-conditioning matrices, and computing the updating direction. Also, the first two of these operations can be amortized by only performing them every T1 and T2 iterations. Lastly, the extra work of Op 1T1m śk i“1 diq required by TNT relative to Shampoo is due to the extra backward pass needed to estimate the true Fisher, as opposed to the EF.
Moreover, although TNT and Shampoo both incur 1T2 řk i“1 d 3 i amortized time to invert the preconditioning matrices, the SVD operation in Shampoo can take much more time than the matrix inverse operation in TNT, especially when the matrix size is large2.
The per-iteration computational complexity of KFAC is more complicated because it depends on the type of the layer/variable. For linear layers, TNT and KFAC both uses two matrices, whose sizes are the number of input nodes and output nodes, respectively. For convolutional layers, TNT uses three matrices, whose sizes are the size of filter, number of input channels, and number of output channels, whereas KFAC uses two matrices whose sizes are the size of filter times number of input channels, and number of output channels. As a result, the first KFAC matrix requires much more memory. In general, the per-iteration complexity of KFAC is no less than that of TNT.

heading: 6 Experiments
text: In this section, we compare TNT with some state-of-the-art second-order (KFAC, Shampoo) and first-order (SGD with momentum, Adam) methods (see Section D.1 in the Appendix on how these methods were implemented). The Hessian-based K-BFGS method [15, 38] is another state-of-the-art Kronecker-factored second-order method for training deep learning models. Since our focus is on optimizers that use Fisher or empirical Fisher as the preconditioning matrix, we did not include K-BFGS in our comparison.
Our experiments were run on a machine with one V100 GPU and eight Xeon Gold 6248 CPUs using PyTorch [37]. Each algorithm was run using the best hyper-parameters, determined by an appropriate grid search (specified below), and 5 different random seeds. In Figures 2 and 3 the performance of each algorithm is plotted: the solid curves give results obtained by averaging the 5 different runs, and the shaded area depicts the ˘standard deviation range for these runs. Our code is available at https://github.com/renyiryry/tnt_neurips_2021.

heading: 6.1 Optimization: Autoencoder Problems
text: We first compared the optimization performance of each algorithm on two autoencoder problems [21] with datasets MNIST [27] and FACES3, which were also used in [32, 34, 5, 15] as benchmarks to compare different algorithms. For each algorithm, we conducted a grid search on the learning rate and damping value based on the criteria of minimal training loss. We set the Fisher matrix update frequency T1 “ 1 and inverse update frequency T2 “ 20 for all of the second-order methods. Details of our experiment settings are listed in Section D.2 in the Appendix. From Figure 2, it is clear that TNT outperformed SGD with momentum and Adam, both in terms of per-epoch progress and process time. Moreover, TNT performed (at least) as well as KFAC and Shampoo, with a particularly strong performance on the FACES dataset. We repeated these experiments using a grid search on more hyper-parameters, and obtained results (see Figure 6 in Sec D.5) that further support our observations based on Figure 2.
2In [3] it is shown that replacing the SVD operation by a coupled Schur-Newton method saves time for matrices of size greater than 1000 ˆ 1000. In our experiments, we used the coupled Newton method implementation of Shampoo.
3https://cs.nyu.edu/~roweis/data.html

heading: 6.2 Generalization: Convolutional Neural Networks
text: We then compared the generalization ability of each algorithm on two CNN models, namely, ResNet32 [19] (with CIFAR10 dataset [24]) and VGG16 [42] (with CIFAR100 dataset [24]). The first-order methods were run for 200 epochs during which the learning rate was decayed by a factor of 0.1 every 60 epochs, whereas the second-order methods were run for 100 epochs during which the learning rate was decayed by a factor of 0.1 every 40 epochs; (these settings are the same as in [45]). Moreover, as indicated in [29, 45], weight decay, different from the L2 regularization added to the loss function, helps improve generalization across different optimizers. Thus, for each algorithm, we conducted a grid search on the initial learning rate and the weight decay factor based on the criteria of maximal validation classification accuracy. The damping parameter was set to 1e-8 for Adam (following common practice), and 0.03 for KFAC4. For TNT and Shampoo, we set “ 0.01. We set T1 “ 10 and T2 “ 100 for the three second-order methods (same as in [45]). Details of our experiment settings and a further discussion of the choice of damping hyper-parameters can be found in Section D.3 in the Appendix.
The results in Figure 3 indicate that, with a proper learning rate and weight decay factor, second-order methods and Adam exhibit roughly the same generalization performance as SGD with momentum, which corroborate the results in [29, 45]. In particular, TNT has a similar (and sometimes better) generalization performance than the other methods. For example, comparing TNT with SGD-m, TNT (SGD-m) achieves 93.08% (93.06%) validation accuracy with ResNet32 on CIFAR10 and 73.33% (73.43%) validation accuracy with VGG16 on CIFAR-100, after 100 (200) epochs (see Table 3 in the Appendix for the accuracy achieved by the other algorithms). Moreover, in terms of process time, TNT is roughly twice (equally) as fast as SGD with momentum on ResNet32/CIFAR10 in Figure 3a (on VGG16 on CIFAR-100 in Figure 3b). This illustrates the fact that TNT usually requires only moderately more computational effort per-iteration but fewer iterations to converge than first-order methods. Also, as shown on the VGG16 model, KFAC seems to be much slower than TNT and Shampoo on larger models. This is because the most recent version of KFAC, which we implemented, uses SVD (i.e., eigenvalue decomposition) to compute inverse matrices (see Section D.1.2 in the Appendix for a discussion of this). In contrast, TNT does not need to use SVD, and the most recent version of Shampoo replaces SVD with a coupled Newton method in [3].
We also compared TNT with a variant of it that uses the empirical rather than the true Fisher as the preconditioning matrix. The results of this comparison, which are presented in Section D.4 in the Appendix, suggest that it is preferable to use Fisher rather than empirical Fisher as pre-conditioning matrices in TNT.
4The value of 0.03 is suggested in https://github.com/alecwangcq/KFAC-Pytorch, a github repo by the authors of [45].

heading: 7 Conclusion and Further Discussions
text: In this paper, we proposed a new second-order method, and in particular, an approximate natural gradient method TNT, for training deep learning models. By approximating the Fisher matrix using the structure imposed by the tensor normal distribution, TNT only requires mild memory and computational overhead compared with first-order methods. Our experiments on various deep learning models and datasets, demonstrate that TNT provides comparable and sometimes better results than the state-of-the-art (SOTA) methods, both from the optimization and generalization perspectives.
Due to space and computational resource constraints, we did not run experiments on even larger models such as ImageNet and advanced models for NLP tasks. However, the results in this paper already show very strong evidence of the potential of the TNT method. We also did not explore extending our method to a distributed setting, which has been shown to be a promising direction for second-order methods such as KFAC and Shampoo [4, 3]. Since TNT already performs very well on a single machine, we expect that it will continue to do so in a distributed setting. These issues will be addressed in future research. We did not compare TNT with the SOTA Kronecker-based quasi-Newton methods [15, 38], since they are not as closely related to TNT as are Shampoo and KFAC. Their performance relative to TNT can be inferred from the comparisons here combined with those reported in [15, 38, 3].
As a final note5, the preconditioning matrices of TNT (as well as those of Shampoo) are derived from the specific shape of the (tensor) parameters of the particular deep learning model that is being trained. One can, of course, reshape these parameters, e.g., by flattening the tensors into vectors, which gives rise to very different preconditioning matrices.
The method proposed in this paper can be applied to any deep learning or machine learning model. If the model and/or data has a flawed design or contains bias, this could potentially have negative societal impacts. However, this possibility is beyond the scope of the work presented in this paper.
5We thank the program chair for pointing this out.
Acknowledgments and Disclosure of Funding
We would like to thank the anonymous reviewers for their very helpful comments and suggestions.
The research efforts of D. Goldfarb and Y. Ren on this paper were supported in part by NSF Grant IIS-1838061.
We acknowledge computing resources from Columbia University’s Shared Research Computing Facility project, which is supported by NIH Research Facility Improvement Grant 1G20RR03089301, and associated funds from the New York State Empire State Development, Division of Science Technology and Innovation (NYSTAR) Contract C090171, both awarded April 15, 2010.

Review summaries:
Paper proposes a tractable natural gradient approximation. It uses the shampoo approximation algebra combined with exponents of -`1 as its approximating the Fisher. Empirical experiments were carried out in two tasks: 1) autoencoder as well 2) convolutional nets.

This paper presents TNT, a variant of Shampoo, but based on natural gradients. The key difference is that while Shampoo uses the gradient itself to compute the preconditioning matrices, TNT uses the gradients based on the predicted distribution to compute preconditioners - this is similar to KFAC. The paper compares TNT with Shampoo and KFAC, and the main differences are: 1. Shampoo preconditioners correspond to Empirical Fisher, whereas TNT approximates the true Fisher, at the cost of an extra backprop step. There is some evidence that the true Fisher captures curvature information better. 2. Shampoo preconditions the gradient as G x_1 U1^{-1/2k} x_2 U2^{-1/2k} ..., whereas TNT preconditions it as G x_1 U1^{-1} x_2 U2^{-1} ...  (modulo some damping).  Thus TNT needs only inverses of matrices, which is an easier operation than finding matrix roots. 3. KFAC approximates each layer's Fisher matrix as a Kronecker product using a unique algorithm, thus KFAC needs to be re-derived for each layer type. TNT like Shampoo works for all layer types.  The paper includes a theoretical analysis of convergence of TNT, as well as some experiments showing that it is competitive on some small datasets.

The method proposes a novel approximate method, for approximating the Fisher matrix of a neural network, which to be used as a preconditioner for second-order optimization during training. The method, termed Tensor Normal Training (TNT) approximates the diagonal blocks of the Fisher via a Tensor Normal distribution. The idea of treating second order methods in this fashion (as having a Gaussian approximation to the parameters) has other connections to Bayesian learning methods and interpretations of Natural Gradient, e.g. like [1]. The key part of the method is to use a Kronecker Factored gaussian distribution over each rank of the actual variable tensor (e.g. Matrix Normal for a matrix weights). The authors derive the exact computation needed for estimating the covariance matrices show the convergence properties of the method. They relate the method to other well known approximate curvature matrix methods - Shampoo and KFAC - and compare on two autoencoder tasks and two classification tasks, where the results suggest that TNT practically outperforms the other methods, with better training and comparable test performance.    [1] - Khan, M.E. and Rue, H., 2021. The Bayesian Learning Rule. arXiv preprint arXiv:2107.04562.