---
title: 'Shampoo: Structure-Aware Deep Learning Optimizer'
url: https://www.emergentmind.com/topics/shampoo
type: topic
---

# Shampoo: Structure-Aware Deep Learning Optimizer

Shampoo is a family of structure-aware preconditioned stochastic gradient optimizers designed to efficiently approximate full-matrix second-order methods for deep learning. By exploiting the matrix and tensor structure of neural network parameter blocks, Shampoo achieves substantially greater curvature adaptation than diagonal or element-wise methods (e.g., Adam), while maintaining much lower memory and computational requirements than full-matrix approaches. Its foundation rests on Kronecker-factored second-moment preconditioners, with rigorous convergence analyses, high empirical efficiency across large-scale deep learning tasks, and a rapidly expanding ecosystem of variants addressing stability, scaling, and hardware constraints [1802.09568, 2309.06497, 2503.10537, 2506.03595, 2602.02016, 2604.17423].

## 1. Mathematical Formulation and Core Algorithm

Shampoo generalizes AdaGrad by constructing a Kronecker product approximation to the empirical covariance of parameter gradients, applied separately to each tensor-shaped parameter block. For a weight matrix $W\in\mathbb{R}^{m\times n}$ and its gradient $G_t$, Shampoo maintains two positive-definite, exponentially weighted accumulators:
\[
L_t = \beta_2 L_{t-1} + (1-\beta_2) G_t G_t^\top \in\mathbb{R}^{m\times m},\quad
R_t = \beta_2 R_{t-1} + (1-\beta_2) G_t^\top G_t \in\mathbb{R}^{n\times n}
\]
with $\beta_2\in[0,1)$ a smoothing parameter [1802.09568, 2506.03595].

The update computes the symmetric inverse $p$-root (often $p=4$) of each factor, and preconditions the gradient:
\[
W_{t+1} = W_t - \eta\, L_t^{-1/4} G_t R_t^{-1/4}
\]
This corresponds, when vectorized, to applying $(L_t \otimes R_t)^{-1/2}$ to the flattened gradient, closely approximating the ideal full-matrix AdaGrad preconditioner with only $O(m^2+n^2)$ memory and $O(m^3+n^3)$ per-factor time, compared to $O((mn)^2)$ and $O((mn)^3)$ for full-matrix AdaGrad [2503.10537, 2604.17423].

In practice, to amortize the cubic cost of eigendecomposition, root-inverses are computed every $F\gg1$ steps ("stale preconditioning"), and step sizes per layer are often grafted to match those of reference optimizers like AdamW [2309.06497, 2506.03595].

## 2. Theoretical Guarantees and Convergence Rates

Shampoo is a specific instance in the class of adaptively preconditioned first-order methods, and its convergence in stochastic online and non-convex optimization is rigorously established. Theoretical results include:

- **Stochastic Convex Regret:** Under convexity, the two-sided variant achieves $O(\sqrt{T})$ regret [1802.09568].
- **Nonconvex Rate:** In general nonconvex settings, the AdamW-style Shampoo achieves
  \[
  \frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left(\frac{\sqrt{m+n}C}{K^{1/4}}\right)
  \]
  where $\|\cdot\|_*$ is the nuclear norm, $K$ is the iteration number, and $C$ encapsulates problem (Lipschitz, smoothness, and variance) parameters [2601.07326]. For practical $m,n$, the $\sqrt{m+n}$ scaling is much smaller than the parameter dimension scaling appearing in vectorized or diagonal methods.

- **Structure Advantage:** Recent unified analyses demonstrate that *one-sided Shampoo* (preconditioning only the row or column dimension) can exhibit sharper regret and finite-time convergence bounds than both two-sided Shampoo and full-matrix AdaGrad, with significantly reduced per-step complexity [2503.10537].

Shampoo's convergence analyses do not require bounded stochastic gradients or extremely small step sizes, unlike several adaptive methods [2604.17423]. Convergence rates are slightly slower than SGD ($O(K^{-1/4})$ nuclear norm vs $O(K^{-1/2})$ Frobenius), but the practical efficiency gains (step size, curvature adaptation) offset this in high-dimensional regimes [2601.07326].

## 3. Numerical, Stability, and Scalability Considerations

Accurate computation of matrix inverse $p$-roots is numerically delicate—small eigenvalues in $L_t$, $R_t$ render the preconditioning operation unstable in low (e.g., float32) precision. Classical implementations require FP64 for root-inverse eigendecompositions on CPUs or fallback routines on accelerators [2305.19416]. To reduce wall time, Shampoo implementations often:

- **Amortize eigendecomposition**: Update roots every $F\gg1$ steps, making preconditioning "stale". However, if $F$ is too large, accumulated staleness can lead to divergence or degraded convergence [2506.07254, 2506.03595].
- **Apply block-wise preconditioning**: For very large layers, block the matrices and batch kernel operations to maximize hardware utilization, as in DASH, which achieves up to $4.8\times$ step speedups compared to distributed Shampoo by stacking blocks into 3D tensors and leveraging batched EVD or polynomial inverse-root solvers [2602.02016].
- **Adaptive damping and refresh scheduling:** Dynamic schemes, such as FOAM, adjust $\epsilon$ and trigger re-computation adaptively, sensing staleness-oriented error to balance numerical stability and computational cost [2606.02365].

Table: Comparison of Key Practical Ingredients

| Feature               | Shampoo         | DASH           | FOAM                | EShampoo/SOAP           |
|-----------------------|----------------|----------------|---------------------|-------------------------|
| FP64 need             | Yes (roots)    | No (FP32/FP16) | No                  | No (rotated basis)      |
| Blocked kernels       | Not required   | Yes            | Optional            | Optional                |
| Adaptive refresh      | No             | No             | Yes                 | Yes (basis criterion)   |
| Grafting/eigenvalue corr. | Often/Yes   | Optional       | Optional            | Yes                     |

Key practical recommendations include block-wise application only to large matrix blocks, periodic (but not too-infrequent) root updates, and, if supported, leveraging batched GPU-friendly root solvers [2602.02016, 2309.06497].

## 4. Algorithmic Extensions and Variants

Several empirically and theoretically motivated extensions have addressed original Shampoo's stability, scale-sensitivity, and tuning complexity [2506.07254, 2409.11321, 2506.03595]:

- **SPlus** [*A Stable Whitening Optimizer for Efficient Neural Network Training*]: Replaces historical eigenvalue normalization with a per-step sign-based update in a fixed eigenbasis, yielding updates with hard-bounded magnitude, shape-aware scaling for width-invariance, and Polyak–Ruppert iterate averaging for noise reduction. SPlus achieves $\sim0.44$ steps-to-Adam and $0.62$ time-to-Adam across diverse large-scale tasks, with robustness to inversion intervals up to 100 steps.
- **SOAP** [*Shampoo with Adam in the Preconditioner's Eigenbasis*]: SOAP demonstrates that Shampoo with exponent $1/2$ is algebraically equivalent to running Adafactor in the preconditioner's eigenbasis and extends this to running AdamW in that basis, updating basis vectors infrequently but per-entry scaling adaptively. SOAP stabilizes performance under infrequent basis refreshes and matches well-tuned AdamW in only $\sim60$\% of the wall-clock time.
- **EShampoo** [*Purifying Shampoo*]: Proposes decoupling the update of preconditioner eigenvalues and eigenvectors. With direct eigenvalue correction per-step (in the rotated basis) and an adaptive criterion for eigenbasis refresh (based on relative off-diagonal error in the stale-projected covariance), it eliminates the need for layerwise norm grafting and increases wall-time efficiency.
- **DASH**: Introduces batched block preconditioning, GPU-friendly Newton–Denman–Beavers and Chebyshev-based root estimation, and power iteration-based scaling for accelerated and scalable distributed training. Maintains Shampoo's per-iteration perplexity while drastically reducing wall-clock time.
- **FOAM**: Dynamically controls damping and refresh frequency via an operator-norm staleness proxy, ensuring stable optimization with reduced eigendecomposition overhead [2606.02365].

## 5. Empirical Performance and Applications

Shampoo and its variants consistently outperform diagonal methods (Adam, AdaGrad) in step- and wall-time efficiency on standard benchmarks (ResNet/ImageNet, ViT, GPT-2/WikiText, Llama/C4), often reducing training steps or time by $1.3$–$1.8\times$ with similar or higher final accuracy [2309.06497, 2509.23500, 2506.07254, 2602.02016]. Key empirical outcomes:

- **Robustness to Quantization**: Shampoo-trained models show the smallest accuracy drop under 4-bit quantization-aware training, with the highest parameter efficiency $\rho_{4\rm bit}\approx0.88$ and minimal zero-shot degradation versus full-precision, outperforming AdamW and Muon even when the max-to-median ratio suggests the opposite [2509.23500].
- **Generalization and Compression**: Shampoo has been shown to yield models with less activation outlier structure, lowering quantization and compression error, and supporting efficient deployment [2602.02016].
- **Distributed Scalability**: DTensor-based sharding and AllGather primitives in distributed PyTorch implementations enable large-scale Shampoo training with at most $10\%$ per-step overhead for up to $1.5\times$ faster convergence in wall-clock time compared to baseline adaptive methods [2309.06497, 2602.02016].

## 6. Interpretations, Connections, and Open Problems

Shampoo's core operation is best understood as a *Frobenius-norm optimal Kronecker-approximation* to the full empirical Fisher (as in full-matrix AdaGrad or Adam), updated with the square-root of the factors to maintain the scaling properties of the vectorized preconditioner [2506.03595]. The update can be decomposed as a spectral descent (polar/singular value factorization) with two adaptive matrix scalings, yielding time-averaged semi-orthogonality in expectation rather than enforcing hard orthogonality (“whitening”) or variance adaptation [2602.09314]. This mechanism is fundamentally distinct from both classical whitening and variance-scaling narratives.

Recent analyses highlight:

- **Structure-exploitation vs. full adaptivity:** Carefully chosen structure (e.g., one-sided Shampoo) can yield both improved regret/convergence constants *and* lower cost than less-structured, fully-adaptive methods—a challenge to the dogma that more adaptivity always confers more efficiency [2503.10537].
- **Elimination of heuristics:** Adaptive eigenbasis frequency and eigenvalue correction can supplant learning rate grafting and ad hoc preconditioner reuse [2506.03595].
- **Momentum and weight decay integration:** Shampoo is commonly equipped with first- and second-moment momentum, decoupled weight decay (AdamW-style), bias correction, and flexible per-layer application (matrix blocks only vs. scalars) [2309.06497, 2601.07326].

Outstanding directions include tight integration of approximation quality into theoretical regret, combining adaptive basis scheduling with hardware-aware batch kernel calls, and extension to more general curvature targets or low-rank approximations [2506.03595].

---

Key references for further details: [1802.09568], [2309.06497], [2503.10537], [2604.17423], [2602.02016], [2506.03595], [2509.23500], [2506.07254], [2409.11321], [2602.09314], [2606.02365], [2601.07326].

Source: https://www.emergentmind.com/topics/shampoo