Beyond AdamW: The Next Generation of Optimizers for Training AI Models 📈

For years, AdamW has been the undisputed workhorse of large language model (LLM) pretraining. 🖥️

AdamW is an element-wise, first-order optimizer that updates every parameter independently: simple, reliable and easy to scale. It powers many frontier models today. However, as researchers push toward larger models, longer training runs, and massive batch sizes, AdamW is beginning to show its weaknesses.

NVIDIA researchers found at global batch sizes approaching 100 million tokens per step, AdamW begins to lose effectiveness. hence a new study from NVIDIA explores a fundamental question:

Can we build optimizers that are smarter than AdamW and allow future AI models to train faster, more efficiently, and more reliably?

The New Contenders: Muon and SOAP 🚀

Unlike AdamW, which treats parameters independently, Muon and SOAP introduce structure-aware optimization, methods that use information about the geometry and relationships within neural network weights.

🌀 Muon — Spectral Orthogonalization: Making Gradient Updates More Balanced

• Matrix-aware updates: Optimizes the geometry of gradient updates rather than scaling parameters independently like AdamW.

• Balanced learning directions: Uses Newton-Schulz iterations to orthogonalize momentum updates and balance singular values.

• Memory efficient: Removes the need for AdamW-style second-moment statistics, reducing optimizer memory.

• Better scaling: Maintains stable training at larger batch sizes and model scales.

🌐 SOAP — Preconditioned Optimization with Adam-Style Adaptivity

• Uses parameter structure: Applies Shampoo-inspired Kronecker preconditioning to capture row and column correlations in weight matrices.

• Better optimization space: Rotates gradients into an eigenbasis where directions become far more effective to update.

• Adam-compatible: Performs adaptive updates like AdamW but in a more informative coordinate system.

• Second-order benefits: Gains curvature awareness without the prohibitive cost of full second-order optimization.

📌 This shift towards new structure-aware optimization approaches would allow for:

1. More efficient AI development

2. Lower computational costs

3. More capable future models

👉 To make this transition possible, NVIDIA has released open-sourced implementations in Megatron-LM and a standalone Emerging-Optimizers library to be available to the public.

AdamW is unlikely to disappear and will remain a staple for small-to-medium training runs. Its simplicity and maturity make it an excellent choice for many applications.

╰┈➤ˎˊ˗ But for frontier-scale training? It still remains to be seen, where it'll go next.