Pretraining Data Scale Reverses the Muon-AdamW Fine-Tuning Gap

Paper: 2611.03633
Authors: Catherine Li, Jingchu Gai, Andrej Risteski, Jacob Mitchell Springer, Aditi Raghunathan (Carnegie Mellon University)
Categories: cs.LG, cs.AI, cs.CL

The Gap

The optimizer landscape for Large Language Model (LLM) pretraining has recently been disrupted by Muon (Momentumized Orthogonalized Newton-Schulz updates). By constraining parameter updates to the Stiefel manifold of orthogonal matrices, Muon consistently achieves faster convergence and lower base pretraining loss compared to AdamW at identical compute budgets.

However, modern foundation models are almost never deployed as raw base checkpoints; they must undergo post-training (Supervised Fine-Tuning, RLHF, and task adaptation). This introduces a critical open question: Does Muon’s pretraining loss advantage translate into superior downstream task performance, or does it degrade the learning-forgetting trade-off?

Prior empirical investigations suffered from significant blind spots:

1. Conflating Base Loss with Downstream Retention: Machine learning practitioners have long treated pretraining validation loss as an infallible proxy for downstream capability.
2. Ignoring Pretraining Scale Interactions: Comparisons between Muon and AdamW were predominantly conducted on small token budgets (sub-10B tokens), overlooking how scaling laws alter optimization dynamics.
3. Absence of a Theoretical Mechanism: Why would an optimizer that produces lower cross-entropy loss perform worse when fine-tuned on new tasks?

                 THE MUON VS ADAMW PRETRAINING PARADOX

  Pretraining Phase (Minimizing Base Cross-Entropy Loss):
  Muon  ───> Consistently lower Loss than AdamW across ALL data scales! ✓
  AdamW ───> Slower convergence, higher training loss.

  Post-Training Phase (Fine-Tuning on Downstream Tasks):
  Small Data Budget (< 5B tokens):
  - Muon dominates AdamW downstream (Better task adaptation & retention).

  Large Data Budget (> 10B-30B tokens):
  - THE REVERSAL: AdamW systematically beats Muon after fine-tuning!
  - As model scale increases, AdamW overtakes Muon at progressively
    smaller token-to-parameter ratios.

---

The Increment

One sentence: The authors discover that scaling pretraining data reverses the downstream Muon–AdamW fine-tuning gap, and explain this phenomenon through the Margin-to-Spread Ratio (MSR)—proving that while Muon creates sharper logit separation at low budgets, AdamW builds broader, more noise-resilient margin distributions at scale.

                    THE MARGIN-TO-SPREAD RATIO (MSR)

                  Correct Logit (y*)
                         │
                         ▼
  ───────────────────────┬────────────────────────────────────────
                         │ <────── Margin h ──────>
                         │
        Wrong Logits Distribution: Spread κ = std(wrong logits)
             ┌───────────┴───────────┐
             │   ░░░░░░░░░░░░░░░░    │
  ───────────┴───────────────────────┴────────────────────────────

  MSR = h / κ = (Distance to incorrect mean) / (Spread of incorrect logits)

  - High MSR: The correct prediction is isolated far beyond the noise floor.
              Robust to output distribution perturbations during SFT.
  - Small Data Scale:  MSR(Muon) > MSR(AdamW)  ==> Muon wins fine-tuning.
  - Large Data Scale:  MSR(AdamW) > MSR(Muon)  ==> AdamW wins fine-tuning!

---

The Anatomy

1. The Learning–Forgetting Frontier

To rigorously evaluate fine-tuning behavior, the authors train multiple OLMo architectures (from 30M to 300M parameters) across pretraining token budgets ranging from 2.4B to 38.4B tokens on the DCLM corpus, followed by extensive hyperparameter grid sweeps during fine-tuning on StackMathQA.

The evaluation measures the Pareto frontier between:

• Target Task Learning Loss: How effectively the model learns new downstream capabilities.
• Pretraining Forgetting Loss: How severely the model degrades on its original pretraining distribution.

Across every model size, as pretraining tokens scale:

1. At small token budgets (e.g., 2.4B tokens for OLMo-60M), the Muon frontier strictly dominates AdamW.
2. At larger token budgets (e.g., 9.6B, 19.2B, and 38.4B tokens), the frontier reverses: AdamW retains more pretraining knowledge for any given level of fine-tuning accuracy.
3. Crucially, as model parameter size grows (from 30M to 300M), the crossover point shifts to smaller token-to-parameter ratios, indicating that large frontier models are especially prone to this reversal.

2. Theoretical Mechanism: The Margin-to-Spread Ratio (MSR)

Fine-tuning can mathematically be framed as an adversarial perturbation $\delta$ applied to the pre-softmax output logits $s(x) = W^\top x$. A base model is resilient to forgetting if its correct predictions can tolerate large output shifts before the argmax ranking flips.

To quantify this tolerance, the authors define the Margin-to-Spread Ratio (MSR): $\rho = \frac{h}{\kappa}$ where:

• $h = s_{y^*} - \frac{1}{|\mathcal{V}| - 1} \sum_{j \ne y^*} s_j$ is the mean margin (distance between the correct logit $y^*$ and the centroid of incorrect logits).
• $\kappa = \sqrt{\frac{1}{|\mathcal{V}| - 1} \sum_{j \ne y^*} (s_j - \bar{s}_{\text{wrong}})^2}$ is the spread (standard deviation among incorrect logits).

A prediction with a large $h$ but an equally large $\kappa$ is fragile: an arbitrary small perturbation can easily push an outlying incorrect logit above $y^*$. True robustness requires maximizing margin $h$ while minimizing spread $\kappa$.

3. Formal Proof in Linear Associative Memory

To analyze why optimizers diverge, the authors study a linear associative memory retrieving $N$ key-value facts $(u_i, v_i) \in \mathbb{R}^d \times \mathbb{R}^d$ parameterized by $W \in \mathbb{R}^{d \times d}$, with data budget ratio $\gamma = N/d$.

Theorem 4.2 (Informal Summary):
Under isotropic Gaussian facts, after an update from zero initialization, the ratio of Gradient Descent’s (AdamW proxy) MSR to Muon’s MSR satisfies: $\mathcal{R}(\gamma) = \frac{\rho_{\text{GD}}(\gamma)}{\rho_{\text{Muon}}(\gamma)} = \frac{1 + \frac{2}{\gamma}}{\sqrt{1 + \frac{1}{\gamma}}}$ which is a strictly increasing function of the data budget $\gamma$.

• When data budget $\gamma$ is small, Muon’s matrix-orthogonalization forces faster spectral alignment, yielding a higher MSR ($\mathcal{R} < 1$).
• As data budget $\gamma$ expands, Gradient Descent / AdamW accumulates coherent gradient signals along the signal subspace while isotropically diffusing noise across orthogonal directions, causing AdamW’s MSR to surpass Muon ($\mathcal{R} > 1$).

---

Quantitative Evaluation

1. MSR Crossover in Language Models

Evaluating pretrained OLMo models on held-out DCLM text reveals that the empirical MSR tracks the downstream fine-tuning gap with high fidelity:

Model Scale	Pretrain Tokens	Muon MSR ($\rho$)	AdamW MSR ($\rho$)	Dominant Optimizer (MSR)	Downstream Fine-Tuning Winner
OLMo-60M	2.4B	2.84	2.51	Muon (+13.1%)	Muon
	4.8B	3.22	3.15	Muon (+2.2%)	Muon
	9.6B	3.65	3.82	AdamW (+4.7%)	AdamW
	19.2B	4.08	4.45	AdamW (+9.1%)	AdamW
	38.4B	4.31	4.92	AdamW (+14.2%)	AdamW
OLMo-100M	4.8B	3.45	3.38	Muon (+2.1%)	Muon
	9.6B	3.88	4.09	AdamW (+5.4%)	AdamW
	19.2B	4.25	4.76	AdamW (+12.0%)	AdamW
OLMo-150M	4.8B	3.61	3.74	AdamW (+3.6%)	AdamW
	19.2B	4.40	5.08	AdamW (+15.5%)	AdamW

Key Findings:

1. The Reversal is Robust Across Scales: In every tested model, Muon leads in MSR at low data volumes, but AdamW overtakes it once token scale increases.
2. Larger Models Flip Sooner: For OLMo-60M, the crossover occurs around ~7B tokens (token
   ratio ~115
   ). For OLMo-150M, AdamW already leads at 4.8B tokens (token
   ratio ~32
   ).
3. Overfitting to Cross-Entropy: Muon’s orthogonalization forces singular values to be uniform, which aggressively squashes cross-entropy loss by sharpening logits, but leaves the underlying representation hypersensitive to the directional weight updates made during fine-tuning.

---

Key Takeaways

1. Pretraining Loss is Not King: An optimizer that achieves a lower validation loss during pretraining does not necessarily yield a better foundation model for downstream deployment.
2. Scale Changes Optimization Geometry: Muon’s superiority in small-scale benchmarks (like Speedrun or sub-10B training) cannot be extrapolated blindly to multi-trillion-token frontier training runs.
3. The Hybrid Compromise: Because Muon converges rapidly in early training while AdamW builds more robust asymptotic representations, hybrid schedules—e.g., pretraining with Muon early on and transitioning to AdamW or annealing with AdamW-like momentum—represent a highly promising frontier.

                 MUON 与 ADAMW 的预训练反转悖论

  预训练阶段（优化目标：最小化基础 Cross-Entropy Loss）：
  Muon  ───> 在所有数据尺度下，Loss 均系统性低于 AdamW！✓
  AdamW ───> 收敛速度较慢，预训练 Loss 相对更高。

  后训练微调阶段（下游任务 SFT 适配）：
  小数据预算规模（< 5B Tokens）：
  - Muon 在下游微调任务中占据优势（收敛快、遗忘少）。

  大数据预算规模（> 10B-30B Tokens）：
  - 发生惊人反转：AdamW 微调后的综合性能全面反超 Muon！
  - 随着模型参数量增大，AdamW 实现反超所需的“Token-参数比”越来越小！

                    边际展宽比（Margin-to-Spread Ratio, MSR）

                   正确预测 Logit (y*)
                          │
                          ▼
  ────────────────────────┼───────────────────────────────────────
                          │ <────── 边际间隔 h ──────>
                          │
          错误 Logit 经验分布: 展宽离散度 κ = std(wrong logits)
              ┌───────────┴───────────┐
              │   ░░░░░░░░░░░░░░░░    │
  ────────────┴───────────────────────┴───────────────────────────

  MSR = h / κ = (正确 Logit 距离错误均值的距离) / (错误 Logit 自身的标准差)

  - MSR 越大：正确预测在几何上被隔绝在噪声基线之上，抗微调扰动能力越强。
  - 小数据规模时：  MSR(Muon) > MSR(AdamW)  ==> Muon 微调表现领先。
  - 大数据规模时：  MSR(AdamW) > MSR(Muon)  ==> AdamW 微调表现反超！