Skip to content

Shampoo optimiser

How Shampoo preconditions matrix gradients on both sides, why inverse fourth roots appear, and how accumulated history and implementation choices change the update.

The Shampoo optimiser is an adaptive gradient descent method that transforms gradients using matrices built from their past outer products, preserving the shape of each parameter tensor. For a quant researcher training a neural forecasting model, it offers richer scaling than coordinatewise methods such as Adam, at additional memory and computation cost. Faster training still needs to be demonstrated on the actual task.

Precondition the rows and the columns

For a parameter matrix Wt∈Rm×nW_t\in\mathbb{R}^{m\times n} and its current loss-function gradient GtG_t, the original matrix version keeps
Lt=ϵIm+∑s=1tGsGs⊤,Rt=ϵIn+∑s=1tGs⊤Gs,Dt=Lt−1/4GtRt−1/4,Wt+1=Wt−ηtDt.\begin{aligned} L_t &= \epsilon I_m+\sum_{s=1}^t G_sG_s^\top, \\ R_t &= \epsilon I_n+\sum_{s=1}^t G_s^\top G_s, \\ D_t &= L_t^{-1/4}G_tR_t^{-1/4}, \\ W_{t+1} &= W_t-\eta_tD_t. \end{aligned}
Here ηt>0\eta_t>0 is the learning rate and ϵ>0\epsilon>0 supplies damping. LtL_t captures row-gradient relationships; RtR_t captures column-gradient relationships. The powers are matrix powers: apply the inverse fourth root to each positive eigenvalue, retaining the eigenvectors. Taking each matrix entry to that power is incorrect.
The two factors need m2+n2m^2+n^2 stored entries, versus (mn)2(mn)^2 for a full preconditioner on flattened parameters, before counting other optimiser state. They represent a structured Kronecker-product preconditioner, not the unrestricted full matrix. See Algorithm 1 of Gupta, Koren and Singer and the scaling paper's preliminaries.

A first step, then the same gradient again

Use a fictional, dimensionless two-by-two gradient

G1=(2004).G_1=\begin{pmatrix}2&0\\0&4\end{pmatrix}.

For this calculation only, set ϵ=0\epsilon=0, start with zero history, and omit momentum and grafting. Both Gram matrices are invertible, so the undamped roots exist:

L1=R1=(40016),L1−1/4=R1−1/4=(1/2001/2).L_1=R_1=\begin{pmatrix}4&0\\0&16\end{pmatrix},\qquad L_1^{-1/4}=R_1^{-1/4} =\begin{pmatrix}1/\sqrt{2}&0\\0&1/2\end{pmatrix}.

Multiplication gives

D1=((1/2) 2 (1/2)00(1/2) 4 (1/2))=I2.D_1=\begin{pmatrix} (1/\sqrt{2})\,2\,(1/\sqrt{2})&0\\ 0&(1/2)\,4\,(1/2) \end{pmatrix}=I_2.

The actual parameter change is −η1I2-\eta_1I_2. The larger raw diagonal gradient has been scaled more strongly.

Now suppose the second gradient happens to be identical: G2=G1G_2=G_1. The accumulators double, giving L2=R2=diag⁡(8,32)L_2=R_2=\operatorname{diag}(8,32). Each inverse fourth-root factor contributes 2−1/42^{-1/4} relative to the first step, so

D2=2−1/2I2≈0.707107I2.D_2=2^{-1/2}I_2\approx0.707107I_2.
Analytic Shampoo example with the same diagonal gradient 2 and 4 at both steps: raw components are 2 and 4; preconditioned components are 1 and 1 at step one, then approximately 0.707 and 0.707 at step two. Learning rate excluded.
Analytic Shampoo example with the same diagonal gradient 2 and 4 at both steps: raw components are 2 and 4; preconditioned components are 1 and 1 at step one, then approximately 0.707 and 0.707 at step two. Learning rate excluded.
Accumulated history changes the scaling even when the current gradient is unchanged. This toy uses no damping, momentum or grafting.

This first-step identity is a property of this example. It is not a rule that Shampoo updates every parameter by one learning rate: even here, the second step differs. Singular gradients, damping, different histories and practical variants also change the calculation.

The implementation is part of the experiment

The authors' scaling work explains the practical costs: forming statistics, computing inverse matrix roots and applying the factors. Large tensors can be partitioned into smaller blocks, trading some cross-block information for manageable memory and computation. Roots can be refreshed less often than the statistics, leaving the update to use older factors between refreshes. Fewer iterations alone therefore do not establish lower elapsed training time.
The Google Research author implementation, inspected on 1 November 2026, exposes separate statistics and preconditioner computation intervals, block size, matrix damping, exponential averaging and momentum. Its grafting step rescales the Shampoo direction to the norm of an SGD or diagonal AdaGrad direction before subsequent momentum handling. This changes the magnitude from the bare formula above. The accompanying README also documents shape merging and exponent overrides. Record these choices and the code revision when comparing runs.
Shampoo is often called a second-order method because it uses second-moment gradient statistics. It does not compute the loss Hessian and is not a full Newton step. Sharing a factored structure with K-FAC does not make their statistics identical: K-FAC targets a Fisher approximation, whereas the update here accumulates gradient Gram matrices. The original paper's convex optimisation analysis does not guarantee a global optimum for a neural network.
For the experimental discipline around this choice, Reading Loss Curves covers logging and validation-based model selection. Its preview is public; the full lesson requires membership. An improved training curve does not by itself establish a better financial forecast.

Sources

Last updated
  • machine learning
  • optimisation
  • shampoo
  • preconditioning
  • deep learning
Keep connecting the ideas
← Back to the glossaryBack to top ↑