SOAP: Improving and Stabilizing Shampoo using Adam

TL;DR

Shampoo offers strong optimization performance but with high computational cost. The authors establish a formal link between Shampoo (power 1/2) and Adafactor in Shampoo’s eigenbasis, then propose SOAP, a streamlined variant that runs AdamW in the rotated space to retain second-moment updates efficiently. Across 360m and 660m language-model pretraining tasks, SOAP achieves substantial speedups over AdamW and Shampoo, with improved robustness to preconditioning frequency and fewer hyperparameters. The work also analyzes space and time overhead, and outlines practical refinements (one-sided variants, low-precision options) and future directions for broader applicability. Overall, SOAP provides a practical, scalable alternative for second-order-inspired optimization in large-scale neural network training.

Abstract

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: $\textbf{S}$hampo$\textbf{O}$ with $\textbf{A}$dam in the $\textbf{P}$reconditioner'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.

Figures (7)

• Figure 1: Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo distributedshampoo implementation) and SOAP. In left and middle figures, Shampoo and SOAP use a preconditioning frequency of 10. The "shorter LR schedule" plot is where we tuned the cosine decay so as to achieve the same terminal performance as AdamW. There we observe a $\geq 40\%$ reduction in the number of iterations and a $\geq 35\%$ reduction in wall clock time compared to AdamW, and approximately a 20% reduction in both metrics compared to Shampoo. In the right figure we ablate preconditioning frequency and observe a slower degradation of performance of SOAP as compared to Shampoo. See \ref{['sec:language']} for a discussion of experimental results and ablation of batch size and \ref{['sec:methods']} for experimental methodology.
• Figure 2: Precise efficiency benefits of SOAP over AdamW and Shampoo for 360m (at 256k and 2m batch size) and 660m (at 2m batch size) model. For the precise methodology, refer to \ref{['sec:methods']}.
• Figure 3: Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo distributedshampoo implementation) and SOAP. Shampoo and SOAP use preconditioning frequency of 10. We observe a $\geq 40\%$ reduction in the number of iterations and a $\geq 35\%$ reduction in wall clock time compared to AdamW, and approximately a $20\%$ reduction in both metrics compared to Shampoo. See \ref{['fig:main']} for 660m results, \ref{['sec:freq', 'sec:critical']} for ablations of preconditioning frequency and batch size respectively, and \ref{['sec:methods']} for detailed calculation of efficiency improvement and experimental methodology.
• Figure 4: (left) Comparing the critical batch size of AdamW vs SOAP. We can see that SOAP improves the critical batch size, by being much closer to the ideal linear scaling with batch size as compared to AdamW. (right) Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo distributedshampoo implementation) and SOAP for token batch size of 256k. Shampoo and SOAP use preconditioning frequency of 80. We observe a $\geq 25\%$ reduction in the number of iterations compared to AdamW, and approximately a 10% reduction compared to Shampoo. See \ref{['fig:prec_overhead']} (right) for wall-clock time improvement and \ref{['sec:methods']} for detailed calculation of efficiency improvement.
• Figure 5: Performance comparison of SOAP and Adam for longer duration training runs.

Theorems & Definitions (2)

• Claim 1
• proof