Shampoo brings second-order curvature to neural network training by keeping small per-dimension preconditioners and combining them through a Kronecker product — instead of ever forming the impossibly large full-matrix preconditioner. This page visualizes the algorithm, the memory/compute tradeoff, the reported results, and where the hosts pushed back on the paper's claims.
Every method in this chain approximates true curvature (the Hessian) more cheaply than the last. Shampoo's move is structural: keep gradients in their natural tensor shape instead of flattening to a vector.
For a 1000×1000 weight layer, a full AdaGrad preconditioner needs (mn)² = 10¹² entries.
Shampoo needs m² + n² ≈ 2×10⁶. Log-scale bars below.
Inverting/rooting the full preconditioner costs O((mn)³). Shampoo's two small
matrices cost O(m³ + n³) — linear in the larger dimension, not quadratic.
A fully-connected layer is a matrix — Shampoo sandwiches the gradient between two small preconditioners,
L_t (rows) and R_t (columns). A conv filter bank is a
higher-order tensor — Algorithm 2 generalizes to one preconditioner per mode.
Click through the four stages of a single optimizer step.
Shampoo runs within a small margin of SGD/AdaGrad/Adam on most workloads, and comes out faster than all three on the 55-layer ResNet — the result the hosts flagged as surprising.
LM1B is the standout: test log-perplexity separates from the pack early and stays lower through all 500K steps — a noticeably bigger gap than either image run produced.
Past a per-dimension size threshold (~1200 in the paper's experiments), Shampoo falls back to a diagonal approximation — plain diagonal AdaGrad with extra bookkeeping. Modern transformer FFN widths (4096–16000+) sit well past that line.
Related Work spends real space on K-FAC (Martens & Grosse, 2015) — the closest existing structure-aware, Kronecker-based method. It never appears in Table 1 or Figures 2–4. Every reported win is against SGD, AdaGrad, and Adam, not the one method built on the same idea.