A New Perspective on Shampoo's Preconditioner

TL;DR

This paper explores the theoretical foundations of Shampoo's Kronecker product preconditioner, revealing its connection to optimal matrix approximations and demonstrating its effectiveness across datasets and architectures.

Contribution

It provides a novel theoretical link between Shampoo's approximation and the optimal Kronecker product, clarifying misconceptions and analyzing practical efficiency tricks.

Findings

Shampoo's approximation is close to the optimal Kronecker product approximation.

The square of Shampoo's approximation relates to power iteration for matrix approximation.

Practical tricks like using batch gradients improve Hessian approximation quality.

Abstract

Shampoo, a second-order optimization algorithm which uses a Kronecker product preconditioner, has recently garnered increasing attention from the machine learning community. The preconditioner used by Shampoo can be viewed either as an approximation of the Gauss--Newton component of the Hessian or the covariance matrix of the gradients maintained by Adagrad. We provide an explicit and novel connection between the $\textit{optimal}$ Kronecker product approximation of these matrices and the approximation made by Shampoo. Our connection highlights a subtle but common misconception about Shampoo's approximation. In particular, the $\textit{square}$ of the approximation used by the Shampoo optimizer is equivalent to a single step of the power iteration algorithm for computing the aforementioned optimal Kronecker product approximation. Across a variety of datasets and architectures we empirically demonstrate that this is close to the optimal Kronecker product approximation. Additionally, for the Hessian approximation viewpoint, we empirically study the impact of various practical tricks to make Shampoo more computationally efficient (such as using the batch gradient and the empirical Fisher) on the quality of Hessian approximation.