Tensor Normal Training for Deep Learning Models
Authors: Yi Ren
Donald Goldfarb
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Optimization for deep networks

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, Shampoo, and K-BFGS, have been proposed and shown to be effective. Based on the so-called tensor normal (TN) distribution, 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.

Tensor Normal Training for Deep Learning Models
Yi Ren, Donald Goldfarb
Department of Industrial Engineering and Operations Research
Columbia University
New York, NY 10027
{yr2322, goldfarr))@columiaa edu
Abstract
Despite the predominant use of first-order methods for training deep learning mod- els, second-order methods, and in particular, natural gradient methods, remain interest because of their potential for accelerating training through the use of cur- vature 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 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 TN distribution, ensures that its covariance has Kronecker separable structure, which leads to 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-a first-order methods, and compara- ble optimization performance to the sttt--ffthee-at second-order methods KFAC and Shampoo. Moreover, TNT demonstrated its ability to generalize as well as first-order methods, while using fewer epochs.
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 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 's method, BFGS [6, 13, 14, 41], and limited-memory BFGS (L-BFGS) [28], without modification, are not practical in 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 memory and computing resources.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Recently, there has been considerable udvancement in the development of second-orde methods that are suitable for deep learning models with huge numbers of parameters. These methods usually ap- proach 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 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- -oder 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- cronecker-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.c. 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 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-norma 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 first attempt to do this and demonstrate its great potential.
Our Contributions. this paper, we propose a brand new approximate natural gradient (NG) method, Gensor-Normal Training (TNT), that makes use of the tensor normal distribution to approxi- mate 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 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 's and no greater than KFAC's, while its per-iteration computational needs no greater than Shampoo's '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, 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 model, by minimizing L(0) = 1. N Σ I(,, fo(x:)), where ((xe, ye) N denotes a given dataset (x; being the input to the model and yı being the target), fa()); denotes the output of the model when 24 is provided as the input, and l denotes loss function
2

---

(e.g. least-squares loss for regression and cross entropy loss for classification) that measures the discrepancy between the model output fo();) and the target y;-
Natural Gradient Method and the Fisher Matrix. In a first-order method, say SGD, the updating direction is always derived from estimate the gradient direction VoL(O). In natural gradient (NG) method [1], however, the Fisher matrix is used as pre-conditioning matrix that is applied to the gradient direction. As shown in [34], the Fisher matrix defined as
Ve log plyle, 0) (Vo logplul 0)) ")'
E 1r
(I)
where Qr is the data distribution of and p(-|2, 0) is the density function of the conditional distribu- tion defined by the model with given input
In many cases, such when p is associated with a Gaussian distribution and the loss function I measures least-squares loss, or when p is associated with multinomial distribution and I is cross- entropy loss, log p is equivalent / (see e.g. [33, 34]). Hence, if DO denotes the gradient of w.r.t. O for a given æ and y, we have that F = pl [D Consequently, one can sample x from Qx and perform forward pass of the model, then sample y from p(-lar 0). and perform backward pass to compute D0, and then use D to estimate F. We call D0 sampling- -based gradient, as opposed to the empirical gradient V (y) fo(2;)) where (Iz, yi) is one instance from the dataset.
It is worth noting that the first moment of DO is zero. To see this, note that, with given
E Ey-p Vo log 0)1
Ve logplu/z 0)l(vle, dy 0)dl
-Vo (f (plllz,0)d = Vol O.
Hence, 1Fx-Q..--ppp [D]] Ez- (Ry-- Vo bogg(yya 0)] x) = 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 tensor- normal distribution [31, 10]:
Definition 1. An arbitrary tensor G E TRd.xx.xx. x-xxxxx is said to follow a tensor normal (TN) distribution with mean parameter M E d xxxx and covariance parameters U1 didix ... Ug E R x if and only if vec(G) ~ Normallved(MM)
@ ... & Uk).
In the above definition, the vec operation refers to the vectorization of a tensor, whose formal definition be found in Sec A in the Appendix. Note that matrix-normal distribution be viewed as a special case of TN distribution, where k = 2. Compared with a regular normal distribution, whose covariance matrix has I de elements, the covariance of k- k-way tensor-normal distribution stored in i smaller matrices with a total number of elements equal to Σ dZ.
To estimate the covariance submatrices Ui...., Uk. the following property (e.g., see [10]) is used:
E G = U. I [Tr(U,),
J#L
where G(+) mGt, dRd denotes the contraction of G with itself along all but the ith dimension and mat, refers to matricization of a tensor (see Section A for the formal definitions). By (2), we can sample G to obtain estimates of the G(1)' ss and hence, estimates of the U.'.. The complexity of computing (+)) d, I[ dj. which is also far less than the complexity of computing vec( (e) needed to estimate the covariance of regular normal distribution.
3 Tensor-Norral 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 censor-normal distribution.
3

---

3.1 Block Diagonal Approximation
In this paper, we consider the case where the parameters of the model consists of multiple tensor variables Wi...., WL. i.e. e (vec(W1) ..... vec(WL) 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 Wi (1 1,-.. ..... I.) come from the weights or biases of layer, whether be feed-forward, convolutional, recurrent, or batch normalization layer, etc. Note that the index I of W, refers to the index of 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 block diagonal. To be more specific, assume that each block corresponds the covariance of a tensor variable in the model. Hence, the approximate Fisher matrix is:
}] = diag/=1 (ar(vve((DW))) The remaining question is how should one approximate Var(vec(DWi)) for 1 1, 1,... L.
F diag (Ex-Qu.
3.2 Computing the Approximate Natural Gradient Direction by TNT
We consider a tensor variable W E dd x.. in the model and assume that G := DW Rd x.,, follows TN distribution with zero mean and covariance parameters Uj,... U3 where U, E Rdxxd, Thus, the Fisher matrix corresponding to W is Fw = Ez-Q... 8.ppp Var Var(vec(G)) = U1 @ ... X 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 Vw L E TR x-xx denote the gradient of L w.r.t. W. The approximate NG updating direction for W is computed as
Fiw vec( Vw = (UT @ & (W = vec (V x UT X2.... ... X U (3)
where (i = 1,... k) denotes 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 See A the Appendix for a proof):
Proposition 1. Let G E Rd xindd and U, Rd xd. for -,..., k. Then, we have
(.-)U.)) vea G) = vec G x U1 x2 U2 .. x Uv).
(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 == 2, do
2: Sample mini-batch M, of size
3: Perform a forward- backward pass over Mj to compute the mini-batch gradient
4: Perform another backward pass over M, with y sampled from the predictive distribution to compute G. DW, (I 1,..., L) averaged across M,
5: for 1, do
Estimate B[G (i I,,,, ki) from G₁
6:
Determine (1) U.., from B[G (1) --., Bio
7:
Compute the inverses of (()) -..,0 ((()
8:
9:
Compute the updating direction py by (3)
Wi W, Pt.
11: end for
end for
3.3 Identifying the Covariance Parameters of the Tensor Normal Distribution
By (2), U, can be inferred from E[G(1] up to constant multiplier. However, different sets of multipliers can generate the same F, i.c. 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 = c(U, @ ... & U&) (see below), and then compute F-1 = + e @ @ (Üx + + with proper choice of e > O, where el plays damping role in the preconditioning matrix. this case, different representations F will lead to different F-1.
The statistics community has proposed various representations for U, For example, [43] imposed that e and the first element of U; to be one for i I,..., k 1, whereas [10] imposed that tr( tr(ū;) for i I,.... k. Although these representations have nice statistical properties, they not ideal from the perspective of inverting the covariance for use in a NG method in optimization.
We now describe one way to determine U1, UR. and e from first set so that F-1 has constant upper bound < We then require that m((:) d: is constant w.r.t i. In other words, the average of the cigenvalues of each of the U, is the same. This helps the U,'s have similar overall "magnitude" so that 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
In particular, we
E G|
U,
(5)
11iped3
[1(E(e((1)) I1, dy
1/k
where co
3.4
Comparison with Shampoo and KFAC
Shampoo, proposed in [18], and later modified and extended in [3], closely related TNT. Both methods use a (lock-diagonal Kronecker-factee the gradient and are able to handle all sorts of tensors, and hence, can be applied to sorts of deep neural network models, easily and seamlessly. The major differences between them are:
preconditioner based on second-order statistics of
(i) The TN distribution cannot be directly applied to EF, which is used in Shampoo, because the empirical gradient does not have zero mean; hence its covariance and second moment are different. is also believed that EF does not capture as much valuable curvature information as Fisher [26]. (ii) Using the statisties ("' 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. 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 NG method using a TN-distributed approximation to the Fisher matrix.
(ii) 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 power' 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 Kronecker factorizations, i.e., KFAC matrices. On the contrary, TNT, like Shampoo, is a model- agnostic method, in the sense that, TNT can be directly applied as long the shape of the tensor variables are specified.
considered, one needs to develop suitable
In [3], for autoencoder problems involving tensors of order 2, the power was set to - where [0,1] treated a hyper- pppre-prraeeter which required tuning, and set = 1 after tuning.
5

---

0.9
0.9
cosine sinbees
0.8
0.8
0.8
TNT
KFAC
0.7
0.7
Shampoo
0.
0.6
1/4 M./
0.6
500pm
0.5
0.5
0.4
500 Rerattoo
1000
500
1000
505
1000
500
1900
Iteraiio
Rteraton
Iteratton
Figure I:
Cosine similarity between the directions produced by the methods shown in the legend and that of a block Fisher method. The algorithms were run on a 16 x 16 down-scaled MNIST [27] dataset and small feed-forward NN with layer widths 236-20-20-20-2000 in [34], we only show the middle four layers.
described in [34]. As
(ii) Each block of TNT corresponds to 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 convo- lutional layer, the KFAC block would correspond 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 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 mhoving-averagge 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, described in Section D in the Appendix, to compute the relevant gradients and approximate block-Fishe matrices used by them, based on values computed at points generated by the block-Fisher method.
As shown in Figure I, 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.
4 Convergence
In this section, present results on the convergence of an idealized version of TNT that uses the actual covariance of DD, rather than 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 stationary point under some mild assumptions. For simplicity, we assume that the model only contains one tensor variable W. However, our results can 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 IR" + R is continuously differentiable. (()) is lower bounded by real number Clow for any A R VL globally Lipschitz continuous with Lipschitz constant L; namely for any 0, 0' E ER. | (()) v [((')] & L| 8'1
Assumption 2. For any iteration t, we have
a) Eg, [VI(O,. [.]] VL(,))
b) Eg. (
where o > 0 is the noise level of the gradient estimation, and Er, 1,2, are independent samples, for given the random variable & is independent of (0,} Assumption 3. Let G Eg (ee-Qu..--p(ve(((( vecc T] is bounded above.
'DO. For any 0 E IR.", the norm of the Fisher matrix F
Since F represents the curvature of the KI. divergence of the model's predictive distribution, As- sumption 3 controls the change of predictive distribution when the model's parameters change; hence,

---

Table 1: Memory and per-iteration time complexity beyond that required by SGD
Name
Memory
Time (per-iteration)
TNT
O k --1 di)
((( m + > d:) di + T2 1
Shampoo
01Σ;=1 * d
0 [T= de + (T2 Ei=1 di- using SVD))
E
Adam-like
O 1= d,)
o dz)
Newton-like
0 1- dc)
depends on specific algorithm
is 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 in the Appendix for proofs.
Lemma 1. [Er-Q..u-PIG [| & (a
|E (
11I, V 1....,k.
Lemma 2. Suppose Assumption 3 holds. Let FINT == (Ū, cT) @ ... & (Ü, + cT) where U, are defined in (5). Then, the norm of PINT is bounded both above and below.
Theorem 1. Suppose that Assumptions 1, 2, and 3 hold for {0e) generated by Algorithm 2 with batch size my = for all t. If we choose Oi,- = LE 1 1-, with 3 E (0.5, 1), then
N 16, [IVC(0)11] 2L (Mc Llow) R²
1
No-1 + (1 (1- B)m ON-8- !
o²
where N denotes the iteration number and the constant Mc 0 depends only L. Moreover, for 6, N needs to be at most given 5 E (0, 1), to guarantee that N 12, [
Implementation Details of TNT and Comparison on Complexity
Implementation Details of TNT. In practice, we compute G = DW averaged over minibatch
of data at every iteration, and use the value of (0)) update moving average estimate G of E[G()] The extra work for these computations (as well for updating the inverses of U.) compared
with stochastic gradient descent method is amortized by only performing them every T; (and T2) iterations, which also the approach used in KFAC and Shampoo, and does not seems to degrade the overall performance of the TNT algorithm. Moreover, we compute E[G(() using the whole dataset at the initialization point as warm start, which is also done in our implementations of Shampoo and KFAC. See Algorithm in the Appendix for the detailed implementation of TNT.
Comparison on Memory and Per-iteration Time Complexity. To compare the memory require- ments and per-iteration time complexities of different methods, consider the case where optimize one tensor variable of size d2 x ... x de using minibatches of size m at every iteration. A plain SGD method requires O d,) to store the model parameters and the gradient, whereas its per-iteration time complexity is (m IT d.). Table lists the memory requirements and per-iteration time complexities in excess of that required by SGD for different methods.
Compared with classic Newton-like method (e.g. BFGS), TNT (as well Shampoo) reduces the dd) to O df), which is comparable to that of Adam-like adaptive gradient methods. In fact, if the d;'s are all equal to d and 3 & k << d, the Kronecker- factored TNT prndinniiion matrix requires kd storage, which is less than that required by the diagonal pre-conditioners used by A Adam-like methods. On the other hand, in terms of per- per-iteration time complexity, TNT (as well as Shampoo) only introduces a mild overhead for estimating the statistics (i)'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 Ti and T2 iterations. Lastly, the extra work of 0 I - d,) required by TNT relative to Shampoo is due to the extra backward pass needed to estimate the true Fisher, opposed to the EF.
memory requirement from O l
7

---

Moreover, although TNT and Shampoo both incur 1 E of amortized time to invert the pre- conditioning matrices, the SVD operation in Shampoo can take much more time than the matrix inverse operation TNT, especially when the matrix size large".
The per-iteration computational complexity of KFAC is more complicated because depends on the type of the layer/variable. For linear layers, TNT and KFAC both uses two matrices, whose sizes the number of imput nodes and output nodes, respectively. For convolutional layers, TNT uses three matrices, whose sizes 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 result, the first KFAC matrix requires much more memory. In general, the per-iteration complexity of KFAC is no less than that of TNT.
6 Experiments
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-e Kronecker-factore second-order method for training deep learning models. Since our focus is on optimizers that Fisher empirical Fisher the preconditioning matrix, did not include K-BFGS in our comparison.
machine with one V100 GPU and eight Xeon Gold 6248 CPUs using PyTorch [37]. Each algorithm was run using the best 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 different runs, and the shaded area depicts the standard deviation range for these runs. Our code is available at https: /gi thub. com/remo/////eeeee com/renyrryry wattttneurrs__000
Our experiments
determined by an appropriate
6.1 Optimization: Autoencoder Problems
TVT
ThT
Adam
100
100
Acem
102
420
100
120
400 Sime (wvors))
500
250
1000
1750
100
Dime Cecoms
2000
a) MNIST autoencoder
FACES autoencoder
Figure 2: Optimization performance of TNT, KFAC, Shampoo, Adam, and SGD-m on two autoen-
coder problems
We first compared the optimization performance of each algorithm on two autoencoder problems [21] with datasets MNIST [27] and FACES*, which were also used [32, 34, 5, 15] as benchmarks to compare different algorithms. For each algorithm, we conducted grid search on the learning rate and damping value based on the criteria of minimal training loss. We set the Fisher matrix update frequency Ti 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 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 the FACES dataset. We repeated these experiments using grid search on more hyper-parameterr, parameters, and obtained results (see Figure See D.S) that further support our observations based on Figure 2.
In [3] it is shown that replacing the SVD operation by coupled NehurnNewton method saves time for matrices of size greater than 1000 1000. In our experiments, we used the coupled Newton method implementation Shampoo. "https://ee nyu. wwtp://ooiis/taa..ttet...
8

---

6.2 Generalization: Convolutional Neural Networks
20P
20-1
20-1
10-1
350
360
202
2rccede DIN reecond
10020 TVT
307
307
TMT
Starcce 5:5-mm
500mm
30
10"11
103
pricess
a) CIFAR-10, 10, ResNet-32
b) C00, VGG16
Figure 3: Generalization ability of TNT, KFAC, Shampoo, Adam, and SGD-m on two CNN models. Upper row depicts the training loss whereas lower row depicts the validation classification error.
We then compared the generalization ability of each algorithm on two CNN models, namely, ResNet32 [19] (with CIFARIO dataset [24]) and VGG16 [42] (with CIFARIOU 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 le-8 for Adam (following common practice), and 0.03 for KFAC For TNT and Shampoo, we
0.01. We set T2 10 and T2 = 100 for the three second-ords -order methods (same as in [45]). Details of our experiment settings and a further discussion of the choice of damping D.3 in the Appendix.
can be found in Section
The results in Figure 3 indicate that, with 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 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 CIFAR 10 and 73.33% (73.43%) validation accuracy with VGG16 on CIFAR-100, after 100 (200) epochs (see Table in the Appendix for the accuracy achieved by the other algorithms). Moreover, terms of process time, TNT is roughly twice (equally) as fast as SGD with momentum on 2//NeC32/CIFARI 10 in Figure 3a (on VGG16 on CIFAR-I00 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.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.
'The value of 0.03 is suggested in htttps://gite com/ Wwleevaagcg/KAA--PPeeeoeee github repo by the authors of [45].

---

7 Conclusion and Further Discussions
In this paper, we proposed a new second-order method, and 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 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 distributed setting, which has been shown be a promising direction for second-order methods such as KFAC and Shampoo [4, 3]. Since TNT already performs very well on single machine, we expect that it will continue do so in 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 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 final note', 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 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 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.
'e thank the program chair for pointing this out.
10

---

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-1838011.
We acknowledge computing resources from Columbia University's Shared Research Computing Facility project, which is supported by NIH Research Facility Improvement Grant 1G20RR030893. 01, 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.
References
[1] S.-I. Amari, H. Park, and K. Fukumizu. Adaptive method of realizing natural gradient learning for multilayer perceptrons. Neural computation, 12(6):1399-1409 2000.
[2] E. Amid, R. Anil, and M. K. Warmuth. Locoprop: Enhancing backprop via local loss optimiza- tion. arXiv preprint arXiv:2106.06 199, 2021.
[3] R. Anil, V. Gupta, T. Koren, K. Regan, and Y. Singer. Scalable second order optimization for
deep learning, arXiv preprint arXiv:2002.09016 2021.
[4] J. Ba, R. Grosse, and J. Martens. Distributed second-order optimization using Kronecker-
factored approximations. 2016.
[5] A. Botev, H. Ritter, and D. Barber. Practical Gauss-Newton optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning-Von Volume 70, pages 557-565. JMLR. org, 2017.
[6] C. G. Broyden. The convergence of a class of double-rank minimization algorithms 1. general considerations. IMA Journal of Applied Mathematics, 6(1):76 90, 1970.
[7] R. H. Byrd, S. Hansen, J. Nocedal, and Y. Singer. A stochastic quasi-Newton method for large-scale optimization. SIAM Journal on Optimization, 26(2): 1008 10000001111111000000111101111111 2016.
[8] D. Choi, C. Shallue, Z. Nado, J. Lee, C. J. Maddison, and G. E. Dahl. On empirical comparisons of optimizers for deep learning, arXiv preprint arXiv: 2019.
[9] A. P. Dawid. Some matrix- distribution theory: notational considerations and a Bayesian application. Biometrika, 68(1):265-274, 1981.
[10] B. S. Dees and D. P. Mandic. A statistically identifiable model for tensor-valued Gaussian random variables. arXiv preprint arXiv:1911.02911 2019.
[11] Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 120u1)2121-2215, 2011.
[12] Dutilleul. The mle algorithm for the matrix normal distribution. Journal of statistical computation and simulation, 64(2):105-123, 1999.
[13] R. Fletcher. A new approach to variable metric algorithms. The computer journal, 13(3): 317-322, 1970.
[14] D. Goldfarb. A family of variable-metric methods derived by variational means. Mathematics of computation, 24(109):23-26, 1970.
[15] D. Goldfarb, Y. Ren, and A. Bahamou. Practical quasi-Newton methods for training deep neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, ed- itors, Advances in Neural Information Processing Systems, volume 33, pages 2386-2396. Curran Associates, Inc., 2020. URL //tp:ee///rrcceeinggs 0tlle/192f044474fffffea 4ew.aa14e49551151955555555
cc/ wwwpcc//pppppeeeeeeeeeeeeeeeeeee /2020/
.df.
[16] R. Gower, D. Goldfarb, and P. Richtárik. Stochastic block BFGS: Squcezing more curvature out of data. International Conference on Machine Learning, pages 1869-1878, 2016.
[17] R. Grosse and J. Martens. A Kronecker-factored approximate Fisher matrix for convolution layers. In International Conference on Machine Learning, pages 573-582, 2016.
[18] V. Gupta, Koren, and Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference Machine Learning, pages 1842-1850 PMLR, 2018.

---

[19] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-
778, 2016.
[20] G. Hinton, N. Srivastava, and K. Swersky. Neural networks for machine learning lecture 6a overview of mini--batc gradient descent. Cited on, 14(8), 2012.
[21] G. E. Hinton and R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. science, 313(5786):504-55577 507, 2006.
[22] M. E. Khan and H. Rue. The Bayesian learning rule. arXiv preprint arXiv:2107.04562 2021.
[23] D. Kingma and J. Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2014.
[24] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[25] A. Krizhevsky, I. Sutskever, and G. Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25:1097-1105, 2012.
[26] F. Kunstner, P. Hennig, and I.. Balles. Limitations of the empirical Fisher approximation for natural gradient descent. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bucc E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL attps://proceediiee wwapiiinnnnnnuuiipgeeee file/ 4445558997954400622222 18 Paper.ddf
[27] Y. LeCun, C. Cortes, and Burges. MNIST handwritten digit database. ATT Labs [Online]. Available: http://yann. www.l//ww....cc///.ee/e
2, 2010.
[28] D. C. Liu and J. Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical programmiing 45(1-3):503-52) 1989.
[29] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Con- ference on Learning Representations, 2019. URL https://openrevier formmiid Bkg6RicqY7.
[30] H. Lu, K. N. Plataniotis, and A. Venetsanopoulos. Multilinear subspace learning: dimensionality reduction of multidimensional data. CRC press, 2013.
[31] A. M. Manceur and P. Dutilleul. Maximum likelihood estimation for the tensor normal dis- tribution: Algorithm, minimum sample size, and empirical bias and dispersion. Journal of Computational and Applied Mathematics, 239:37-49, 2013.
[32] Martens. Deep learning via Hessian-free optimization. ICML, volume 27, pages 735-742, 2010.
[33] J. Martens. New insights and perspectives on the natural gradient method. arXiv preprint 2014.
[34] J. Martens and R. Grosse. Optimizing neural networks with Kronecker- factored approximate curvature. In International conference on machine learning, pages 2408-24882 2287-241, 2015.
[35] J. Martens, J. Ba, and M. Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018. URL attps://openrevic wwtp://e//tommmdtttieiieeoa
[36] A. Mokhtari and A. Ribeiro. Res: Regularized stochastic BFGS algorithm. IEEE Transactions on Signal Processing, 62(23):6089-6104, 2014.
[37] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang. Z. DeVito, M. Raison, A. Tejani, S. CChilamkurhyy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Gamnett, editors, Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL http:: http:://ppppppeeeeerpppppppepepe neurips www.eerss..cpppapeeeeee cc/paper/ 9015-pytorch. an- hrppppati.e..taaaaaaaa.roooroeor wwwpaarnnnnnnnnanarrrreaa pdf.
[38] Y. Ren and D. Goldfarb. Kromecker-factored quasi-Newton methods for convolutional neural networks. arXiv preprint arXiv:2102.06737, 2021.
[39] H. Robbins and S. Monro. stochastic approximation method. The annals of matheraatiaal statistics, pages 400-407, 1951.
12

---

[40] R. M. Schmidt, F. Schneider, and P. Hennig. Descending through a crowded valley- benchmarking deep learning optimizers. In International Conference on Machine Learning. pages 9367-9376. PMLR, 2021.
[41] D. F. Shanno. Conditioning of quasi-Newton methods for function minimization. Mathematics of computation, 24(1 1 1):647-656, 1970.
[42] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1555 2014.
[43] M. Singull, M. R. Ahmad, and D. von Rosen. More on the Kronecker structured covariance matrix. Communications in Statisficx-Theory and Methods, 41(13-14)2512--522) 2012.
[44] X. Wang, S. Ma, D. Goldfarb, and W. Liu. Stochastic quasi-Newton methods for nonconvex stochastic optimization. SIAM Journal Optimization, 27(2)927-9556, 2017.
[45] G. Zhang, C. Wang. B. Xu, and R. Grosse. Three mechanisms of weight decay regularization.
In International Conference on Learning Representations, 2019. URL /opprr/vieeerr...... ht/f//oummiddBllle 3Rct7.
13


Reviews:
[{"review_id": "r3txLXFlTpd", "reviewer": "Reviewer_kZ6w", "summary": "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.  \n\n[1] - Khan, M.E. and Rue, H., 2021. The Bayesian Learning Rule. arXiv preprint arXiv:2107.04562.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is written really well, very clear simple idea which shows some promise. I liked the idea a lot and think it is very clever. The significance of the method is less clear, as although the authors show benefit on several problems, these are achieved under somewhat restricted conditions - small batch size, fixed learning rate and damping throughout optimization, in contrast to for instance the full fledged KFAC algorithm from [33]. Nevertheless, I understand including all that can be quite difficult to do, hence I see this as a minor point, but hope the authors to consider making this comparison in the near future. The actual method, I find very interesting, as for dense layers it is almost exactly the same as KFAC, but not quite. In fact, I would urge the authors to have a plot/figure, where for a single layer where they can compute the Full fisher exactly, what are the difference between the KFAC and TNT approximation (specifically of the eigen directions) e.g. more than the just the cosine similarity. I do not have a clear intuition on this, but I think it would be very useful for understanding why TNT might be a better approximation. \n\nI did not fully check the maths in the appendix regarding the convergence rate theorems and lemmas, but they look correct at a glance.\n\nA few other concerns on the experimental section:\n\n1. The autoencoder examples seem to be with ReLU where previous works have always done this with hyperbolic tangent activation functions. I would urge the authors to rerun the experiments in this way, as one of the key points of these experiments is that they are very difficult from optimization point of view and the activation matters. \n\n2. Could you also include an experiment where you use TNT for the joint weights+biases parameters, where you concatenate the biases to the weights to form a new matrix.\n\n3. For the convolutions could you include a TNT where it treats the parameters the same way as KFAC - as a single matrix. This would server as making more clear to the reader whether the benefit of TNT comes from better factorization of the approximate matrix, or from better estimation of the factors themselves, or from both. \n\nNit pick: Note that with the appropriate normalization of the two factors in KFAC (e.g. rescaling them such that they have the same norms), the way it computes a separate damping for each factor, can be reduced indeed to being the same for both. Would be nice if you can include this in the appendix, to make the point more clear to a curious reader. \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "peV51MVacsb", "reviewer": "Reviewer_5D53", "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:\n1. 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.\n2. 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.\n3. 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.\n\nThe paper includes a theoretical analysis of convergence of TNT, as well as some experiments showing that it is competitive on some small datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The algorithm should be of considerable interest to the ML community, as it seems to combine desirable attributes of both KFAC and Shampoo.\n\nIn section 3.3, it seems that the only reason one has to worry about the covariance parameters is the presence of damping. If damping is small (i.e. can be ignored), then can one can simplify the expression G x_1 U1^{-1} x_2 U2^{-1} ... to G x_1 G^(1)^{-1} x_2 G^(2)^{-1} ... * Tr(G^(1))^{k-1} ?\n\nThe experimental section could use a lot of improvements. It would be preferable to get a stronger CIFAR-10 baseline - SGD-m can reach much better accuracy than the 93.06 considered in this paper. The other datasets are also rather small, and it seems that the Shampoo authors have already addressed the performance issues in [3] --- for some reason this paper used SVD to compute matrix roots (including a data transfer to the CPU), whereas coupled Newton is much faster and can be executed on GPUs, so KFAC and Shampoo look artificially worse than TNT. Also this paper's results on auto-encoders are not nearly as good as [3] for these other optimizers. \n\nMinor nit - the citations in the Appendix do not match with the main paper.\n\nUpdate after author comments: The authors have addressed many of the comments about the experimental section, and will update this section in the paper. Correspondingly the overall assessment has been updated.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "None", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "40Yhxryc-SG", "reviewer": "Reviewer_KJXA", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The approach is clever,  and is also a straightforward extension of Shampoo to use sampled gradients from model's predictive distribution for computing statistics for the preconditioner. Another difference is that of exponents - as Shampoo (Gupta et al) approximates full matrix adagrad and requires the exponents sum to -1/2 for the loewner order to hold, however, as authors note its tuned over in Anil et al 2020. \n \nThe main concern with the work (and reason for score) is in the experiment section on tuning, and implementation of Shampoo. The author's use SVD to compute inverse pth root which is known to be expensive, and Anil et al 2020 already have shown a simple iterate method (coupled newton iterations) that runs the inverse-pth root much faster (~4x) - quite as fast as inverse operation on GPUs/TPUs.  Moreover, switching to eigh will make the implementation faster than SVD as there is no need to compute U, and V^T separately, as the matrices are symmetric, and postive (semi) definite. Thus main concern, is that experimental section results are unfair to the baseline being compared.  Code for inverse pth root ( https://github.com/google-research/google-research/blob/master/scalable_shampoo/jax/shampoo.py#L340 )\n\nAnother question, is comparison against https://arxiv.org/pdf/2006.08877.pdf which is considered to be already faster than K-FAC is missing from the experimental section, which could strengthen the paper.\n\nMinor point (as its a recent work): But mentioning here for correctness (https://arxiv.org/abs/2106.06199) has shown K-FAC to work extremely well on these problems when tuned really well; They find the algorithms are sensitive to how often inverse is run, is another confounding factor in empirical comparisons.\n\n\n== update ==\n\nI am leaning towards acceptance as authors promised to tune hyper-parameters including warm-up and decay for the baselines.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "24"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper develops a new 2nd-order optimization method targeted at neural networks that combines aspects of K-FAC and Shampoo. In particular, it uses a Shampoo-style approximation to estimate the Fisher information matrix instead of the Empirical Fisher (as in Shampoo), and like K-FAC raises the curvature matrix approximation to the power -1 to compute the preconditioner. The reviewers agree that this is a well written paper that makes a solid contribution to the area. 

The main concerns of the reviewers are related to the experiments and the tuning of different optimizers. I share these concerns, but am pleased to see that the authors are quite willing to rerun their experiments based on the reviewer's guidance and update their conclusions accordingly. One thing I would like to bring up again is the use of SVD (which is much slower than inverses) and constant damping in K-FAC. Especially for those autoencoder problems a constant damping is highly suboptimal. The Tensorflow release of K-FAC has a MNIST autoencoder experiment implementation (under tensorflow_kfac/examples/autoencoder_mnist.py) which reproduces the results from the original paper (which gets a loss of 50.25 after 1.5k steps using an increasing batch size schedule). It would be quite interesting to see how well the proposed preconditioner would performed as a drop-in replacement for K-FAC in that setting, or even one with a fixed batch size of 60k.

Another thing probably worth pointing out is that this method, like Shampoo, uses the tensor shape of the network's parameters to construct its approximations. That this works well for convnets has entirely to do with the specific role of these parameters play in the network, and how this just happens to be related to their shape in typical implementations. In general, shape can be totally arbitrary, and there is nothing preventing me from parameterizing all of the convolutional layers in my network with flat vectors, or with tensors of the same shape but with randomly permuted entries.
Decision: Accept (Spotlight)