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