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

> Spectral probing reveals Transformer loss landscapes have a volatile head and tolerant bulk, motivating SAMuon which amplifies bulk step sizes to beat Muon by 13-24% in token efficiency.

- **Source:** [arXiv](https://arxiv.org/abs/2608.25990)
- **Published:** 2026-09-05
- **Permalink:** https://picx.dev/p/IhaG0E
- **Whiteboard:** https://picx.dev/p/IhaG0E/image

## Summary

# 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**

| Model | Batch | Tokens | AdamW | Muon | SAMuon loss | SAMuon improv. | SAMuon-lite loss | SAMuon-lite improv. |
|-------|-------|--------|-------|------|-------------|----------------|------------------|---------------------|
| 124M  | 1024  | 10B    | 3.2105 | 3.1622 | **3.1435** | **20.3%** | 3.1438 | 20.3% |
| 124M  | 2048  | 10B    | 3.2410 | 3.1726 | **3.1484** | **23.8%** | 3.1511 | 21.5% |
| 124M  | 4096  | 10B    | 3.3050 | 3.1953 | **3.1658** | **24.0%** | 3.1678 | 22.1% |
| 300M  | 1024  | 10B    | 3.0297 | 2.9836 | **2.9640** | **17.7%** | 2.9663 | 15.9% |
| 300M  | 2048  | 10B    | 3.0658 | 2.9941 | **2.9713** | **18.8%** | 2.9738 | 17.1% |
| 1B    | 1024  | 20B    | 2.7599 | 2.7251 | **2.7114** | **13.3%** | 2.7117 | 13.3% |
| 1B    | 4096  | 20B    | 2.8067 | 2.7413 | **2.7222** | **15.6%** | 2.7240 | 14.4% |

### 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

---

_Markdown view of https://picx.dev/p/IhaG0E, served by PicX — AI-generated visual whiteboard summaries of research papers._