Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: 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.  

[1] - Khan, M.E. and Rue, H., 2021. The Bayesian Learning Rule. arXiv preprint arXiv:2107.04562.

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. 

I did not fully check the maths in the appendix regarding the convergence rate theorems and lemmas, but they look correct at a glance.

A few other concerns on the experimental section:

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

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

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

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


limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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. 
 
The 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 )

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

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


== update ==

I am leaning towards acceptance as authors promised to tune hyper-parameters including warm-up and decay for the baselines.

limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
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.

In 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} ?

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

Minor nit - the citations in the Appendix do not match with the main paper.

Update 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.
limitations_and_societal_impact: None
review_title: null

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)