# Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization

arXiv: https://arxiv.org/abs/2509.03378 · HTML: https://mlsanity.com/paper/2509.03378/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 64 references that are in the map.
- First posted on arXiv 3 Sep 2025.

## Builds on (most cited first, 19 of 19)

- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [Manopt, a Matlab toolbox for optimization on manifolds](https://mlsanity.com/paper/1308.5200/) (2013)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Conjugate-Computation Variational Inference : Converting Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models](https://mlsanity.com/paper/1703.04265/) (2017)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [ASGO: Adaptive Structured Gradient Optimization](https://mlsanity.com/paper/2503.20762/) (2025)
- [Benchmarking Neural Network Training Algorithms](https://mlsanity.com/paper/2306.07179/) (2023)
- [Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations](https://mlsanity.com/paper/1906.02914/) (2019)
- [Disentangling Adaptive Gradient Methods from Learning Rates](https://mlsanity.com/paper/2002.11803/) (2020)
- [Structured Preconditioners in Adaptive Optimization: A Unified Analysis](https://mlsanity.com/paper/2503.10537/) (2025)
- [Accelerating Neural Network Training: An Analysis of the AlgoPerf Competition](https://mlsanity.com/paper/2502.15015/) (2025)
- [Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective](https://mlsanity.com/paper/2402.03496/) (2024)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [Faster Stochastic Variational Inference using Proximal-Gradient Methods with General Divergence Functions](https://mlsanity.com/paper/1511.00146/) (2015)
- [Homeomorphic-Invariance of EM: Non-Asymptotic Convergence in KL Divergence for Exponential Families via Mirror Descent](https://mlsanity.com/paper/2011.01170/) (2020)
- [Simplifying Momentum-based Positive-definite Submanifold Optimization with Applications to Deep Learning](https://mlsanity.com/paper/2302.09738/) (2023)

## Built on by (most cited first, 16 of 16)

- [Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization](https://mlsanity.com/paper/2606.16899/) (2026)
- [Preconditioned Norms: A Unified Framework for Steepest Descent, Quasi-Newton and Adaptive Methods](https://mlsanity.com/paper/2510.10777/) (2025)
- [The Quantization Benefits of Residual-Free Transformers](https://mlsanity.com/paper/2605.25880/) (2026)
- [SOAP-Bubbles: Structured Weight Uncertainty for Neural Networks](https://mlsanity.com/paper/2606.23357/) (2026)
- [DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers](https://mlsanity.com/paper/2602.02016/) (2026)
- [FOAM: Frequency and Operator Error-Based Adaptive Damping Method for Reducing Staleness-Oriented Error for Shampoo](https://mlsanity.com/paper/2606.02365/) (2026)
- [Convergence Rate Analysis of the AdamW-style Shampoo: Unifying One-Sided and Two-Sided Preconditioning](https://mlsanity.com/paper/2601.07326/) (2026)
- [Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization](https://mlsanity.com/paper/2605.06316/) (2026)
- [Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training](https://mlsanity.com/paper/2605.16184/) (2026)
- [Double Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation Loss](https://mlsanity.com/paper/2606.06418/) (2026)
- [Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage](https://mlsanity.com/paper/2605.26327/) (2026)
- [Exploiting weight-space symmetries for approximating curvature](https://mlsanity.com/paper/2606.00442/) (2026)
- [Rethinking Bregman Divergences in Kronecker-Factored Optimizers](https://mlsanity.com/paper/2606.00542/) (2026)
- [ShamAN-Q: Shampoo Augmented NanoQuant for Sub-1-bit LLM Weights](https://mlsanity.com/paper/2609.38521/) (2026)
- [Exact information accounting for SGD methods](https://mlsanity.com/paper/2610.00446/) (2026)
- [How Inner Optimizers Shape Replica Scaling](https://mlsanity.com/paper/2610.19064/) (2026)

## Abstract (arXiv)

Shampoo and its efficient variant, SOAP, employ structured second-moment estimations and have shown strong performance for training neural networks (NNs). In practice, however, Shampoo typically requires step-size grafting with Adam to be competitive, and SOAP mitigates this by applying Adam in Shampoo's eigenbasis -- at the cost of additional memory overhead from Adam in both methods. Prior analyses have largely relied on the Frobenius norm to motivate these estimation schemes. We instead recast their estimation procedures as covariance estimation under Kullback-Leibler (KL) divergence minimization, revealing a previously overlooked theoretical limitation and motivating principled redesigns. Building on this perspective, we develop KL-Shampoo and KL-SOAP, practical schemes that match or exceed the performance of Shampoo and SOAP in NN pre-training while achieving SOAP-level per-iteration runtime. Notably, KL-Shampoo does not rely on Adam to attain competitive performance, eliminating the memory overhead introduced by Adam. Across our experiments, KL-Shampoo consistently outperforms SOAP, Shampoo, and even KL-SOAP, establishing the KL-based approach as a promising foundation for designing structured methods in NN optimization. An implementation of KL-Shampoo/KL-SOAP is available at https://github.com/yorkerlin/KL-Methods