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

Summary (Overview)

• Core contribution: The paper introduces an out-of-sample spectral probing framework that measures the loss-optimal step size along each singular direction of the momentum buffer in Transformer pretraining, revealing a stable, highly anisotropic spectral profile.
• Key finding: The loss landscape exhibits a "volatile head" (leading singular direction) operating at the Edge-of-Stability with a much smaller optimal step size than the "tolerant bulk," which permits substantially larger steps. This provides a unified spectral allocation account of why Muon outperforms Adam, which outperforms SGD.
• Proposed method: Spectral-Aware Muon (SAMuon) and its simplified variant SAMuon-lite, which hold the head at Muon's scale while amplifying the bulk using static spectral priors derived from the measured profile.
• Empirical results: Across "modded-nanogpt" models from 124M to 1B parameters, SAMuon achieves 13.3%–24.0% token-efficiency improvement over Muon (Scion implementation), with SAMuon-lite retaining most of this gain at nearly zero wall-clock overhead.
• Theoretical guarantee: Both variants retain Muon's asymptotic convergence rate ($\mathcal{O}(T^{-1/4})$ stochastic, $\mathcal{O}(T^{-1/2})$ deterministic) under standard assumptions.

Introduction and Theoretical Foundation

Background and Motivation

The pretraining of Transformer-based decoder-only LLMs is constrained by optimiser efficiency. Orthogonal optimisers such as Muon achieve substantial convergence acceleration over Adam, particularly during the transient (early, non-asymptotic) phase of pretraining. However, the mechanism behind this acceleration remains incompletely understood.

Theoretical Context

The dominant view frames Muon as spectral-norm steepest descent, which accounts for stability and learning-rate transfer but does not explain the convergence advantage on highly curved objectives, since whitening "effectively removes all curvature information." Recent work by Su [2025] showed orthogonalisation is loss-optimal under an isotropic curvature model, but this assumption is violated in practice. Concurrent work by Wang et al. [2026] attributes Muon's advantage to a lower directional-sharpness penalty but relies on aggregate in-sample Hessian diagnostics that compress the update's spectral geometry into a single scalar.

Key Mathematical Foundation

The Muon optimiser operates on each 2D weight matrix $\mathbf{W}^{(t)} \in \mathbb{R}^{m \times n}$ (with $m \leq n$):

$$\mathbf{M}^{(t)} = \mu \mathbf{M}^{(t-1)} + (1-\mu) \mathbf{G}^{(t)} \tag{1}$$ $$\mathbf{O}^{(t)} = \mathrm{NS}\left(\mathbf{M}^{(t)}\right) \approx \left(\mathbf{M}^{(t)} \mathbf{M}^{(t)\top}\right)^{-1/2} \mathbf{M}^{(t)} = \mathbf{U}\mathbf{V}^{\top} \tag{2}$$ $$\mathbf{W}^{(t+1)} = \mathbf{W}^{(t)} - \eta \kappa \mathbf{O}^{(t)} \tag{3}$$

where $\mu$ is the momentum coefficient, $\eta$ the learning rate, and $\kappa = \sqrt{d_{\mathrm{out}}/d_{\mathrm{in}}}$ the Scion scaling factor. The Newton-Schulz iteration $\mathrm{NS}(\cdot)$ approximates the polar factor $\mathbf{U}\mathbf{V}^{\top}$ without explicit SVD.

Methodology

Spectral Probing Framework

The paper proposes an out-of-sample spectral probing procedure that measures the loss landscape along each direction in the spectral basis of the momentum buffer:

1. Per-Rank Probe Construction: For each of the $N_W = 72$ 2D weight matrices, the gradient of a batch $B_g$ is integrated into the momentum buffer, and the response is measured on a disjoint out-of-sample batch $B_o \neq B_g$.

2. Probe Definition: The rank-$d$ probe collects the $d$-th singular component of every probed matrix:

$$\mathbf{P}_{j,d} = -\kappa_j \mathbf{u}_{j,d} \mathbf{v}_{j,d}^{\top}, \quad \Delta\boldsymbol{\theta}^{(d)} = \operatorname{vec}\left(\{\mathbf{P}_{j,d}\}_{j=1}^{N_W}\right) \tag{4}$$

3. Loss-Optimal Step Estimation: Using a local-quadratic approximation (LQA):

$$\mathcal{L}(\boldsymbol{\theta}_t + \eta \Delta\boldsymbol{\theta}; \mathcal{B}_o) \approx \mathcal{L}(\boldsymbol{\theta}_t; \mathcal{B}_o) + \eta \mathbf{g}^{\top} \Delta\boldsymbol{\theta} + \frac{1}{2}\eta^2 \Delta\boldsymbol{\theta}^{\top} \mathbf{H} \Delta\boldsymbol{\theta} \tag{5}$$

The optimal step size is:

$$\eta_d^* = -\frac{\mathbf{g}^{\top} \Delta\boldsymbol{\theta}^{(d)}}{\Delta\boldsymbol{\theta}^{(d)\top} \mathbf{H} \Delta\boldsymbol{\theta}^{(d)}} \tag{6}$$

The directional curvature is estimated via finite differences:

$$\Delta\boldsymbol{\theta}^{\top} \mathbf{H} \Delta\boldsymbol{\theta} \approx \frac{2}{\varepsilon^2}\left[\mathcal{L}(\boldsymbol{\theta}_t + \varepsilon \Delta\boldsymbol{\theta}; \mathcal{B}_o) - \mathcal{L}(\boldsymbol{\theta}_t; \mathcal{B}_o) - \varepsilon \mathbf{g}^{\top} \Delta\boldsymbol{\theta}\right] \tag{7}$$

SAMuon Algorithm

The SAMuon update applies a head-anchored spectral allocation:

$$\mathbf{O}_{\mathrm{SA}}^{(t)} = \gamma_t \mathrm{NS}\left(\mathbf{M}^{(t)}\right) - \sum_{i=1}^{k} (\gamma_t - s_i(t)) \mathbf{u}_i \mathbf{v}_i^{\top} \tag{8}$$

SAMuon follows the measured log-rank-linear profile:

$$s_i = 1 + (\gamma - 1) \frac{\log i}{\log k}, \quad 1 \leq i \leq k \ (k \geq 2) \tag{9}$$

with $k(d_{\mathrm{model}}) = \left\lfloor 32\sqrt{d_{\mathrm{model}}/512} \right\rfloor$ (width rule).

SAMuon-lite uses $k=1$ (two-level approximation via power iteration), amplifying the bulk uniformly to scale $\gamma$ while pinning only the head.

Both variants use a cosine-scheduled warmup: $\gamma_t = 1 + (\gamma - 1)w_t$, with $w_t \in [0,1]$ rising from 0 to 1 over the warmup horizon. Setting $\gamma = 1$ recovers standard Muon exactly.

Empirical Validation / Results

Experimental Setup

• Models: "modded-nanogpt" at 124M (width 768), 300M (width 1280), and 1B (width 2560) parameters
• Data: FineWeb dataset, 10B tokens (124M/300M) or 20B tokens (1B)
• Batch sizes: 1024, 2048, 4096 sequences
• Baselines: Tuned AdamW and Muon (Scion implementation)
• Hardware: NVIDIA A100 (80GB) GPUs, bfloat16 mixed precision

Main Results

Table 1: Final validation loss and token-efficiency improvement across the model-scale × batch-size grid

Key Findings

1. Both SAMuon variants outperform Muon in all configurations: SAMuon improves final loss by 0.0137–0.0295 over Muon, translating to 13.3%–24.0% token-efficiency improvement.

2. Improvement grows with batch size: Token-efficiency gains increase from 20.3% (batch 1024) to 24.0% (batch 4096) on 124M, likely because larger batches provide less noisy spectral estimates.

3. SAMuon matches or outperforms SAMuon-lite: The gap generally widens with batch size (0.0–2.3 percentage points), suggesting the finer-grained log-rank-linear profile carries meaningful signal.

4. Computational overhead: SAMuon-lite adds only 0.5% to Muon's iteration time (power iteration), while SAMuon adds 7.4% (randomised SVD). Neither adds persistent optimiser state.

Theoretical and Practical Implications

Theoretical Implications

• Unified spectral allocation account: The paper provides the first unified explanation of why Muon > Adam > SGD in Transformer pretraining, viewing optimisers as spectral allocators:

  • SGD allocates step size proportional to singular values (opposite to what the landscape demands)
  • Adam dampens but cannot overcome this misallocation
  • Muon whitens to uniform scale, reallocating step mass from head to bulk

• Convergence guarantee: For idealised exact-whitening versions, both variants preserve Muon's convergence rates: $\mathcal{O}(T^{-1/4})$ stochastic and $\mathcal{O}(T^{-1/2})$ deterministic, up to $\gamma$-dependent constants.

Practical Implications

• Memory efficiency: Both variants maintain exactly one dense momentum buffer per weight matrix ($mn$), identical to Muon and strictly smaller than Adam ($2mn$) and SOAP ($6mn$).

• FLOPs overhead: Negligible—$O(mn)$ for SAMuon-lite's power iteration and $O(mnk)$ for SAMuon's rank-$k$ randomised SVD, both small beside the $O(mnr)$ Newton-Schulz whitening.

• Offline probing methodology: The spectral probing framework turns optimiser design into a measurable problem, enabling researchers to probe loss landscapes with saved checkpoints efficiently.

Conclusion

The paper demonstrates that the Transformer pretraining loss landscape has a stable, highly anisotropic spectral profile: mostly flat across the bulk, with a sharp decline near a single volatile head at the Edge-of-Stability. This profile provides a unified account of optimiser performance and reveals that Muon's uniform whitening underutilises the tolerant bulk.

SAMuon and SAMuon-lite exploit this headroom by holding the head at Muon's scale while amplifying the bulk, achieving 13.3%–24.0% token-efficiency improvements over Muon across all evaluated configurations. The gains validate the spectral allocation analysis and suggest that aligning optimiser updates with the measured loss-optimal step profile translates directly into better convergence.

Future directions include:

• Investigating cross-probe curvature interactions for refined profiles
• Adapting the profile to drift over training trajectories
• Exploring block-dependent allocations (preliminary evidence shows inter-block differences)
• Improving the svd_lowrank implementation to reduce SAMuon's wall-clock cost
• Properly tuning learning rate schedules specifically for SAMuon variants