Shampoo optimizer

Adam and its relatives keep one scalar per parameter, a diagonal preconditioner that ignores how parameters move together. A truly full-matrix preconditioner would capture all those correlations but would need a matrix the size of the parameter count squared. Shampoo is the structured compromise: instead of one giant matrix, it keeps a separate, modest preconditioner for each axis of each weight tensor, so it captures correlations along rows and along columns without ever forming the full thing.

For a weight matrix of shape m by n, Shampoo accumulates a left statistic L of size m by m from G G^top and a right statistic R of size n by n from G^top G, where G is the gradient. The update is preconditioned as L^{-1/4} G R^{-1/4}, using inverse fourth roots of the two factor matrices; for higher-order tensors there is one such factor per dimension. The fractional matrix roots are what give it its full-matrix flavor along each axis.

Shampoo gives stronger, correlation-aware preconditioning than diagonal Adam and has proven effective for large-batch and large-language-model training, especially in its distributed form that shards the factor computations. The costs are the memory of the per-axis factor matrices and the periodic expense of computing matrix inverse roots, and it is usually grafted onto a baseline optimizer's step size for stability.