Why Muon Outperforms Adam: A Curvature Perspective

TL;DR

Muon is the new star in LLM pretraining, often training models twice as fast as Adam. While we knew it worked, we didn't quite know why from a geometric standpoint. This paper reveals that Muon’s secret weapon is its ability to find "flatter" paths. Specifically, Muon minimizes Normalized Directional Sharpness (NDS), allowing it to take larger effective steps without hitting the "curvature walls" that slow down Adam and SGD.

The Problem: The Hidden Wall of Curvature

When we train LLMs, we want the largest possible loss decrease at every step. Using a second-order Taylor expansion, the loss change $\Delta L$ is roughly: $\Delta L\approx ext(GradientGain)-ext(CurvaturePenalty)$

Common wisdom suggests better optimizers get more "Gradient Gain" (aligning better with the steepest descent). However, this research finds that Muon and Adam actually have similar first-order gains. The real difference? Adam pays a much higher "tax" in the form of a second-order curvature penalty. Adam's updates accidentally point toward "sharp" directions in the landscape where the loss shoots back up quickly, canceling out the progress.

Methodology: Entering the Spectral Domain

The authors decompose the curvature penalty into two parts: the scale of the update and the Normalized Directional Sharpness (NDS).

• The Insight: Muon and Adam take similar-sized steps, but Muon’s direction is smarter.
• The Mechanism: LLM landscapes are heterogeneous—they have few directions with massive curvature (sharp) and many with low curvature (flat). Adam and Gradient Descent (GD) tend to over-invest update energy into those sharp directions because they often align with high gradients.
• The Fix: Muon utilizes spectral normalization (Orthogonalization). By setting all non-zero singular values of the momentum matrix to the same scale, Muon forces the update to distribute energy evenly across all curvature modes, effectively "ignoring" the bias toward sharp directions.

Figure 1: Comparison showing Muon (orange) incurring a significantly lower curvature penalty (c) than Adam (blue) at matched validation loss.

Key Findings & Experiments

1. The Imbalance Amplifier

The study used synthetic Zipf-PCFG data to control how "unbalanced" the training data is (a common trait of real-world text). As data becomes more imbalanced, the "sharpness" of the landscape increases. Crucially, the gap between Adam and Muon widens: Adam struggles significantly more with "jagged" landscapes created by rare tokens, while Muon remains robust.

2. Within-Layer vs. Cross-Layer Interactions

By breaking down the Hessian matrix, the authors found that Muon’s advantage is most prominent within individual layers. As training progresses, Muon's NDS advantage shifts almost entirely to within-layer blocks, meaning its matrix-aware normalization is perfectly suited for the layer-wise structure of Transformers.

Figure 2: Ratio analysis showing that while update norms (green) are similar, the NDS (cyan) is the primary driver of Muon's lower penalty.

3. Theoretical Proof: Balancing Energy

The authors proved via a "Structured Quadratic Model" that Muon achieves a smaller average NDS than GD by balancing energy. If the curvature heterogeneity is strong enough (which it is in LLMs), Muon is mathematically guaranteed to reach a lower loss than GD in the same number of steps by avoiding the "overshoot" in high-curvature directions.

Critical Analysis & Future Outlook

Takeaway: This work provides a rigorous geometric foundation for Muon. It tells us that we should look beyond just "gradient alignment" and start focusing on "curvature avoidance" when designing next-generation optimizers.

Limitations: The study focuses primarily on Causal LLMs (NanoGPT). Whether these insights apply to Diffusion Models or other architectures where traditional Adam is still king remains an open question.

Future Work: The research suggests that any Muon variant that further shifts updates toward "flatter" directions (lower complexity regions) could yield even higher efficiency. We are likely entering an era where "spectral awareness" becomes a standard feature of high-performance optimizers.

Conclusion

Muon isn't just a heuristic trick; it is a mathematically sound way to navigate the treacherous, sharp-ridged landscapes of large-scale neural networks. By "flattening" the update through spectral normalization, it avoids the curvature tax that has been slowing down our training for years.