# SOAP: Improving and Stabilizing Shampoo using Adam

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

## Key facts

- Cited by 192 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 47 references that are in the map.
- First posted on arXiv 17 Sep 2024.

## Builds on (most cited first, 30 of 37)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [The Road Less Scheduled](https://mlsanity.com/paper/2405.15682/) (2024)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [Full Parameter Fine-tuning for Large Language Models with Limited Resources](https://mlsanity.com/paper/2306.09782/) (2023)
- [Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model](https://mlsanity.com/paper/1907.04164/) (2019)
- [Resolving Discrepancies in Compute-Optimal Scaling of Language Models](https://mlsanity.com/paper/2406.19146/) (2024)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [Preconditioned Stochastic Gradient Descent](https://mlsanity.com/paper/1512.04202/) (2015)
- [Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures](https://mlsanity.com/paper/2311.00636/) (2023)
- [No Train No Gain: Revisiting Efficient Training Algorithms For Transformer-based Language Models](https://mlsanity.com/paper/2307.06440/) (2023)
- [Deconstructing What Makes a Good Optimizer for Language Models](https://mlsanity.com/paper/2407.07972/) (2024)
- [Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks](https://mlsanity.com/paper/1811.12019/) (2018)
- [AdaLomo: Low-memory Optimization with Adaptive Learning Rate](https://mlsanity.com/paper/2310.10195/) (2023)
- [Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective](https://mlsanity.com/paper/2402.03496/) (2024)

## Built on by (most cited first, 30 of 192)

- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Practical Efficiency of Muon for Pretraining](https://mlsanity.com/paper/2505.02222/) (2025)
- [MARS: Unleashing the Power of Variance Reduction for Training Large Models](https://mlsanity.com/paper/2411.10438/) (2024)
- [ASGO: Adaptive Structured Gradient Optimization](https://mlsanity.com/paper/2503.20762/) (2025)
- [Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective](https://mlsanity.com/paper/2502.00604/) (2025)
- [Deep Learning is Not So Mysterious or Different](https://mlsanity.com/paper/2503.02113/) (2025)
- [Cautious Optimizers: Improving Training with One Line of Code](https://mlsanity.com/paper/2411.16085/) (2024)
- [Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training](https://mlsanity.com/paper/2505.13738/) (2025)
- [Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data](https://mlsanity.com/paper/2502.04664/) (2025)
- [Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise](https://mlsanity.com/paper/2506.04192/) (2025)
- [Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful](https://mlsanity.com/paper/2507.07101/) (2025)
- [Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal?](https://mlsanity.com/paper/2511.00674/) (2025)
- [Convergence of Muon with Newton-Schulz](https://mlsanity.com/paper/2601.19156/) (2026)
- [Structured Preconditioners in Adaptive Optimization: A Unified Analysis](https://mlsanity.com/paper/2503.10537/) (2025)
- [COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs](https://mlsanity.com/paper/2502.17410/) (2025)
- [The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training](https://mlsanity.com/paper/2502.19002/) (2025)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates](https://mlsanity.com/paper/2509.02981/) (2025)
- [Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales](https://mlsanity.com/paper/2512.05620/) (2025)
- [SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training](https://mlsanity.com/paper/2412.13148/) (2024)
- [Moonshine: Speech Recognition for Live Transcription and Voice Commands](https://mlsanity.com/paper/2410.15608/) (2024)
- [SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training](https://mlsanity.com/paper/2505.24749/) (2025)
- [A Stable Whitening Optimizer for Efficient Neural Network Training](https://mlsanity.com/paper/2506.07254/) (2025)
- [The Newton-Muon Optimizer](https://mlsanity.com/paper/2604.01472/) (2026)
- [The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton](https://mlsanity.com/paper/2510.09378/) (2025)
- [Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning](https://mlsanity.com/paper/2603.09697/) (2026)
- [MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration](https://mlsanity.com/paper/2603.28254/) (2026)
- [Muon2: Boosting Muon via Adaptive Second-Moment Preconditioning](https://mlsanity.com/paper/2604.09967/) (2026)
- [NeoBERT: A Next-Generation BERT](https://mlsanity.com/paper/2502.19587/) (2025)

## Abstract (arXiv)

There is growing evidence of the effectiveness of Shampoo, a higher-order preconditioning method, over Adam in deep learning optimization tasks. However, Shampoo's drawbacks include additional hyperparameters and computational overhead when compared to Adam, which only updates running averages of first- and second-moment quantities. This work establishes a formal connection between Shampoo (implemented with the 1/2 power) and Adafactor -- a memory-efficient approximation of Adam -- showing that Shampoo is equivalent to running Adafactor in the eigenbasis of Shampoo's preconditioner. This insight leads to the design of a simpler and computationally efficient algorithm: ShampoO with Adam in the Preconditioner's eigenbasis (SOAP). With regards to improving Shampoo's computational efficiency, the most straightforward approach would be to simply compute Shampoo's eigendecomposition less frequently. Unfortunately, as our empirical results show, this leads to performance degradation that worsens with this frequency. SOAP mitigates this degradation by continually updating the running average of the second moment, just as Adam does, but in the current (slowly changing) coordinate basis. Furthermore, since SOAP is equivalent to running Adam in a rotated space, it introduces only one additional hyperparameter (the preconditioning frequency) compared to Adam. We empirically evaluate SOAP on language model pre-training with 360m and 660m sized models. In the large batch regime, SOAP reduces the number of iterations by over 40% and wall clock time by over 35% compared to AdamW, with approximately 20% improvements in both metrics compared to Shampoo. An implementation of SOAP is available at https://github.com/nikhilvyas/SOAP.