PaperlayerView exact revision on arXiv

Source-linked AI summary

Spectral Allocation: Why Muon Outperforms Adam, and How to Improve Muon

Xiaodong Wu, Wenyi Yu, Chao Zhang, Philip Woodland

arXiv:2608.25990v1cs.LG

TL;DR

The paper asks why Muon accelerates Transformer pretraining despite incomplete understanding of its mechanism. It probes momentum-buffer singular directions on held-out batches, finds a stable anisotropic step profile, and uses that profile to design SAMuon. Across evaluated modded-nanogpt scales and batch sizes, SAMuon variants outperform tuned AdamW and Muon, with SAMuon requiring 13.3%–24.0% fewer tokens than Muon to reach the same validation loss.

  • Problem

    The source of orthogonal optimisers’ fast convergence remains an open question, despite Muon’s acceleration relative to Adam.

  • Method

    The paper performs out-of-sample spectral probing of momentum-buffer singular directions and designs static head-anchored SAMuon priors from the measured loss-optimal step profile.

  • Results

    13.3%–24.0% fewer training tokens let SAMuon reach Muon’s validation loss, while both SAMuon variants outperform tuned AdamW and Muon across all evaluated configurations.

  • Takeaways & Limitations

    Muon’s uniform whitening underutilises the tolerant spectral bulk, motivating low-overhead spectral allocation through SAMuon and SAMuon-lite.

  • Takeaways & Limitations

    The reported SAMuon variants are under-tuned relative to Muon because they use Muon’s schedule and transfer 124M hyperparameters to larger models without re-tuning.

Abstract

from arXiv · show

Orthogonal optimisers such as Muon can substantially accelerate large language model pretraining relative to Adam, yet the mechanism remains incompletely understood. We investigate this through an out-of-sample spectral probing analysis of Transformer loss landscapes. At checkpoints along real training trajectories, we decompose each momentum buffer into its singular directions and estimate the loss-optimal step size along each direction on held-out data. The resulting spectral profile is anisotropic yet stable across batches and training stages, and consistent across the optimisers and model scales: a volatile head operating at the Edge-of-Stability supports a much smaller step size than the tolerant bulk, which permits substantially larger steps. This profile provides a unified spectral allocation account of why Muon outperforms Adam, which outperforms SGD. It also exposes a limitation of Muon's uniform scaling: it still underutilises the bulk. Guided by this finding, we introduce Spectral-Aware Muon (SAMuon), which holds the head at the Muon scale and amplifies the bulk using a static spectral prior. We provide two variants: the complete SAMuon follows the measured profile using a low-rank randomised SVD and the simplified SAMuon-lite uses a two-level approximation via rank-one power iteration. Neither method adds persistent optimiser state or notable extra FLOPs beyond Muon at scale, and the idealised exact-whitening versions of both retain Muon's asymptotic convergence rate under standard assumptions. Across "modded-nanogpt" models from 124M to 1B parameters, both variants outperform tuned AdamW and Muon (Scion implementation) baselines in all evaluated model-scale and batch-size configurations. SAMuon requires 13.3% to 24.0% fewer training tokens to reach the same validation loss as Muon, while SAMuon-lite retains most of this gain with near-zero wall-clock overhead.

1 Introduction

The paper uses out-of-sample spectral probing to explain Muon’s advantage and proposes SAMuon variants that amplify the tolerant spectral bulk. Across modded-nanogpt settings, both variants outperform tuned AdamW and Muon, with SAMuon achieving substantial token savings.

  • Motivation: Muon’s fast convergence remains theoretically incompletely explained, motivating an out-of-sample analysis of Transformer loss landscapes.Existing explanations include spectral-norm steepest descent and curvature-based accounts, but the paper identifies limitations in these views.
  • Spectral probing: The measured loss-optimal step profile is anisotropic and stable, with a volatile head permitting smaller steps than the tolerant bulk.The profile is mostly flat across the bulk and drops sharply near the head, with a close-to-linear decline over leading log ranks.
  • Spectral allocation: Muon outperforms Adam, which outperforms SGD, because whitening reallocates step mass from the volatile head toward the more tolerant bulk.SGD allocates steps in proportion to singular values, Adam dampens this misallocation, and Muon applies a uniform scale.
  • SAMuon: SAMuon holds the volatile head at Muon’s scale while amplifying the remaining spectrum using static spectral priors.Complete SAMuon uses rank-k randomised SVD; SAMuon-lite uses a two-level power-iteration approximation, with no persistent state and negligible extra FLOPs.

2 Related work

Related work interprets Muon through norm geometry, curvature, spectral analysis, and adaptive variants. The paper positions its spectral-allocation account as addressing curvature-aware optimisation while preserving Muon’s low-memory appeal.

  • Theoretical interpretations: Norm-geometric theories characterize orthogonalisation as spectral-norm steepest descent with robust update scales and learning-rate transfer.These theories also describe trust-region interpretations and spectral-norm max-margin bias, but do not address curvature inside the norm ball.
  • Curvature explanations: Curvature-based work links Muon’s advantage to isotropic curvature or lower directional-sharpness penalties, while this paper reports that isotropic curvature is violated in practice.The paper also contrasts its out-of-sample spectral probing with aggregate in-sample Hessian diagnostics.
  • Spectral analysis: Empirical Hessian studies establish highly anisotropic neural-network landscapes with bulk spectra near zero and a small number of large outliers.Transformer analyses further examine anisotropy at parameter-block granularity.
  • Adaptive methods: Muon’s single momentum buffer gives it a memory advantage over Adam and many adaptive or second-order methods that add model-scaling state.Several Muon variants improve adaptivity but sacrifice this memory advantage; COSMOS instead restricts adaptive preconditioning to a leading subspace.

3 Preliminaries

Muon updates each two-dimensional weight matrix by integrating gradients into a momentum buffer and orthogonalising that buffer. Its implementation uses fixed matrix-dependent scaling and Newton–Schulz approximation rather than an explicit SVD.

  • Muon update: Muon operates independently on each two-dimensional weight matrix, integrating the mini-batch gradient into a momentum buffer before producing the update.The buffer, rather than the raw gradient, is the object that gets orthogonalised.
  • Scaling: The Scion implementation uses a fixed scale κ = sqrt(dout/din), where dout and din are the matrix’s output and input dimensions.The scale is applied per weight matrix.
  • Spectral representation: The exposition assumes the momentum buffer has full row rank, with rank r = m, and represents matrices using reduced singular value decompositions.The singular values are ordered by spectral rank, with corresponding left and right singular vectors.
  • Orthogonalisation: Newton–Schulz approximates the polar factor UV^T using a small number of matrix products, avoiding an explicit SVD.Under exact whitening, every non-zero singular value is mapped to one.

4 Spectral probing for the loss-optimal step profile

Spectral probing measures held-out loss-optimal steps along momentum-buffer singular directions and finds a stable, strongly anisotropic profile: a volatile head limits Muon while a tolerant bulk permits larger steps.

  • Methodology: Spectral probing constructs one model-wide probe per singular rank from the updated momentum buffers of the probed weight matrices.The probes preserve the orientation and Scion scaling of Muon’s corresponding singular components, enabling direct comparison with Muon’s effective learning rate.
  • Methodology: The loss-optimal step size is estimated by a local-quadratic approximation to held-out loss rather than exhaustive line searches.The procedure averages measurements over held-out batches and tracks the resulting profile across checkpoints.
  • Observations: The measured profile spans nearly an order of magnitude, remaining mostly flat across the bulk before dropping sharply near the head.Its shape is highly anisotropic yet stable across batches, training stages, and model scales.
  • Observations: The leading ranks carry most spectral energy but have optimal steps nearly an order of magnitude smaller than the bulk, with the head pinned at Muon’s Edge-of-Stability scale.This makes the volatile head the limiting direction for Muon’s uniform update.
  • Observations: The tolerant bulk permits steps several times larger than the head, but Muon applies the same conservative scale to every rank.The resulting idealised cumulative loss-reduction headroom motivates amplifying the bulk, although cross-rank interactions make its recoverable size uncertain.
  • Why Muon outperforms Adam: Viewing optimisers as spectral allocations gives a unified account of the reported ordering Muon over Adam over SGD.Muon improves bulk utilisation through whitening, while the analysis identifies uniform scaling as approximately correct but still structurally suboptimal.

5 Spectral-Aware Muon (SAMuon)

SAMuon retains Muon’s momentum and whitening while applying a static, head-anchored spectral allocation that boosts the tolerant bulk. Its complete and lite variants use low-rank estimates, with the lite version targeting near-Muon runtime and both idealised variants preserving Muon’s convergence orders.

  • Design: SAMuon holds the volatile head at Muon’s scale and increases the remaining whitened spectrum toward the bulk’s measured tolerance.The allocation is static and head-anchored, avoiding online profile learning and persistent optimiser state.
  • Variants: SAMuon follows the measured log-rank-linear profile, while SAMuon-lite uses a two-level profile with one anchored leading direction and a uniformly boosted remainder.SAMuon estimates the leading k singular pairs with randomised low-rank SVD; SAMuon-lite uses rank-one power iteration.
  • Efficiency: Both variants retain one dense momentum buffer per weight matrix, matching Muon and using less persistent memory than Adam or SOAP.The head estimate adds only a transient low-rank factor pair rather than optimiser state.
  • Efficiency: 0.5% extra iteration time is measured for SAMuon-lite on the 1B model, compared with 7.4% for SAMuon’s current low-rank SVD implementation.The SVD overhead is left for future implementation improvements despite its low theoretical FLOP count.
  • Convergence: Under exact whitening, both idealised variants preserve Muon’s convergence orders: O(T^-1/4) stochastically and O(T^-1/2) deterministically.The guarantee applies to head-anchored allocations with scales in [1, γ], up to γ-dependent constants.

6 Experiments

Experiments evaluate SAMuon and SAMuon-lite against tuned AdamW and Muon across model scales and batch sizes. Both variants improve validation loss and token efficiency, while SAMuon-lite preserves most gains at lower cost.

  • Both SAMuon variants attain lower final validation loss than tuned Muon across every trained model-scale and batch-size cell.
  • 13.3%–24.0% token-efficiency improvement over Muon is achieved by SAMuon across the evaluation grid.
  • 13.3%–22.1% token-efficiency improvement over Muon is achieved by SAMuon-lite across model scales and batch sizes.
  • SAMuon improvements generally grow with batch size, including 20.3% to 24.0% on 124M and 13.3% to 15.6% on 1B.
  • SAMuon matches or outperforms SAMuon-lite in every reported cell, with a token-efficiency advantage of 0.0–2.3 percentage points.

7 Discussion

The discussion presents offline spectral probing as a practical route for optimizer design and interprets the SAMuon results as evidence for measured spectral allocation. It also identifies tuning and profile-design boundaries for future work.

  • Offline spectral probing of saved checkpoints can provide design signals while reducing the need for exhaustive full-training sweeps during early optimizer development.
  • SAMuon variants are under-tuned because they use a Muon-optimized schedule and direct 124M-to-larger-scale hyperparameter transfer.
  • SAMuon-lite achieves 13.3%–22.1% token-efficiency improvement, while SAMuon adds up to 2.3 percentage points through its richer log-rank-linear profile.
  • Future profiles could model cross-probe curvature interactions, trajectory drift, and block-dependent spectral differences.

8 Conclusion

The conclusion explains Muon’s behavior through a stable, anisotropic spectral profile and proposes SAMuon to allocate larger steps to the tolerant bulk. Validation supports substantial token savings, while probing assumptions and scope conditions are explicitly checked.

  • The measured profile is mostly flat across the tolerant bulk but declines sharply near a volatile head at the Edge-of-Stability.
  • Uniform whitening under-exploits the bulk because the volatile head constrains the shared Muon scale.
  • SAMuon uses a static head-anchored prior, while SAMuon-lite approximates the profile with two spectral levels and nearly Muon-speed execution.
  • 13.3%–24.0% token-efficiency improvement lets SAMuon reach Muon’s final loss across all trained model and batch-size combinations.
  • The probing analysis assumes local quadraticity, held-out descent directions, and positive directional curvature, filtering 0.52% of sampled probes.
  • The local-quadratic estimate is supported by close agreement with exact line search across 24 analyses and 773 probes, with median relative deviation near 5%.

B Spectral Probing of an AdamW Trajectory

Out-of-sample probing on AdamW reveals the same volatile-head and tolerant-bulk structure found along Muon trajectories, indicating that the profile is not specific to Muon’s trajectory.

  • The AdamW profile rises by more than an order of magnitude from the volatile head to a mostly flat, tolerant bulk.
  • The deep tail contributes only a modest residual share of cumulative out-of-sample loss reduction.
  • The local-quadratic estimate agrees with golden-section line search, with median deviation near 5% in the operationally relevant range.
  • Muon’s head is pinned at the shared global scale because whitening applies the same scale to every singular direction, whereas AdamW’s unnormalised head step varies across checkpoints.

C Spectral Probing at Width 768

The spectral structure persists at width 768: a volatile head near the Muon scale transitions into a mostly flat tolerant bulk across checkpoints.

  • The width-768 Muon analysis reproduces the qualitative structure observed in the width-512 analysis.
  • At width 768, the profile remains highly anisotropic and stable across checkpoints, with a volatile head, tolerant bulk, and approximately log-rank-linear transition.
  • The approximate plateau-onset rank shifts from about 32 at width 512 to about 40 at width 768.
  • With only two directly probed widths, the observed agreement is not treated as evidence for a general scaling law.

D In-Sample Spectral Probing

In-sample probing makes the bulk and tail appear more productive than they are out-of-sample, reflecting batch refitting rather than transferable loss reduction and motivating conservative spectral allocation comparisons.

  • D In-Sample Spectral Probing: In-sample probing measures each probe on the same batch that formed the momentum buffer, unlike the held-out-data protocol.
  • D In-Sample Spectral Probing: The in-sample bulk is elevated above its out-of-sample level, and its anisotropy spans two to three orders of magnitude.
  • D In-Sample Spectral Probing: In-sample cumulative loss reduction accrues steadily from the tail, whereas out-of-sample reduction concentrates toward the head.
  • E Construction of the Spectral Allocation Comparison: The allocation comparison treats SGD as singular-value proportional, Adam as partially damped, Muon as uniform, SAMuon-lite as uniformly bulk-amplified, and SAMuon as log-rank-linear.
  • E Construction of the Spectral Allocation Comparison: SAMuon-lite uses k = 1 and γ = 5, while SAMuon uses γ = 7.07 to follow the measured near-head decline.
  • F.1 Setup and notation: The idealised exact-whitening convergence guarantee covers head-anchored scaling with amplification bounded by γ and specializes to Muon when γ = 1.

F.3 Main convergence theorems

The idealised exact-whitening SAMuon update admits an ergodic convergence bound under assumptions F.1–F.3 and recovers Muon’s stochastic and deterministic convergence orders. The analysis also makes explicit that bulk amplification increases noise and deterministic penalties, while warmup controls amplified tracking error.

  • Under assumptions F.1–F.3, idealised exact-whitening SAMuon with constant learning rate η yields a bound on the expected nuclear norm of the true gradient.
  • The convergence-rate choices in Corollary F.5 recover Muon’s stochastic and deterministic convergence orders.The result uses the same zero-buffer initialisation as Algorithm 1.
  • Setting γ = 1 specialises the theorem to Muon with zero-buffer initialisation.The resulting initial-gradient term is vanishing, and both conventions give the convergence orders stated in Corollary F.5.
  • Amplifying the tail by γ linearly increases the stochastic noise floor and quadratically increases the deterministic penalty.The descent bound exposes this dependence through the amplified tracking error.
  • The head anchor is empirically motivated by Edge-of-Stability measurements, while the convergence bound is shape-agnostic for scales in [1, γ].Warmup withholds amplification while momentum tracking error is largest early in training.

G.3 Training hyperparameters

Training uses tuned optimizer settings transferred across model scales, with SAMuon variants fixed at spectral radius 50 and warmed into their spectral shaping. The custom Newton–Schulz implementation is accurate over the operative singular-value range, and SAMuon achieves the best-of-sweep loss at least as low as SAMuon-lite.

  • All experiments use FineWeb with global batch sizes of 1M, 2M, and 4M tokens, while training budgets range from approximately 10B to 20B tokens.The 124M budget is approximately four times the Chinchilla-optimal allocation.
  • Hyperparameters are tuned per batch size on the 124M model and transferred across scales, except AdamW’s learning rate, which is halved at each scale step.SAMuon uses fixed spectral radius 50 and sweeps only the tail boost γ.
  • The spectral profile is warmed over the first 30% of training, interpolating from Muon scaling to the target amplification before holding it fixed.This schedule allows the model to settle onto a stable Muon trajectory before full anisotropic shaping.
  • For normalised input singular values σ ≥ 0.02, the four-iteration 4 × 3 Newton–Schulz scheme has maximum response error 0.0069.
  • SAMuon attains a best-of-sweep validation loss at least as low as SAMuon-lite at all three batch sizes, with both effectively tied at batch size 1024.SAMuon’s optimum is γ = 7.07 in all cases, whereas SAMuon-lite’s optimum varies.

I Per-Module Sign Alignment

Per-module sign alignment measures how consistently each momentum-buffer singular direction agrees with held-out gradients. The head is reliably aligned across weight matrices, whereas tail alignment depends strongly on module type and is weakest for attention query and key projections.

  • Sign alignment is estimated by applying each rank-d direction in the optimizer’s fixed orientation and averaging agreement across weight matrices.This operationally measures whether a direction makes transferable progress on unseen batches.
  • Cross-matrix agreement determines whether contributions add or cancel within a rank and contributes to the mostly flat bulk.
  • The top few ranks are almost always aligned across all weight matrices, making the head orientation effectively deterministic.
  • In the tail, attention projections drift toward or below 0.5 alignment, while MLP projections remain strongly aligned across nearly the entire spectrum.
  • Query and key projections have the worst tail alignment, so their low-energy tail directions are the least trustworthy.
Loading 2608.25990v1…