PaperlayerView exact revision on arXiv

Source-linked AI summary

Why Muon Outperforms Adam: A Curvature Perspective

Shuche Wang, Fengzhuo Zhang, Jiaxiang Li, Dirk Bergemann, Zhuoran Yang

arXiv:2606.04662v1cs.LGcs.AI

TL;DR

The paper asks what landscape property explains Muon’s advantage over Adam. Using curvature analysis, controlled-imbalance experiments, layer-wise decomposition, and quadratic theory, it finds that Muon’s lower NDS drives a smaller curvature penalty and larger one-step loss decrease.

  • Problem

    The landscape property supporting Muon’s advantage over Adam remains unclear.

  • Method

    The paper combines second-order loss analysis, NDS decomposition, controlled-imbalance and layer-wise experiments, and heterogeneous-curvature quadratic modeling.

  • Results

    Muon’s lower NDS produces a smaller curvature penalty and larger one-step loss decrease than Adam, while quadratic analysis establishes lower average NDS than GD.

  • Takeaways & Limitations

    Muon’s advantage is associated with update directions that encounter less curvature rather than with smaller update norms.

  • Takeaways & Limitations

    The study focuses on causal LLMs, while its quadratic analysis relies on four empirically verified assumptions about gradient and Hessian structure.

Abstract

from arXiv · show

Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective. First, we apply a second-order Taylor approximation to the training landscape and show that Muon achieves a larger one-step loss decrease than Adam at matched validation loss. The two optimizers have comparable first-order gains, but Muon consistently incurs a smaller second-order curvature penalty. Second, we decompose this curvature penalty into the squared update norm and Normalized Directional Sharpness (NDS). We find that Muon and Adam have comparable update norms, so Muon's smaller curvature penalty is driven by lower NDS, not update scale. Third, we study how training data and model structure shape Muon's NDS advantage. Using Zipf-Probabilistic Context-Free Grammar (PCFG) data with controlled imbalance, we show that data imbalance amplifies Muon's NDS advantage over Adam. A within-/cross-layer decomposition further shows that, in the middle and late stages of training, Muon's lower NDS is mainly sustained by smaller within-layer curvature. Beyond empirical evidence, we analyze stylized quadratic problems with heterogeneous curvature and gradient alignment toward high-curvature modes. We prove that Muon attains a smaller average NDS than GD by balancing update energy across curvature groups; when curvature heterogeneity is sufficiently strong, this also yields lower local quadratic loss after the same number of steps.

1 Introduction

Muon is a matrix-aware optimizer that can train LLMs up to about 2× faster than Adam. This section frames the paper’s goal as identifying the optimization-landscape property behind that advantage.

  • 1 Introduction: Muon spectrally normalizes the gradient momentum matrix, setting its nonzero singular values to the same scale.This design exploits the matrix structure of model parameters.
  • 1 Introduction: Up to about 2× faster training than Adam is reported for Muon across various LLM pretraining model scales.The comparison concerns Large Language Model pretraining.
  • 1 Introduction: The paper takes a first step toward explaining Muon’s advantage over Adam through the optimization landscape.It asks which landscape property supports Muon’s advantage.

2. How do pretraining factors, such as the training data and model structure, influence this property?

Muon’s NDS advantage over Adam grows with training-data imbalance, while quadratic analyses attribute its lower average NDS to balancing update energy across heterogeneous curvature groups. Together, these findings identify lower NDS and its reduced second-order curvature cost as a mechanism behind Muon’s larger one-step loss decrease.

  • Training data: Muon’s NDS advantage over Adam becomes larger as Zipf-PCFG training data becomes more imbalanced.The experiments use synthetic data with controlled imbalance levels.
  • Model structure: Under heterogeneous curvature and gradient alignment toward high-curvature directions, Muon has lower average NDS than GD by balancing update energy across curvature groups.The analysis uses stylized quadratic optimization problems designed to mirror LLM-training landscapes.
  • Mechanism: The empirical and theoretical findings identify a concrete mechanism behind Muon’s advantage over Adam.This mechanism centers on lower normalized directional sharpness and the resulting reduction in curvature penalty.
  • Mechanism: Muon’s spectrally normalized update direction yields lower NDS than Adam, producing a smaller second-order curvature cost and larger one-step loss decrease.The two optimizers have comparable first-order gains, so the difference is primarily attributed to curvature cost.

2 Related Work

Prior work has established Muon’s empirical gains and examined its practical, algorithmic, scalability, convergence, and geometric properties. Related literature also covers structured or matrix-aware updates, Adam and adaptive methods, and the roles of sharpness and curvature in optimization and generalization.

  • Muon and structured matrix updates: Recent studies examine Muon’s empirical behavior, algorithms, scalability, convergence properties, and geometric interpretation of its orthogonalized updates.Jordan et al. (2024b) established Muon’s gains on matrix-valued parameters.
  • Muon and structured matrix updates: A broader structured-update perspective connects Muon with orthogonalized updates and matrix- or block-aware preconditioning methods.This literature places Muon within a larger geometric family of structured optimization methods.
  • Adam and adaptive optimization: The paper’s Adam comparison builds on literature spanning Adam, adaptive gradients, decoupled regularization, and convergence theory.The cited foundations include Kingma and Ba (2014) and Duchi et al. (2011).
  • Sharpness, curvature, and optimization geometry: Related work studies sharpness and curvature in deep-learning optimization and generalization, including parameterization sensitivity of naive sharpness measures.This line connects large-batch training, minima geometry, and generalization.

3 Preliminaries

This section introduces Adam’s coordinate-wise moment normalization and Muon’s matrix-structure-aware spectral normalization, then establishes notation for gradients, Hessians, and matrix perturbations. Muon’s update direction is scale-invariant and can be computed efficiently with Newton–Schulz iterations.

  • Adam: Adam normalizes parameter updates coordinate-wise using exponential moving averages of stochastic-gradient first and second moments.After bias correction, square roots and division are applied element-wise.
  • Muon: Muon explicitly exploits matrix-gradient spectral structure by accumulating momentum and spectrally normalizing it through its singular vectors.For momentum matrix B_t = U_tS_tV_t^⊤, Muon sets O_t = U_tV_t^⊤.
  • Muon: Muon’s update direction is scale-invariant, so multiplying B_t by any positive scalar does not change it; Newton–Schulz iterations efficiently approximate the spectral normalization.This avoids computing an exact singular value decomposition in practice.
  • Notation and assumptions: The paper assumes a twice continuously differentiable empirical training loss and defines matrix gradients, Hessian actions, vectorized Hessian representations, Frobenius inner products, and norms.For perturbation Z, the Hessian operator is defined by differentiating the gradient along W + ϵZ.

4 Main Results

Muon’s advantage over Adam is consistently linked to a smaller second-order curvature penalty rather than a larger first-order gain. This curvature advantage is driven by lower normalized directional sharpness, amplified by data imbalance, and increasingly associated with within-layer curvature during training.

  • 4.1 One-step loss decrease: Muon achieves a larger one-step loss decrease than Adam at matched validation-loss levels, primarily because its second-order curvature cost is smaller.The optimizers have comparable first-order decreases, while Muon consistently incurs a smaller curvature penalty.
  • 4.2 Curvature decomposition: Muon and Adam have comparable update norms, so Muon’s smaller curvature penalty is driven by notably lower Normalized Directional Sharpness (NDS).The average NDS ratio is 1.76, and the curvature-penalty and NDS ratios track closely.
  • 4.3 Data imbalance: As Zipf imbalance increases from s=0 to s=1, Adam’s normalized NDS rises from 1.63 to 2.38, whereas Muon’s increases from 1.00 to 1.25.Imbalance increases NDS for both optimizers but widens the gap between them more strongly for Adam.
  • 4.4 Layer-wise curvature: Muon’s within-layer sharpness share rises from about 14% early to about 44% later in training, while Adam’s fluctuates around 30%.The increasing within-layer contribution helps keep Muon’s full-model NDS low in the middle and late stages.
  • 4.4 Layer-wise curvature: Both within-layer and cross-layer directional-sharpness components remain smaller for Muon than for Adam throughout training.Muon’s within-layer component becomes increasingly dominant as training proceeds, while Adam’s sharpness composition remains comparatively stable.

5 A Case Study of Structured Matrix-Block Quadratic Models

This section develops a structured matrix-block quadratic model with low-Kronecker-rank, approximately jointly diagonalizable Hessians and heterogeneous curvature. Under these assumptions, Muon has lower finite-horizon averaged NDS than GD and, with sufficient curvature heterogeneity, achieves lower quadratic loss.

  • Quadratic Model: The local quadratic model is Q(Y) = L(W0) −⟨G, Y⟩+ 1/2 · ⟨Y, H[Y ]⟩ for a matrix-block update Y.The model uses the loss, gradient, and Hessian operator at a fixed parameter W0.
  • Structural Assumptions: The Hessian is assumed to have small Kronecker rank, with factors that approximately share common orthogonal eigenbases.Empirical results support low-rank Kronecker approximations and approximate simultaneous diagonalization in attention-matrix Hessians.
  • Curvature Heterogeneity: Positive paired curvatures are modeled as heterogeneous two-level spectra, with a high-curvature group occupying α = m/q < 1/2.Observed paired curvatures are strongly long-tailed, with values spanning more than six orders of magnitude.
  • Theorem 5.5: For every finite horizon T ≥1, Muon has a smaller finite-horizon averaged NDS than GD.The proof attributes this to Muon assigning exactly an α fraction of update energy to high-curvature modes, while GD’s time-averaged weight is always larger than α.
  • Theorem 5.5: Under ρ + 1 > 1/α > 1 + σH/σL, Muon achieves lower quadratic loss than GD for every finite horizon T ≥1.Here ρ = wH/wL; the mechanism is that Muon spreads update energy across active singular modes instead of repeatedly favoring sharp directions.

6 Conclusion · A Experimental Details

The paper explains Muon’s advantage over Adam through lower Normalized Directional Sharpness, which reduces curvature penalties and improves one-step loss decrease. It further links this advantage to data imbalance, layer-wise interactions, and quadratic-model theory.

  • 6 Conclusion: Muon’s lower NDS reduces its curvature penalty and yields a larger one-step loss decrease than Adam.This is the paper’s central curvature-based explanation of Muon’s superiority.
  • 6 Conclusion: Data imbalance shapes Muon’s NDS advantage over Adam.The conclusion identifies data imbalance as one factor influencing the optimizer comparison.
  • 6 Conclusion: Layer-wise interactions shape Muon’s NDS advantage over Adam.The conclusion highlights model layer structure as another factor affecting the advantage.
  • 6 Conclusion: A quadratic-model analysis theoretically establishes Muon’s advantages in NDS and one-step loss decrease.The theoretical analysis complements the empirical curvature findings.
  • 6 Conclusion: The work takes a first step toward understanding Muon’s superiority over Adam from a curvature perspective.The conclusion characterizes the study as an initial contribution rather than a complete account.
  • 6 Conclusion: The provided conclusion passage identifies a limitation of the work but does not specify it fully.The passage ends mid-sentence after introducing the limitation.

A.1 FineWeb Main-Text Experiments … A.2.2 Experiment Setup

The experiments combine a 124M-parameter NanoGPT trained on FineWeb-10B with a controlled synthetic Zipf-PCFG setup using a smaller 9M-parameter model. The synthetic corpus varies token-distribution imbalance, topic persistence, grammatical structure, and independently configurable training and validation Zipf exponents.

  • A.1 FineWeb Main-Text Experiments: FineWeb experiments use a 124M-parameter NanoGPT with 12 Transformer layers, 12 attention heads, hidden dimension 768, and a 50,257-token GPT-2 vocabulary.Training uses FineWeb-10B with sequence length 1024.
  • A.1 FineWeb Main-Text Experiments: Muon is applied to all matrix parameters except token embeddings and the language-model head, with momentum µ = 0.95 and no weight decay.The momentum coefficient warms up linearly from 0.85 to 0.95 over the first 300 steps.
  • A.1 FineWeb Main-Text Experiments: Directional sharpness is computed every 500 steps through Hessian–vector products, using experiments run on 4 A100 GPUs with 80 GB memory each.
  • A.2.1 Dataset Construction: The synthetic Zipf-PCFG corpus combines latent topic structure with explicit grammatical constraints and independently configurable training and validation Zipf exponents.Its main comparison fixes the validation distribution at s = 0, removing rank-based Zipf bias from validation token emission.
  • A.2.1 Dataset Construction: The synthetic vocabulary contains 4,411 tokens partitioned into C = 20 grammatical classes, plus a separate EOS token slot.
  • A.2.1 Dataset Construction: The corpus instantiates K = 30 latent topics with sparse class-token affinities and strong within-document topic persistence from self-transition probability Tkk = 0.85.Topic affinities use αϕ = 0.3 and the topic prior uses απ = 1.0.
  • A.2.1 Dataset Construction: Increasing imbalance level s concentrates token probability on higher-ranked tokens, producing heavier-tailed class-conditional distributions.At s = 0, emission follows only the topic-specific base distribution ϕk,c.
  • A.2.2 Experiment Setup: The Zipf-PCFG experiments use a 9M-parameter NanoGPT with 4 Transformer layers, 4 attention heads, model dimension 256, sequence length 1,024, and effective batch size 8,192 tokens per step.Validation uses 2,457,600 tokens with sequence length 1,024; total training is 10,000 optimization steps.

B Additional Experimental Results to Section 4.2

Across training, Muon consistently has lower NDS than Adam while maintaining comparable update norms, and its NDS advantage remains robust under different alignment choices. The mean NDS ratio is 2.94 when aligned by training step, versus 1.76 at matched validation loss.

  • Muon maintains consistently smaller NDS than Adam throughout the entire training trajectory, from early to late training.
  • The Adam-to-Muon update norm ratio remains close to 1 at every training step, indicating comparable update norms.
  • 2.94 is the mean Adam-to-Muon NDS ratio over matched training steps, exceeding the 1.76 ratio reported at matched validation loss.The difference arises because Muon reaches lower validation loss than Adam at the same training step, so step alignment compares different optimization stages.
  • The qualitative conclusion of Observation 2 is robust to whether comparisons are aligned by training step or matched validation loss.Matched-validation-loss comparison assesses per-step curvature differences, while step-aligned comparison provides a complementary assessment.

C Layerwise NDS contribution … D.3 Curvature Heterogeneity (Assumption 5.3)

Muon’s within-layer NDS advantage is concentrated in boundary Transformer layers, while the appendix verifies the curvature assumptions through low-rank Hessian structure, approximate diagonalization, and strongly heterogeneous curvature spectra. These analyses provide the empirical and structural basis for the paper’s curvature perspective.

  • C Layerwise NDS contribution: 70% of Muon’s within-layer NDS gap versus Adam comes from boundary layers L1 and L12, compared with 28% from L8–L11 and 2% from L2–L7.The boundary layers interact most directly with token embeddings and output logits, where the data distribution has the most immediate influence.
  • D Empirical Verification Details for Section 5: The empirical verification uses the last Zipf-PCFG checkpoint trained with Muon and computes dense block Hessians for four 256×256 attention matrices.Each dense Hessian has dimension 65,536 × 65,536.
  • D.1 Effective low-Kronecker-rank block Hessian (Assumption 5.1): ξ(4) equals 0.75, 0.95, 0.87, and 0.71 for WQ, WK, WV, and WO, respectively, supporting a rank-4 Kronecker approximation.The WV residual is small and approximately diagonal, consistent with Assumption 5.1; analogous heatmaps show similarly small residuals for WQ, WK, and WO.
  • D.2 Simultaneous Diagonalization (Assumption 5.2): At rank 4, the Kronecker approximation captures ξ(4) ≥0.87 of Frobenius energy for three of four attention matrices while preserving computational tractability.This choice keeps the matrix-pair count small enough for reliable JADE joint diagonalization.
  • D.2 Simultaneous Diagonalization (Assumption 5.2): JADE joint diagonalization leaves only about 11–15% of squared Frobenius energy off-diagonal, indicating substantial approximate simultaneous diagonalization.The reported objective value is 0.845 for the aligned matrix family.
  • D.3 Curvature Heterogeneity (Assumption 5.3): Among the top-128 curvature modes, 88 are positive, and the spectrum is strongly heterogeneous, with w1/w88 ≈2.59 × 10^6.The remaining entries are negligible: the 128-th largest |wi| is less than 10−6 of |w1|, while w1 accounts for a large fraction of the positive mode trace.

D.4 Gradient alignment (Assumption 5.4) · D.5 Synthetic GD vs. Adam vs. Muon under the structured quadratic model

The gradient-alignment diagnostic supports the structured quadratic model as a reasonable idealization: most gradient energy lies in the positive subspace, and paired-diagonal curvature dominates despite off-diagonal components. Synthetic structured-quadratic experiments then show Muon substantially outperforms GD and Adam in normalized directional sharpness and optimal-step loss decrease.

  • D.4 Gradient alignment (Assumption 5.4): JADE-based mode projections and top positive Hessian eigenpairs yield nearly identical cumulative energy ratios.This verifies the alignment diagnostic using the same checkpoint and dense Hessian, with q = 88 positive eigenpairs.
  • D.4 Gradient alignment (Assumption 5.4): ζ(q) = 0.871 indicates that most gradient energy lies in the positive row/column subspace.The diagnostic supports the positive-subspace component of Assumption 5.4, but not the stronger paired-diagonal gradient form.
  • D.4 Gradient alignment (Assumption 5.4): The paired-diagonal positive component accounts for 88.0% of ⟨G, H[G]⟩, while the off paired-diagonal component accounts for 12.0%.Thus, curvature cost remains primarily controlled by the paired-diagonal mechanism even though the gradient is not exactly paired-diagonal.
  • D.5 Synthetic GD vs. Adam vs. Muon under the structured quadratic model: Synthetic experiments use random structured-quadratic instances with defaults d1 = d2 = 256, q = 88, αw = 1.3, and ασ = 0.5.Only q = 88 modes carry positive curvature, matching the effective dimensionality observed empirically in Appendix D.3.
  • D.5 Synthetic GD vs. Adam vs. Muon under the structured quadratic model: 0.09×: Muon’s NDS is roughly 11× smaller than GD’s, whereas Adam’s is 0.50× GD’s.The ratios are normalized by GD’s NDS value; all optimizers start at Y0 = 0 and use exact line-search step sizes.
  • D.5 Synthetic GD vs. Adam vs. Muon under the structured quadratic model: 5.60×: Muon’s optimal-step quadratic-model loss decrease is larger than GD’s, while Adam’s advantage over GD is 1.22×.Adam and GD exhibit similar NDS and loss-decrease profiles in these synthetic experiments.

E Proof of Theorem 5.5 … F.2 Proof of Proposition E.2

The proof reduces the quadratic comparison to heterogeneous high- and low-curvature modes, derives exact Muon and GD dynamics, and shows Muon has lower average NDS and suboptimality under stated conditions. Supporting proofs establish the exact line-search formula and GD’s alternating energy-share mechanism.

  • E Proof of Theorem 5.5: The model uses positive modes with two heterogeneous curvature groups, wH > wL > 0, while higher-curvature modes carry larger gradient energy.The groups contain m high-curvature modes and q − m low-curvature modes, with σH > σL > 0.
  • F Proofs of Supporting Propositions and Lemmas: The proof tracks both optimizers within the positive-mode subspace using scalar residual coefficients and exact line-search updates.Induction shows iterates, residuals, and update increments remain in that subspace.
  • F.1 Proof of Lemma E.1: The exact line-search step size follows by expanding the quadratic loss along any direction and optimizing the resulting linear-minus-quadratic decrease.This lemma supplies the step-size formula used throughout the Muon and GD proofs.
  • E Proof of Theorem 5.5: Muon weights all modes equally through spectral normalization, whereas GD concentrates its direction on the group with greater residual energy.Consequently, Muon’s NDS is constant across steps, while GD’s NDS varies with its high-curvature energy share.
  • F.2 Proof of Proposition E.2: GD’s high-curvature share alternates between p and 1 − p, so its NDS alternates between p wH + (1 − p) wL and (1 − p) wH + p wL.This oscillation underlies the average-NDS comparison with Muon.
  • E Proof of Theorem 5.5: For every T ≥ 1, Muon’s suboptimality is ab/[2(a + b)] Γ^2(T − 1)d0^2, establishing linear convergence to Y⋆.The proof then compares this closed form with GD’s contraction and concludes Muon’s gap is strictly smaller at every horizon.
  • E Proof of Theorem 5.5: Under ρ + 1 > 1/α > 1 + σH/σL, Muon has lower loss after the first step and a smaller multiplicative contraction coefficient than GD thereafter.Here ρ = wH/wL.
  • F.2 Proof of Proposition E.2: GD’s high-curvature energy alternates because exact line search overshoots high-curvature modes and undershoots low-curvature modes.The resulting opposite contraction signs make energy shift between the two groups across steps.

F.3 Proof of Proposition E.3 … F.6 Proof of Proposition E.6

The appendices prove the propositions by reducing Muon and GD dynamics to scalar residuals and contraction factors in a two-curvature-group quadratic model. Under the stated curvature, imbalance, and alignment conditions, the proofs establish Muon’s residual contraction and the comparison R > Γ^2.

  • F.3 Proof of Proposition E.3: Muon’s suboptimality is expressed through residual coefficients, reducing its analysis to scalar residual dynamics across curvature groups.The residual representation tracks suboptimality through r^Muon_i,t and then introduces scaled residuals for the dynamics.
  • F.3 Proof of Proposition E.3: Under Assumption 5.3, Muon’s scaled residuals are constant within the high- and low-curvature groups and contract by Γ at each subsequent step.The groupwise forms are c^Muon_H,t and c^Muon_L,t, with the inter-group gap contracting by Γ.
  • F.4 Proof of Proposition E.4: GD’s suboptimality evolves through alternating group proportions, with P_t+1 = 1 − P_t and S_t+1 = C(P_t)S_t.The proof uses the symmetry of the suboptimality expression and the identity C(p)C(1 − p) = R.
  • F.5 Proof of Proposition E.5: The proof of Proposition E.5 establishes positivity of the relevant expression by rewriting it into curvature- and group-weighted terms.It treats separately the cases ρ ≥ 2 and 1 < ρ < 2 after using 1/(1+ρ) < α to obtain 1 < τ < ρ.
  • F.6 Proof of Proposition E.6: For Proposition E.6, R is rewritten as F(p(1 − p)), where F(v) = (ρ − 1)^2v / (ρ + (ρ − 1)^2v) is strictly increasing for v > 0.This converts the contraction comparison into a monotonicity statement in p(1 − p).
  • F.6 Proof of Proposition E.6: Under the stated assumptions, p lies strictly between α and β, implying p(1 − p) > αβ.The argument uses τ < β/α and the strict concavity and symmetry of v(x) = x(1 − x).
  • F.6 Proof of Proposition E.6: The final comparison yields R > F(αβ) > Γ^2, completing the proof of Proposition E.6.The positivity follows from α < 1/2, ρ > 1, and αρ > β.
Loading 2606.04662v1…