Shampoo and Second-Order Preconditioning

Second-order optimization methods use curvature information — how the loss surface bends, not just how steeply it slopes — to take smarter, better-scaled steps than plain gradient descent, and they generally converge in far fewer iterations. The problem is that the full curvature matrix for a neural network with millions of parameters is far too large to ever compute or store directly. Shampoo works around this for weight matrices (which are how most neural network layers are organized) by building a separate, much smaller preconditioner for each dimension of the matrix — one for its rows, one for its columns — rather than one preconditioner covering every entry jointly.

This is a much coarser approximation than the true curvature matrix, but it's cheap enough to maintain at every training step, and it captures enough of the real curvature structure to meaningfully reduce the number of training iterations needed, in exchange for more computation per iteration than a first-order method like Adam.

A worked example

For a weight matrix with 1,000 rows and 1,000 columns, exact second-order preconditioning would require inverting a matrix over a curvature space with a million dimensions — infeasible. Shampoo instead maintains and inverts two much smaller 1,000×1,000 matrices, one per dimension, which is computationally tractable, and in large-scale training benchmarks this has let models reach a target accuracy in meaningfully fewer training steps than Adam, at the cost of more expensive per-step computation.