Summary (Overview)

• Identifies and fixes SOAP instability: The paper diagnoses a "stale preconditioner" effect in SOAP at large batch sizes, causing loss spikes and divergence. They propose per-step QR orthogonalization with current gradient inclusion and KL-divergence-based covariance estimation to eliminate these issues.

• Large batch scaling advantage: SOAP and Muon consistently outperform AdamW on multi-billion-parameter models (up to 72B total parameters) trained on trillions of tokens. Both maintain stability and quality at global batch sizes up to 100M tokens, while AdamW degrades.

• Fair optimizer comparison framework: The authors use an update-RMS matching framework to ensure fair learning rate transfer between AdamW, Muon, and SOAP, and empirically evaluate Muon's orthogonalization quality.

• Layer-wise distributed optimizer: A new distributed implementation compatible with Megatron-LM balances memory and hides communication overhead while avoiding approximations to optimizer computations, enabling efficient large-scale training.

• Open-source contribution: Release of the Emerging-Optimizers codebase (https://github.com/NVIDIA-NeMo/Emerging-Optimizers) containing implementations of all optimizers and experimental variants.

Introduction and Theoretical Foundation

The choice of optimizer is central to frontier model training, serving as the link between systems engineering and algorithm convergence. Optimizers dictate distributed execution constraints (optimizer states consume significantly more memory than parameters) and govern data efficiency, convergence rates, and generalization.

The paper addresses the tension between:

• Element-wise scalar optimizers (AdamW, RMSProp): computationally simple, scalable, but ignore correlational structure between gradients
• Higher-order methods (Shampoo, SOAP, Muon): capture loss landscape geometry, theoretically allowing faster convergence, but historically lack scalability

Key theoretical relationships: The paper establishes that Shampoo's preconditioned update reduces to Muon's polar factor when EMA parameters $\beta_1, \beta_2 = 0$:

$$(\mathbf{G}\mathbf{G}^T)^{-1/4}\mathbf{G}(\mathbf{G}^T\mathbf{G})^{-1/4} = UV^T$$

where $UV^T$ comes from the SVD $\mathbf{M}_t = U\Sigma V^T$. This shows SOAP/Shampoo-style whitening and Muon-style orthogonalization are closely related—both reduce the dominance of large singular directions.

SOAP update rule (combining Shampoo preconditioning with Adam-like updates in the eigenbasis):

$$u_t = Q_L \mathrm{Adam}(Q_L^T m_t Q_R) Q_R^T \tag{1}$$

Muon applies momentum then orthogonalizes via Newton-Schulz iterations to approximate the polar factor, avoiding explicit preconditioner storage.

Methodology

Large Batch Size Scaling for MoE Models

For MoE models, effective batch size per expert is:

$$B_{\mathrm{eff}}^{\mathrm{expert}} = B_{\mathrm{Global}} \times \frac{k}{N} \tag{2}$$

Since $k \ll N$ (e.g., top-8 gating with 256 experts gives ratio 0.03125), expert parameters remain in a low-batch regime while dense/shared parameters see the full global batch—making dense components the bottleneck for batch size scaling.

Square Root Scaling Rule

For fair comparison when changing batch sizes, the learning rate transfer rule is derived from preserving update variance:

$$\eta' = \eta \sqrt{\frac{B'}{B}} \tag{3}$$

Update-RMS Matching

To fairly compare optimizers, the update RMS norm is matched. For SOAP, since rotation matrices are orthonormal:

$$\|\Delta\mathbf{W}\|_F^2 = \|\mathbf{N}\|_F^2$$

For Muon, a correction factor depending on the EMA damping factor is applied:

$$\sqrt{\frac{1 - \beta_1}{1 + \beta_1}} \approx 0.2$$

Experimental Setup

Models: 8B dense GPT, 3B active/30B MoE, and 8B active/72B hybrid Mamba-Transformer MoE models, trained on 1T-3T token subsets of Nemotron-3 dataset.

Architecture configurations:

Learning rates: WSD (Warmup-Stable-Decay) schedule with minus-sqrt decay. Max LR ranges from $8.0 \times 10^{-4}$ to $1.6 \times 10^{-3}$.

Empirical Validation / Results

1. Muon vs. AdamW at Large Batch Sizes

• Muon achieves lower training loss with fewer spikes at baseline batch size (1× = 25M tokens)
• AdamW fails to scale effectively at larger batch sizes, hitting "critical batch size" limits
• Muon maintains robust convergence at 2× and 3× batch sizes (50M and 75M tokens)

Downstream benchmark results (Table 5): Muon at 3× batch size achieves higher MMLU (74.00 vs 73.38), better coding scores (Coding Avg: 61.89 vs 62.22), and better math scores (Math Avg: 81.47 vs 79.65) compared to AdamW at 1×.

2. SOAP Instability Diagnosis

• "Slingshot" instability: Infrequent eigenbasis refreshes (every 10 steps) combined with omitting current gradient from recalculation cause oscillating gradient norms and loss spikes in early training
• For an 8B dense model, this lag becomes catastrophic, leading to divergence
• Fix: Per-step QR updates + inclusion of current gradient in eigenbasis calculation eliminates loss spikes

3. KL-Shampoo Covariance Estimation

Integrating KL-divergence-regularized covariance estimation improves stability by reducing the condition number of Kronecker factors:

$$\kappa(S_{\mathrm{KL-Shampoo}}) = \sqrt{\kappa(S_{\mathrm{Shampoo}})} \tag{4}$$

This improves numerical stability of eigenbasis recomputation.

4. Muon vs SOAP Comparison

• KL-SOAP maintains a consistent, slight edge over Muon in CE loss during pretraining on Qwen-3-30B-A3B
• MOP (Muon with exact polar decomposition via SVD) also slightly outperforms standard Muon
• Both SOAP and Muon outperform AdamW at all batch sizes tested

5. Layer-Wise Distributed Optimizer

Key features:

1. Load balancing: Full parameter matrices distributed round-robin by size across GPUs
2. Parameter update: Each GPU independently updates assigned parameters
3. Overlapped all-gather: Asynchronous variable-sized all-gather-V collectives hide communication behind computation

Theoretical and Practical Implications

Theoretical Implications

• Preconditioner freshness: The paper demonstrates that stale preconditioners cause instability at scale, highlighting the importance of real-time gradient statistics in second-order methods
• Condition number reduction: KL-Shampoo's square-root reduction of condition number provides a principled mechanism for improving numerical stability in Kronecker-factored preconditioners
• Unified view: The relationship between Shampoo, SOAP, and Muon (reducing to polar factor when EMA=0) provides theoretical unification of spectral and preconditioning approaches

Practical Implications

• Large batch scaling: Muon and SOAP enable effective scaling to batch sizes where AdamW fails, potentially reducing communication overhead and improving GPU utilization
• Production deployment: The layer-wise distributed optimizer enables deployment of these algorithms at frontier model scales
• Memory tradeoffs: KL-SOAP recommended when memory is not limiting; Muon has lower footprint than AdamW (no second moments)
• Architectural sensitivity: Muon should be applied selectively (e.g., not to Mamba's conv1D weights), suggesting hybrid recipes for optimal performance

Conclusion

Main takeaways:

1. Both SOAP and Muon consistently outperform AdamW at multi-billion parameter scales
2. KL-SOAP is the most effective approach overall, recommended when memory is not limiting
3. Stability requires real-time eigenbasis updates and KL-divergence-based covariance estimation
4. The open-source Emerging-Optimizers codebase provides implementations for the research community

Future directions:

• SOAP TP support: Native Tensor Parallelism support and handling fused tensors
• Optimizer-driven memory layout: Refactoring DDP buffer allocation to be optimizer-aware
• Architecture co-design: Mapping how architectural components (MLA, LoRA) interact with full-rank preconditioning
• Scaling laws for batch size: Combining scaling-law predictions with systems-level batch size optimization
• More accurate orthogonalization: Addressing numerical challenges from "heavy-tail" small singular values, especially in BF16 Newton-Schulz iterations where machine precision is $2^{-7} = 0.0078125$ vs FP32's $2^{-23} \approx 1.19 \times 10^{-7}$