Summary
Training modern deep neural networks is challenging due to their size and the inherent non-convexity of the training objective.
The most popular optimizers for training these models are first-order methods such as SGD, Adam, and AdamW.
In classical convex optimization, it is well-established empirically and theoretically that methods that incorporate second-order curvature information outperform first-order methods.
However, directly extending these methods to the Deep Learning setting leads to prohibitive computational and memory costs.
Fortunately, more computationally friendly approximate second-order optimizers have emerged in recent years, and they have shown promising performance relative to methods like Adam and AdamW.
Foremost among these methods is Shampoo, which recently won the Algoperf benchmark for training neural networks commonly deployed in practice.
Despite its excellent empirical performance, the reasons underlying Shampoo's success have remained a mystery.
This paper helps explain the method's success by providing a better theoretical understanding of what Shampoo's preconditioner approximates.
The paper provides experiments that corroborate their theory and show it can explain the observed practical performance of variants of Shampoo.
Finally, the paper provides empirical studies investigating how common tricks employed in practice for computational efficiency impact the quality of the Hessian approximation.
Strengths
**Clarity:**
Overall, I found the paper easy and enjoyable to read.
The paper is very well-organized, and the authors make all their points clearly.
For instance, usually, each theoretical result comes with a plot demonstrating the result empirically. This was very nice, made the presentation more concrete, and helped to immediately validate the theory.
**Contribution:**
The theoretical insights into the Shampoo preconditioner are valuable and timely, as Shampoo is arguably the most practically effective approximate second-order method currently available.
This paper's results can possibly lead to improved convergence analyses of Shampoo-like optimizers and inspire the design of new optimizers in this area.
This paper will be of interest to anyone in the ICLR community who is interested in more efficient model training and deep learning optimization.
The experiments in Section 4 give useful practical insights into how computational tricks commonly employed in practice impact the approximation quality of $H_{GN}$.
I think these results can help practitioners to develop better heuristics that balance the tradeoff in the preconditioner delivering better approximation quality vs. computation time.
**Correctness:**
I read through all the proofs of the various results in the paper. They are clear and correct.
**Overall:**
I enjoyed the paper and think it makes a nice contribution to the deep learning optimization literature, so I recommend acceptance.
Weaknesses
**Introduction**
The first paragraph's discussion of the cost of 2nd-order methods should be conditioned a bit; otherwise, it is a bit misleading.
A quadratic space requirement and a cubic computational complexity only arise if you naively try to apply classical techniques like Newton's method to Deep Learning.
By leveraging automatic differentiation, we can compute hvps (which is usually all we need) without forming the Hessian at a cost of $\mathcal O(np)$, where $n$ is the number of samples and $p$ is the number of parameters.
The same holds for the Fisher and Gauss-Newton matrices.
Of course, if no minibatching is done, this too is prohibitively expensive as n and p are both very large.
However, in practice, good performance is often obtained using subsampling.
Many methods in the literature exploit this trick and are called Hessian-free.
The authors are aware of such methods, as they discuss them in the appendix.
Given this, in the final version, I would like the statements of the first paragraph to be qualified.
**Experiments**
One area where the paper could improve is if, in addition to plotting the approximation quality against the iteration count, it also gave plots showing how well the optimizer with that approximation performed on the actual objective loss.
Preferably, these would appear side-by-side.
This would give better insight into the impact of the preconditioner's approximation quality on the optimizer's convergence speed.
For instance, in the setting of Fig 3. it would be interesting to see if $Shampoo^2$-Real Labels exhibited much worse performance than $Shampoo^2$-Sampled Labels, as $Shampoo^2$-Real Labels provides a much worse approximation to $H_{GN}$ in the minibatch setting.
**Related Work**
The authors have been thorough in their list of related work in the appendix, but I feel they should add the citation:
Yao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K. and Mahoney, M., 2021, May. AdaHessian: An adaptive second order optimizer for machine learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 35, No. 12, pp. 10665-10673).
AdaHessian is another popular stochastic second-order optimizer based on a stochastic approximation to the Hessian diagonal combined with Adam-style momentum and weighted averaging.
Questions
1. In the revision, could you also include plots showing training loss, as per my comment in Weaknesses?
2. Do the authors have any thoughts on why $Shampoo^2$-Real Labels provides a much better approximation quality in the minibatch setting on CIFAR-10 than on ImageNet and Binary MNIST? This is not an essential question; I'm merely curious.