Shi et al., Meta (2023) — ARXIV-2309.06497
Key takeaways
- Shampoo belongs to the AdaGrad family and builds a block-diagonal preconditioner, each block a coarse Kronecker-product approximation to full-matrix AdaGrad for one parameter. The approximation is what makes the second-order idea tractable at all.
- The contribution here is not the algorithm but the distribution: memory and computation for each parameter’s blocks are spread across ranks via PyTorch’s DTensor, with an AllGather on the computed search directions each step.
- That brings the cost to at most a 10% per-step wall-clock penalty against standard diagonal adaptive methods — the number that decides whether a preconditioning optimizer is usable, and the reason the paper is cited by practitioners rather than only by theorists.
- Validated by ablation on ImageNet ResNet-50, with minimal tuning.
Standing in the anthology
Filed as the predecessor LIT-157 simplifies and the practical baseline the optimizer comparison in LIT-156 includes. It is here mainly for the 10% number: the record’s optimizer practices trade convergence against per-step cost constantly — SOTA-121’s Muon, LIT-131’s per-head variant that is cheaper because the blocks are smaller — and this is where the overhead budget for a matrix preconditioner was first shown to be affordable.
No practice of its own, and the reason is altitude rather than adoption. The contribution is not an algorithm but a distribution strategy for one, and its transferable content — that sharding a preconditioner’s blocks across ranks holds the penalty to 10% — is a condition on a recommendation rather than a recommendation. It is filed as a source of SOTA-165 and stated there, in the section about what the class costs, which is where a reader choosing an optimizer would look for it.