Shampoo: Preconditioned Stochastic Tensor Optimization

The loss landscapes of neural networks are far from isotropic: some directions have sharp curvature while others are nearly flat. A first-order optimizer like SGD uses the same learning rate for every direction, so it oscillates in the steep directions and crawls in the flat ones. This is the fundamental motivation for preconditioning — multiplying the gradient by a matrix that rescales each direction according to the local curvature.

Full-matrix AdaGrad does exactly this: it accumulates the outer products of past gradients into a single matrix $H_t = \sum_{s=1}^{t} g_s g_s^\top$ and uses $H_t^{-1/2}$ to precondition the gradient. The result is an adaptive step size for each direction in parameter space that provably achieves optimal regret. The problem? For a network with $d$ parameters, $H_t$ is a $d \times d$ matrix. A modest layer with 1000 inputs and 1000 outputs has $d = 10^6$ parameters, making $H_t$ a $10^{6} \times 10^{6}$ matrix — a trillion entries. Storing it, let alone inverting it, is completely impossible.

Practitioners therefore use diagonal approximations: Adam tracks only the element-wise second moment of each parameter. This is cheap — $O(d)$ memory — but it throws away all information about how parameters interact. The cross-parameter curvature that makes full preconditioning so powerful is entirely lost.

The key observation behind Shampoo is deceptively simple: neural network parameters are naturally organized as matrices and tensors, not as flat vectors. A fully connected layer's weights form an $m \times n$ matrix. A convolutional layer's weights form a 4D tensor (input-depth × width × height × output-depth). Flattening these into a single vector, as full-matrix methods require, discards the structural information that could make preconditioning tractable.

Shampoo asks: what if we precondition each dimension of the tensor separately? For an $m \times n$ weight matrix $W$, instead of one $mn \times mn$ preconditioner, maintain two smaller matrices — a left preconditioner $L_t \in \mathbb{R}^{m \times m}$ that captures row-wise (output) curvature, and a right preconditioner $R_t \in \mathbb{R}^{n \times n}$ that captures column-wise (input) curvature.

The update rule accumulates gradient statistics separately for each dimension. On each step $t$, given the gradient matrix $G_t$:

• Update the left preconditioner: $L_t = L_{t-1} + G_t G_t^\top$ (captures row correlations)
• Update the right preconditioner: $R_t = R_{t-1} + G_t^\top G_t$ (captures column correlations)
• Take the preconditioned step: $W_{t+1} = W_t - \eta \, L_t^{-1/4} \, G_t \, R_t^{-1/4}$

Notice that $G_t G_t^\top$ is just an $m \times m$ matrix and $G_t^\top G_t$ is $n \times n$ — both are easy to store and manipulate. The key is the $-1/4$ power (rather than $-1/2$ as in full AdaGrad): since there are two preconditioners contributing, each takes a square root of the square root, resulting in the familiar $O(1/\sqrt{t})$ effective step-size decay.

The theoretical justification for Shampoo comes from a beautiful connection to the Kronecker product. When you flatten the weight matrix $W$ into a vector $w = \overline{\text{vec}}(W)$, the Shampoo update becomes a standard preconditioned gradient step with preconditioner $H_t = L_t^{1/4} \otimes R_t^{1/4}$, where $\otimes$ denotes the Kronecker product.

The critical insight is Lemma 8 from the paper: the Kronecker product of the per-dimension statistics lower-bounds the full statistics matrix. Formally, for gradients of rank at most $r$:

$\epsilon I + \frac{1}{r} \sum_t g_t g_t^\top \;\leq\; L_T^{1/2} \otimes R_T^{1/2}$

This means Shampoo's implicit preconditioner is at least as strong as the full-matrix preconditioner (up to a rank factor). The small eigenvalues — which are the most important for effective preconditioning because they correspond to poorly-conditioned directions — are preserved. Shampoo doesn't just approximate the full preconditioner; it provably dominates it in the operator sense.

The matrix case is clean and intuitive, but Shampoo's real power emerges with higher-order tensors. A convolutional layer's weights are 4-dimensional: input-channels × width × height × output-channels. Shampoo maintains a separate preconditioner for each of these four dimensions.

The algorithm generalizes cleanly. For an order-$k$ tensor of dimensions $n_1 \times \cdots \times n_k$, Shampoo maintains $k$ preconditioner matrices $H_t^i \in \mathbb{R}^{n_i \times n_i}$. Each is updated using the contraction of the gradient with itself along all but the $i$-th dimension: $H_t^i = H_{t-1}^i + G_t^{(i)}$, where $G_t^{(i)} = \text{mat}_i(G_t)\text{mat}_i(G_t)^\top$. The preconditioned gradient is obtained by multiplying each dimension by $(H_t^i)^{-1/2k}$ using the tensor-matrix product $\times_i$.

This extension is non-trivial: the convergence analysis requires proving that the Kronecker product of $k$ dimension-wise preconditioners still lower-bounds the full preconditioner. The proof uses the operator monotonicity of geometric means of commuting positive semidefinite matrices — a deep result from matrix analysis.

Shampoo's analysis is set in the online convex optimization framework. The learner sees a sequence of convex loss functions $f_1, \ldots, f_T$ and wants to minimize regret — the difference between its cumulative loss and the loss of the best fixed point in hindsight.

The main theorem (Theorem 10) bounds the regret of the general tensor version:

$\text{Regret}_T \leq \sqrt{2r}\, D \prod_{i=1}^{k} \text{Tr}\big((H_T^i)^{1/2k}\big)$

where $r$ is the geometric mean of the gradient ranks across dimensions, $D$ is the diameter of the iterates, and $H_T^i$ is the accumulated statistics for dimension $i$. Under standard assumptions, each trace term scales as $O(T^{1/2k})$, so the product of $k$ such terms gives $O(T^{k \cdot 1/2k}) = O(\sqrt{T})$ — the optimal rate for stochastic convex optimization.

The raw algorithm requires computing matrix powers $(H_t^i)^{-1/2k}$ every step, which involves an eigendecomposition or SVD. Three practical tricks make Shampoo competitive in wall-clock time:

• Delayed preconditioner updates. Recompute the matrix roots every 20–100 steps instead of every step. Between updates, reuse the stale roots. The authors found almost no accuracy loss from this delay.
• Momentum. Incorporate standard momentum: $\bar{G}_t = \alpha \bar{G}_{t-1} + (1-\alpha) G_t$ with $\alpha = 0.9$. The preconditioners are still updated with the raw gradient, but the step is taken with the momentum-averaged gradient.
• Diagonal fallback for large dimensions. When a dimension exceeds roughly 1200, switch to a diagonal preconditioner for that dimension only. Other dimensions continue to use full matrices. This makes Shampoo applicable even when some tensors have one very large dimension.

Surprisingly, despite its richer computation, Shampoo's runtime per step is comparable to SGD, Adam, and AdaGrad in practice. The matrix operations are highly parallelizable on GPUs, and the preconditioner matrices are small enough for efficient eigendecomposition.

The paper evaluates Shampoo on image classification (CIFAR-10/100 with ResNets and Inception networks) and language modeling (LM1B with attention models). Across all experiments, Shampoo converged significantly faster than SGD, Adam, and AdaGrad in terms of training steps.

On CIFAR-10 with a 32-layer ResNet, Shampoo reached lower training loss in roughly half the iterations needed by Adam. On LM1B language modeling, Shampoo achieved the same test perplexity as Adam in a fraction of the steps. Crucially, the runtime per step was comparable — the extra computation for maintaining and applying preconditioners was negligible compared to the forward and backward passes of the network itself.

Perhaps the most striking result is the runtime comparison: despite performing SVD-based matrix operations, Shampoo processes roughly the same number of batches per second as SGD on a Tesla K40 GPU. In some configurations (ResNet-55 on CIFAR-100), Shampoo was actually faster per step than the baselines.

Both Shampoo and K-FAC use Kronecker-factored preconditioners, but they come from fundamentally different motivations:

• K-FAC approximates the Fisher information matrix of a neural network by assuming statistical independence between the activations and the backpropagated gradients. It is deeply tied to the structure of backpropagation in feed-forward networks and needs to sample from the model's predictive distribution.
• Shampoo makes no assumptions about the model architecture. It simply observes that parameters live in tensor spaces and uses the gradient's own statistics to build per-dimension preconditioners. The theoretical guarantees hold for any convex loss, not just neural network losses.

This architectural agnosticism is Shampoo's key practical advantage: it works as a drop-in optimizer replacement. You don't need to modify the model or the training code — just swap the optimizer. K-FAC, by contrast, requires knowledge of each layer's type and structure, making it harder to implement as a general-purpose tool.

Shampoo opened the door to a family of practical second-order optimizers for deep learning. Its influence radiates in several directions:

• SOAP (2024) showed that Shampoo with the 1/2 power is equivalent to running Adafactor in the eigenbasis of the preconditioner, and stabilized Shampoo using Adam-like running averages instead of cumulative sums.
• Muon (2024) simplifies Shampoo by dropping the temporal accumulation entirely: when you take just the current gradient's SVD and map it to its polar decomposition $UV^\top$, you get an update that equalizes all singular values. Muon approximates this with Newton-Schulz iterations, which scale better than SVD, making it practical for very large models.
• muP (Maximal Update Parameterization) research showed that Shampoo and K-FAC both connect to the deeper theoretical question of how to scale neural network parameterizations so that optimal hyperparameters transfer across model sizes. Proper scaling of Shampoo's damping terms enables width-invariant training.

A Shampoo variant won the inaugural AlgoPerf optimization competition in 2023, demonstrating that structured preconditioning can outperform Adam-family optimizers in a standardized benchmark setting.