Tensor Normal Training for Deep Learning Models

1 Introduction . 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 ) . 2 Preliminaries . 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 CAN NOT 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 . 3 Tensor-Normal Training . 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 . 3.1 Block Diagonal Approximation . 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 .

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