Shampoo

The guarantee (Gupta et al., 2018). Gupta et al. analyze Shampoo as an algorithm for online convex optimizationa repeated game: at each step the learner commits to $W_t$, then a convex loss $f_t$ is revealed and the learner incurs $f_t(W_t)$; stochastic optimization is the case of losses drawn independently from one distribution: at each step $t$ a convex loss $f_t$the loss revealed at step $t$, for example the loss on step $t$'s minibatch, convex in $W$ arrives after the learner has chosen $W_t$, and regret compares the learner's total loss with that of a fixed matrix $W^*$any fixed comparison matrix, typically a minimizer of the total loss. An online-to-batch conversion turns a regret bound into a stochastic optimization error of order regret divided by $T$. The theorem says Shampoo's regret grows like $\sqrt T$, the best possible rate in this setting, with a constant made of the traces of the final fourth roots. If every $f_t$ is convexwhy: the regret bound for preconditioned online gradient steps (Step 3) uses convexity to bound $f_t(W_t)-f_t(W^*)$ by $\langle G_t,W_t-W^*\rangle$ and every $G_t$ has rank at most $r$why: Step 1 splits the gradient into $r$ rank-one pieces, and the Cauchy-Schwarz inequality used there loses a factor $r$; with no rank bound one takes $r=\min(m,n)$, then Shampoo with $\eta=D/\sqrt{2r}$ satisfies, for every $W^*$,$$\begin{aligned}&\underbrace{\sum_{t\le T}f_t(W_t)-\sum_{t\le T}f_t(W^*)}_{\text{regret}}\\&\quad\le\sqrt{2r}\,D\,\underbrace{\mathrm{tr}(L_T^{1/4})}_{\text{left}}\,\underbrace{\mathrm{tr}(R_T^{1/4})}_{\text{right}},\end{aligned}$$where $D$${}=\max_{t\le T}\|W_t-W^*\|_F$. With $\epsilon=0$ and gradients of spectral norm at most 1, $G_tG_t^\top\preceq I_m$ gives $\mathrm{tr}(L_T^{1/4})\le mT^{1/4}$, and likewise $\mathrm{tr}(R_T^{1/4})\le nT^{1/4}$, so the regret is at most $\sqrt{2r}\,D\,mn\sqrt T$. For a tensor with $k$ dimensions the bound is a product of $k$ traces of $1/(2k)$ powers.

One gradient's outer product is bounded by each one-sided factor. Write the gradient in its singular value decomposition $G=USV^\top=\sum_{i\le r}\sigma_iu_iv_i^\top$, with $u_i,v_i$ the columns of $U,V$, so $g=\mathrm{vec}(G)=\sum_i\sigma_i\,u_i\otimes v_i$. For any vectors $w_1,\dots,w_r$, $(\sum_iw_i)(\sum_iw_i)^\top\preceq r\sum_iw_iw_i^\top$, because $(\sum_ix^\top w_i)^2\le r\sum_i(x^\top w_i)^2$ for every vector $x$ by Cauchy-Schwarz; hence $gg^\top\preceq r\sum_i\sigma_i^2(u_iu_i^\top)\otimes(v_iv_i^\top)$. Bounding $v_iv_i^\top\preceq I_n$ leaves $GG^\top\otimes I_n$, and bounding $u_iu_i^\top\preceq I_m$ leaves $I_m\otimes G^\top G$:$$\begin{aligned}\tfrac1r\,gg^\top&\preceq(GG^\top)\otimes I_n,\\\tfrac1r\,gg^\top&\preceq I_m\otimes(G^\top G).\end{aligned}$$Each bound alone discards the structure of one side; the factor $r$ is the Cauchy-Schwarz loss over $r$ rank-one pieces.

The geometric mean of the two bounds is the Kronecker bound. Sum Step 1 over $s\le t$ and add $\epsilon I_{mn}$. With $X=\epsilon I_{mn}+\frac1r\sum_{s\le t}g_sg_s^\top$, this gives $X\preceq \htmlData{gr=g11}{L_t}\otimes I_n$ and $X\preceq I_m\otimes \htmlData{gr=g12}{R_t}$. The two right sides commute, and for commuting positive semidefinite matrices the geometric mean is monotone in each argument (Ando et al.), so$$\begin{aligned}X&=X^{1/2}X^{1/2}\\&\preceq(L_t\otimes I_n)^{1/2}(I_m\otimes R_t)^{1/2}\\&=L_t^{1/2}\otimes R_t^{1/2}.\end{aligned}$$The square roots come from splitting the two bounds evenly. Anil et al. (2020) split unevenly: weights $\alpha$ and $1-\alpha$ give $X\preceq L_t^{\alpha}\otimes R_t^{1-\alpha}$ for any $\alpha\in[0,1]$, and $\alpha=1$ keeps only the left factor, so AdaGrad's square root gives the step $L_t^{-1/2}G$, which is how they precondition layers too large for a right factor.

AdaGrad's regret argument goes through with the Kronecker preconditioner. Write $S_t$${}=L_t^{1/4}\otimes R_t^{1/4}$, so the step is $w\leftarrow w-\eta S_t^{-1}\htmlData{gr=g7}{g_t}$; $S_t$ increases with $t$ because $L_t$ and $R_t$ do. The standard bound for preconditioned online gradient steps with nondecreasing preconditioners gives$$\begin{aligned}\mathrm{Regret}\le{}&\frac{D^2}{2\eta}\,\mathrm{tr}\,S_T\\&+\frac\eta2\sum_{t\le T}g_t^\top S_t^{-1}g_t.\end{aligned}$$The sum is where Step 2 enters. Let $\hat H_t=(r\epsilon I_{mn}+\sum_{s\le t}g_sg_s^\top)^{1/2}$, full-matrix AdaGrad's preconditioner with $\epsilon$ multiplied by $r$. Step 2 times $r$, and the operator monotonicity of the square root (Loewner), give $\hat H_t\preceq\sqrt r\,S_t$. AdaGrad's own potential argument gives $\sum_tg_t^\top\hat H_t^{-1}g_t\le2\,\mathrm{tr}\,\hat H_T$. Chaining the two, $\sum_tg_t^\top S_t^{-1}g_t\le\sqrt r\cdot2\,\mathrm{tr}\,\hat H_T\le2r\,\mathrm{tr}\,S_T$. Choosing $\eta=D/\sqrt{2r}$ balances the two terms at $\sqrt{2r}\,D\,\mathrm{tr}\,S_T$, and $\mathrm{tr}(\htmlData{gr=gn5}{A\otimes B})=\mathrm{tr}A\,\mathrm{tr}B$ gives the theorem. The tensor case is the same argument with a geometric mean of $k$ one-sided bounds.

Roots on idle CPUs, partial preconditioners for large layers, and grafting make Shampoo run at scale. The inverse roots were the bottleneck, up to 100 times a training step. For a Transformer layer the factor's condition number was $10^7$ to $10^9$ during training, and an inverse root loses about $\log_2$ of the condition number in bits of precision, so the roots need float64, which accelerators do slowly. Anil et al. compute them with the coupled Newton iterationan iteration for inverse matrix roots that uses only matrix products and stays numerically stable (Guo and Higham; Iannazzo); here started from the previous root, warm-started from the previous root, in float64 on the CPUs attached to the accelerators, asynchronously, and swap in new roots every few hundred steps; recomputing them that rarely had no significant effect on quality. Large layers get partial preconditioners: embedding and softmax layers use one side only ($L^{-1/2}G$, from the unequal split in Step 2), and large dense layers are cut into blocks, each its own Shampoo matrix. Grafting from AdaGrad, or from the layer-wise scaling of LAMB, sets each layer's step length, because the preconditioners' norms spread over a wide range across layers and caused instability. They also report that treating the power as a hyperparameter sometimes helped, and they used the square root on each side in their autoencoder experiments. On a Transformer for WMT'14 English to French (93M parameters) Shampoo needed 1.95 times fewer steps at about 16% more time per step; on a click-through-rate model on the 1 TB Criteo data it reached the target AUC of 80.25% in 30,970 steps against 64,000; BERT-Large needed 16% fewer steps than LAMB; and ResNet-50 at batch size 32,768 reached 75.9% top-1 accuracy in 1,729 steps against 2,512.

Spreading the preconditioners across GPUs keeps the step within 10% of Adam's, and the result won AlgoPerf. In standard data-parallel training every worker repeats the whole optimizer step, which is cheap for Adam but would make some large models' Shampoo steps 50% to 75% slower. Shi et al. assign each layer's factors and grafting state to one worker, balanced greedily by size, let each worker compute only its layers' preconditioned directions, and gather the directions with one AllGather; the step then costs at most about 10% more than a diagonal method's. The roots use a symmetric eigendecomposition in float32 (retried in float64 if it fails), with the eigenvalues shifted to be at least $\epsilon$. They describe grafting as one of the main requirements for Shampoo to work in practice, and report that the square-root power (or the fourth-root power multiplied by 1.82) is often effective for networks made mostly of linear layers. On ResNet-50 for ImageNet, with roots recomputed every 50 steps, blocks of at most 2048 dimensions and grafting from SGD with the tuned Nesterov recipe's hyperparameters unchanged, Shampoo reached 77.44% validation accuracy in 90 epochs against 76.85% (means over five seeds) at 8% more time, and the Nesterov 90-epoch accuracy with 1.5 times fewer steps and 1.35 times less time. With Adam grafting and roots every 100 steps, the implementation won the external-tuning track of AlgoPerf, about 28% faster on average than the NAdamW baseline and about 2% faster than the second-place entry, though five different entries were fastest on at least one workload (Kasimbeg et al., 2025).

$L\otimes R$ is one power-iteration step toward the best Kronecker approximation, and one step is enough in practice. Rearrange $H$ so that the entry $\mathbb E[G_{ij}G_{kl}]$ sits at row $(i,k)$ and column $(j,l)$. A Kronecker product $A\otimes B$ then becomes $\mathrm{vec}(A)\mathrm{vec}(B)^\top$, the Frobenius distance is preserved, and multiplying the rearranged $H$ by $\mathrm{vec}(B)$ gives $\mathrm{vec}(\mathbb E[GBG^\top])$; this is the matrix form of power iteration in the mechanism above. Starting from $B=I_n$ and $A=I_m$ gives $L$ and $R$. Why one step suffices: write the singular value decomposition of the rearranged $H$ as $\sum_i\theta_i\,\mathrm{vec}(A_i)\mathrm{vec}(B_i)^\top$, with unit-norm matrices $A_i,B_i$. One step from the identity gives $\sum_i\theta_i(\mathrm{tr}\,B_i)A_i$, so it lands near the top factor $A_1$ when $\theta_1\mathrm{tr}\,B_1$ is large next to the other $\theta_i\mathrm{tr}\,B_i$. The top factor $B_1$ is positive semidefinite, and among positive semidefinite matrices of unit Frobenius norm the scaled identity has the largest worst-case overlap with the others, while if $B_1$ is positive definite the later $B_i$ are not positive semidefinite, so their traces can cancel. In their measurements on binary MNIST, CIFAR-5M and ImageNet, the cosine similarity of $L\otimes R$ with $H$ stayed close to that of the optimal Kronecker product throughout training, and Shampoo's $L^{1/2}\otimes R^{1/2}$ was clearly lower, for $H$ both AdaGrad's statistic and the Gauss-Newton matrixthe Hessian of the loss with the network linearized at the current weights, $J^\top\nabla_f^2\ell\,J$, where $J$ is the Jacobian of the network output with respect to the weights and $\nabla_f^2\ell$ the Hessian of the loss in the output; for cross-entropy it equals $\mathbb E[gg^\top]$ with labels sampled from the model. If $H=A\otimes B$ exactly, then $L=(\mathrm{tr}\,B)A$ and $R=(\mathrm{tr}\,A)B$, so $H=L\otimes R/\mathrm{tr}\,L$. Since practitioners tune the power anyway, the paper notes that squaring adds no new power to try; it changes which power to expect. For the Gauss-Newton reading they also find that averaging gradients over a batch is harmless with labels sampled from the model; with the real labels, as Shampoo uses, the approximation is still good at batch size 1 but degrades at larger batches, because real-label gradients do not average to zero.

In Shampoo's eigenbasis the square-root form is Adafactor, so SOAP runs Adam there instead. Take the idealized Shampoo with averages over the data, the square-root power and the trace factor, so entry $(i,j)$ of $\tilde G$ is divided by $\sqrt{\htmlData{gr=g14}{\lambda_i}\mu_j/\mathrm{tr}\,L}$. Adafactor run on $\tilde G$ estimates $\mathbb E[\tilde G\odot\tilde G]$ by $a\,c^\top/\sum_ia_i$ from its row sums $a$ and column sums $c$, and$$\begin{aligned}a_i&=\textstyle\sum_j\mathbb E[(q_i^\top G\bar q_j)^2]\\&=\mathbb E\|q_i^\top G\|^2=q_i^\top Lq_i=\lambda_i,\end{aligned}$$with $q_i$ the $i$-th column of $Q_L$ and $\bar q_j$ the $j$-th column of $Q_R$; likewise $c_j=\htmlData{gr=g15}{\mu_j}$, and $\sum_ia_i=\mathrm{tr}\,L$. The two scales are the same. Shampoo can update its scale only when it recomputes $Q_L,Q_R$, but the per-entry average is cheap, so SOAP keeps it as Adam does, with momentum $M$SOAP's momentum, an exponential moving average of the gradient with rate $\beta_1$, kept in the original coordinates and rotated each step and second moment $V$SOAP's per-entry second moment, an exponential moving average with rate $\beta_2$ of the squared rotated gradient $\tilde G\odot\tilde G$ updated every step:$$\begin{aligned}V&\leftarrow\beta_2V+(1-\beta_2)\,\tilde G\odot\tilde G,\\W&\leftarrow W-\eta\,Q_L\Big(\frac{Q_L^\top MQ_R}{\sqrt V+\epsilon}\Big)Q_R^\top.\end{aligned}$$Here $\beta_1,\beta_2$the exponential-moving-average rates of the momentum and of the second moment (and of $L$ and $R$), as in Adam are Adam's rates, and $L$ and $R$ are themselves moving averages with rate $\beta_2$. The basis is refreshed every 10 steps in the main runs by one step of power iteration with a QR decomposition, started from the previous $Q_L,Q_R$; one-dimensional parameters run plain Adam, and a very large dimension keeps the identity on its side. Grafting goes away, and the only hyperparameter beyond Adam's is how often to refresh the basis. Anil et al. (2020) had described essentially this algorithm for matrices in an appendix and reported no gain in unpublished experiments; Vyas et al. attribute the difference to implementation details. On 360M- and 660M-parameter language models with batches of 2 million tokens, SOAP needed at least 40% fewer steps and 35% less time than AdamW to reach AdamW's final loss, and about 20% less of both than the distributed Shampoo. As the refresh interval grew from 1 to 100 steps both beat AdamW throughout, they were close at 1, and SOAP lost quality much more slowly. At 256,000-token batches the gain fell to 25% fewer steps than AdamW and about 10% fewer than Shampoo. SOAP stores $2m^2+2n^2+3mn$ numbers per matrix, as the distributed Shampoo does, against AdamW's $3mn$ (both counts include the gradient). The times were measured on one GPU with gradient accumulation.

Grafting corrects Shampoo's eigenvalues, and a correct, current scale removes the need for it. Without grafting and with roots recomputed every 100 steps, Shampoo trained worse than AdamW at every learning rate and $\epsilon$ tried, for a vision transformer and a ConvNeXt V2 on Imagewoof, and many settings diverged. Any step of the form $Q_L[\tilde G_{ij}/d_{ij}^{\,p}]Q_R^\top$, with orthogonal $Q_L,Q_R$, per-entry scales $d_{ij}>0$ and a power $p>0$, has Frobenius norm between $(\max d)^{-p}\|G\|_F$ and $(\min d)^{-p}\|G\|_F$, so only the scale sets a step's length; the basis never does. For full-matrix Adam, diagonal Adam and the per-entry-corrected Shampoo, all with $p=1/2$, the limits can be taken as $\lambda_{\max}(\mathbb E[gg^\top])^{-p}\|G\|_F$ and $\lambda_{\min}(\mathbb E[gg^\top])^{-p}\|G\|_F$; for Shampoo, with independent gradient entries, they carry the extra factor $(mn)^{-1/4}$ of the example in the mechanism above, and stale eigenvalues can push the step outside them. With roots recomputed every step, the trace-scaled square-root Shampoo matched or beat grafted Shampoo; with the basis recomputed every 100 steps and a per-entry scale updated every step, as in SOAP, the result matched or beat grafted Shampoo, and Adam's best learning rate carried over. On AlgoPerf workloads it needed 2,536 steps against 4,301 on FastMRI (10.4 against 14.0 minutes) and 8,320 against 12,574 on OGBG (means over the runs that reached the target), and about the same on ImageNet with a vision transformer. For the basis they propose recomputing a factor only when the rotated factor $\hat\Lambda=\hat Q^\top L\hat Q$, with $\hat Q$ the last computed eigenvectors, has drifted from diagonal; while $\|\hat\Lambda-\mathrm{diag}(\hat\Lambda)\|_F\le$ $\tau$the tolerance of the adaptive rule: a factor's eigenbasis is recomputed only when the off-diagonal part of the rotated factor exceeds this fraction of its Frobenius norm$\|\hat\Lambda\|_F$, the old basis is kept. On the Imagewoof vision transformer, recomputing every 100 steps took 20% more time than this rule checked every 50 steps with $\tau=0.1$. On the same vision transformer, an accurate basis mattered mostly early in training: a basis computed at the first step and never updated still beat AdamW (final training loss about 0.12 against 0.3), and the whole gap to AdamW came from the bases of two-dimensional parameters, so vectors can be left to Adam.