SOAP: Improving and Stabilizing Shampoo using Adam

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 LLM 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.

• The paper demonstrates that SOAP reduces training iterations by over 40% by performing updates in a rotated eigenbasis that combines Shampoo's second-order information with Adam's adaptive learning rates.
• It introduces a dual-layer update mechanism that maintains preconditioner matrices while executing efficient AdamW updates, simplifying hyperparameter tuning with a single additional parameter.
• Experimental results show SOAP outperforms AdamW and the original Shampoo, achieving up to 20% fewer iterations and 15% less wall clock time through optimized preconditioning frequency.

SOAP: Improving and Stabilizing Shampoo using Adam

Introduction

The paper "SOAP: Improving and Stabilizing Shampoo using Adam" proposes a novel optimization algorithm aimed at improving the efficiency of LLM training. This algorithm, referred to as SOAP (ShampoO with Adam in the Preconditioner's eigenbasis), builds on the efficacy of the Shampoo optimizer by addressing its computational overhead and extending its applicability through leveraging the adaptive qualities of Adam in a rotated coordinate space. The synthesis of Shampoo's second-order information with Adam's adaptive learning rates, achieved via SOAP, offers significant reductions in both the number of training iterations and wall clock time required for LLM pre-training, when compared to traditional methods like AdamW and the original Shampoo.

Technical Contributions

One of the paper's primary contributions is establishing a formal connection between Shampoo and Adafactor by demonstrating their equivalence in the eigenspace provided by Shampoo's preconditioner. This insight underpins the derivation of SOAP, which effectively executes AdamW updates in this eigenbasis—a more computationally feasible and efficient method. Crucially, SOAP introduces a single additional hyperparameter, preconditioning frequency, simplifying hyperparameter tuning compared to the multiple parameters required by Shampoo.

Figure 1: Comparing performance of tuned runs for AdamW, Shampoo, and SOAP, highlighting SOAP's efficiency advantages.

SOAP's performance was empirically validated on LLM tasks, showcasing a reduction in the required number of iterations by over 40% and a decrease in wall clock time by over 35% compared to AdamW. These improvements are achieved without sacrificing training outcomes, further emphasizing SOAP's potential in real-world applications.

Algorithmic Design

The SOAP algorithm is underpinned by a dual-layer update mechanism: it maintains Shampoo-like second-order information while benefiting from Adam-like adaptivity. Specifically, SOAP performs updates in a rotated space defined by the eigenvectors of the Shampoo preconditioner matrices, $L$ and $R$. These matrices are updated less frequently to reduce computational cost, while Adam's running averages are continually adapted, ensuring robustness across different preconditioning frequencies.

Figure 2: Precise efficiency benefits of SOAP over AdamW and Shampoo for different model and batch sizes.

Maintaining $L$ and $R$ requires storing and manipulating additional matrices, but the resultant benefits in optimization efficiency justify this overhead. The trade-off between computational savings and memory usage can be managed through various implementation strategies, including the use of reduced precision for certain operations.

Experimental Analysis

SOAP's efficacy was demonstrated through extensive experiments on LLM pre-training tasks, involving models with 360 million and 660 million parameters. Comparisons against AdamW and DistributedShampoo highlighted significant efficiency improvements:

• Up to 20% fewer iterations and 15% less wall clock time than Shampoo.
• A consistent performance advantage over a range of batch sizes and preconditioning frequencies.

These results underscore SOAP's robustness and adaptability, particularly in scenarios with large batch sizes where second-order methods are traditionally disadvantaged due to higher per-iteration costs.

Figure 3: Performance of SOAP variants focused on optimizing space or time usage, offering insights into potential efficiency improvements.

Implementation Considerations

Implementing SOAP involves careful attention to computational details, such as eigenvector computation frequency and the precision of matrix operations. The algorithm benefits from modern hardware configurations, such as distributed computing environments where the overhead of matrix operations can be distributed across multiple GPUs. Additionally, the paper suggests potential for further optimization through one-sided projections or low-rank approximations for matrix updates, which could yield further savings in both time and space complexities.

Conclusion

The SOAP algorithm offers a compelling advancement in the field of deep learning optimization, particularly for LLM training tasks. By marrying the strengths of Shampoo's second-order preconditioning with Adam's adaptivity, SOAP addresses key challenges in optimization efficiency. Future work could explore enhancements in SOAP's implementation, such as integrating distributed computation or exploring alternative low-precision storage techniques, to further extend its applicability and resource efficiency in large-scale AI applications.