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

> Muon's pretraining loss advantage over AdamW reverses after fine-tuning at larger data budgets, with the crossover shrinking as model size grows.

- **Source:** [arXiv](https://arxiv.org/abs/2611.03633)
- **Published:** 2026-11-07
- **Permalink:** https://picx.dev/p/ubDabG
- **Whiteboard:** https://picx.dev/p/ubDabG/image

## Summary

## Summary (Overview)

- **Key Finding**: The paper reveals that Muon's pretraining advantage over AdamW does not persist after fine-tuning at larger pretraining data budgets. Beyond a crossover token-parameter ratio (TPR), AdamW-pretrained models achieve better learning-forgetting tradeoffs after fine-tuning, despite Muon achieving lower pretraining loss.
- **Scaling Trend**: As model size increases (30M→600M parameters), the crossover where AdamW overtakes Muon occurs at progressively smaller pretraining data budgets (from 320 TPR down to 10 TPR).
- **Theoretical Contribution**: The authors introduce the **margin-to-spread ratio (MSR)** to quantify robustness to output perturbations, proving in a linear associative memory model that gradient descent's MSR gains on Muon's as the data budget grows.
- **Empirical Validation**: The MSR trend in real language models (30M–100M parameters) mirrors the fine-tuning reversal—Muon has larger MSR at small budgets, AdamW at large budgets.
- **Practical Implication**: Pretraining loss is a poor proxy for downstream adaptability, suggesting that optimizer selection and pretraining objectives should explicitly consider post-training properties.

---

## Introduction and Theoretical Foundation

### Background and Motivation

The Muon optimizer has recently demonstrated strong empirical performance for pretraining large language models, challenging AdamW's status as the standard optimizer (Jordan et al., 2024). However, prior work shows that pretraining optimization choices minimizing pretraining loss do not necessarily yield the best model after fine-tuning (Watts et al., 2026). The central question:

> Does Muon's pretraining advantage over AdamW extend to the learning–forgetting tradeoff?

### Learning-Forgetting Tradeoff

Fine-tuning improves downstream performance but may degrade pretrained knowledge. For a pretrained model $\theta_{PT}$ fine-tuned to $\theta_{FT}$, the paper measures:
- **Learning**: fine-tuning loss $\mathcal{L}_{FT}(\theta_{FT})$
- **Forgetting**: pretraining loss $\mathcal{L}_{PT}(\theta_{FT})$

The set of achievable outcomes is formalized as:

$$
\mathcal{T}(\theta_{\mathrm{PT}}) := \left\{(\mathcal{L}_{\mathrm{FT}}(\theta_{\mathrm{FT}}), \mathcal{L}_{\mathrm{PT}}(\theta_{\mathrm{FT}})) \mid \theta_{\mathrm{FT}} \in \Theta_{\mathrm{FT}}(\theta_{\mathrm{PT}})\right\}
$$

Base models are compared via their **learning-forgetting frontiers** (lower-left Pareto frontiers).

### Fine-Tuning Gap

The **fine-tuning gap** $G_{FT}$ is the signed mean minimum Euclidean distance between AdamW and Muon frontiers:
- $G_{FT} > 0$: AdamW dominates
- $G_{FT} < 0$: Muon dominates

### Crossover Token-Parameter Ratio (TPR)

The **crossover TPR** is the minimum TPR (pretraining tokens per model parameter) required for the AdamW frontier to overtake Muon's frontier.

---

## Methodology

### Experimental Setup

**Pretraining**: OLMo-style models (30M–600M parameters) trained on 150M to 154B DCLM tokens using warmup–stable–decay (WSD) schedules. Nearly 600 pretraining runs total.

**Fine-tuning**: Four datasets evaluated—GSM8K, StackMathQA, Alpaca, MusicPile. Sweeps across:
- AdamW learning rates: $10^{-4}$ to $10^{-2}$
- Muon learning rates: $6\times10^{-4}$ to $2\times10^{-2}$

Total: ~29,000 fine-tuning runs.

**Two evaluation settings**:
1. **Lowest-pretraining-loss setting**: Tune pretraining hyperparameters to minimize pretraining loss, then sweep fine-tuning configurations.
2. **Best-frontier setting**: Sweep all pretraining and fine-tuning configurations for the best possible learning-forgetting frontier.

### Theoretical Framework: Linear Associative Memory

The paper analyzes a linear associative memory where a weight matrix $W \in \mathbb{R}^{d\times d}$ associates key-value pairs $(u_i, v_i) \in \mathbb{R}^d \times \mathbb{R}^d$ drawn from $\mathcal{N}(0, \hat{I}_d)$. Logits are:

$$s_{ij}(W) := v_j^\top W u_i$$

Training uses softmax cross-entropy loss:

$$\hat{p}_W(\cdot \mid i) := \operatorname{softmax}(s_{i1}(W), \dots, s_{iN}(W)), \quad \mathcal{L}(W) := -\sum_{i=1}^{N} q_i \log \hat{p}_W(i \mid i)$$

**Data budget**: $\gamma = N/d > 1$ (facts per dimension), mirroring token-parameter ratio. As $d \to \infty$, $\log N = o(d)$.

**Optimizers compared**:
- **Muon** (exact orthogonalization): $W_{t+1} = W_t - \eta \operatorname{polar}(\nabla \mathcal{L}(W_t))$ where $\operatorname{polar}(A) = PQ^\top$ from SVD $A = P\Sigma Q^\top$
- **Gradient Descent**: $W_{t+1} = W_t - \eta \nabla \mathcal{L}(W_t)$

### The Margin-to-Spread Ratio (MSR)

**Definition 4.1**: For fact $i$, define pairwise margins $m_{ij} := s_{ii}(W) - s_{ij}(W)$ for $j \neq i$, with:

$$h_i := \underset{j \neq i}{\operatorname{mean}} m_{ij}, \qquad \kappa_i^2 := \underset{j \neq i}{\operatorname{mean}} (m_{ij} - h_i)^2$$

The per-fact MSR is $\rho_i := h_i / \kappa_i$. Model-level MSR averages means and variances:

$$\bar{h} := \underset{i}{\operatorname{mean}} h_i, \qquad \bar{\kappa}^2 := \underset{i}{\operatorname{mean}} \kappa_i^2, \qquad \rho := \frac{\bar{h}}{\bar{\kappa}}$$

**Interpretation**: MSR measures the average margin (correct minus incorrect logits) in units of the incorrect logits' spread. Larger MSR = more robust to output perturbations.

---

## Empirical Validation / Results

### Main Result: Fine-Tuning Gap Reversal

**Key finding**: At small pretraining budgets (2.4B tokens), Muon Pareto-dominates AdamW. Beyond 2.4B tokens, AdamW achieves lower fine-tuning loss with less forgetting (Figure 1). This holds in both evaluation settings and across all four fine-tuning datasets.

### Crossover TPR Decreases with Model Size

| Model Size | Crossover TPR (StackMathQA) |
|------------|---------------------------|
| 30M | 320 TPR |
| 60M | ~40 TPR |
| 600M | 10 TPR |

The specific crossover varies by dataset, but the trend is consistent: **as model size grows, Muon's advantage reverses at substantially smaller data budgets**.

### Driving Factors

The reversal is driven by two factors (Figure 3b, c):
1. **Shrinking pretraining loss advantage**: Muon's advantage over AdamW in pretraining loss narrows as model size increases.
2. **Increased degradation**: For fixed TPR and fine-tuning loss, Muon forgets increasingly more than AdamW as model size grows.

### Theoretical Result: One-Step Analysis

**Theorem 4.2 (Informal)**: In the Gaussian associative memory with data budget $\gamma = N/d$, for sufficiently large $\gamma$:

- GD has a larger MSR than Muon
- Increasing $\gamma$ increases the ratio $\mathcal{R} = \rho^{GD}/\rho^{Muon}$
- The ratio approaches $3\pi/8 \approx 1.18$ from below:

$$-\frac{C}{\gamma} \leq \mathcal{R} - \frac{3\pi}{8} \leq -\frac{c}{\gamma}$$

where $0 < c < C$ are absolute constants.

**Key intermediate results** (from proof sketch):

$$\bar{h}^{GD} = \frac{\|G\|_F^2}{N-1}, \qquad \bar{h}^{Muon} = \frac{\|G\|_*}{N-1}$$

$$(\bar{\kappa}^{GD})^2 = d^3(\gamma + 2)(1 + o_p(1))$$

$$(\bar{\kappa}^{Muon})^2 = d\left[1 + O(\gamma^{-3/2}) + o_p(1)\right]$$

Leading-order MSRs:

$$\rho^{GD} \approx \sqrt{\frac{d}{\gamma}}\left(1 - \frac{1}{\gamma}\right), \qquad \rho^{Muon} \approx \frac{8}{3\pi}\sqrt{\frac{d}{\gamma}}\left(1 - \frac{1}{10\gamma}\right)$$

### Simulations Beyond One Step

After $t = 10^4$ steps with $d = 256$, GD's MSR overtakes Muon's as $\gamma$ grows, and perturbation robustness ($\sigma_{50}$—the perturbation size at which half of facts are lost) follows the same ordering (Figure 4).

### Language Model Validation

Measuring MSR on held-out DCLM text for 30M–100M models (Figure 5):

- **Muon has larger MSR at small data budgets**
- **AdamW has larger MSR at sufficiently large budgets**
- The MSR ratio $\rho^{AdamW}/\rho^{Muon}$ crosses one as budget grows, at every model size

This mirrors the fine-tuning reversal, though crossovers need not occur at exactly the same budget since the learning-forgetting tradeoff also depends on pretraining loss.

---

## Theoretical and Practical Implications

### Theoretical Implications

1. **Robustness is an intrinsic property**: The MSR captures a checkpoint's robustness to output perturbations that training loss does not capture. A larger MSR provably guarantees more competitors stay ranked below the correct answer under isotropic Gaussian noise (Theorem B.3).

2. **Optimizer-specific generalization**: Muon and GD (and by extension AdamW) reach qualitatively different solutions rather than faster/slower routes to the same minimum. The MSR difference is fundamental to the checkpoint, not specific to fine-tuning or architecture.

3. **Connection to perturbation robustness**: MSR strongly correlates with robustness to Gaussian weight perturbations (Section B.2.2), linking to prior work (Springer et al., 2025) on how pretrained parameter sensitivity affects fine-tuning.

### Practical Implications

1. **Pretraining loss is a poor proxy for adaptability**: Lower pretraining loss does not by itself yield a model that adapts better, especially as data budgets grow. This extends beyond Muon to any pretraining choice (hyperparameters, data, architecture).

2. **Optimizer selection requires more than loss comparison**: The choice between Muon and AdamW should consider downstream behavior, not just pretraining performance.

3. **MSR as a design target**: The MSR is a potential candidate for explicitly optimizing adaptability during pretraining.

4. **Scaling warning**: Since crossover TPR decreases with model size, Muon's fine-tuning disadvantage becomes more pronounced at scale—a critical consideration for frontier model development.

---

## Conclusion

### Main Takeaways

1. **Pretraining data scale reverses the Muon–AdamW fine-tuning gap**: Beyond a crossover budget, AdamW outperforms Muon after fine-tuning despite Muon achieving lower pretraining loss.

2. **The crossover shrinks with model size**: From 320 TPR (30M) to 10 TPR (600M), making the issue increasingly relevant at scale.

3. **The margin-to-spread ratio explains the reversal**: MSR tracks robustness to output perturbations; Muon has larger MSR at small budgets, AdamW at large budgets, mirroring fine-tuning behavior.

4. **The trend is theoretically grounded**: Proven in a linear associative memory model and validated in real language models.

### Future Directions

- **Explicitly optimize for adaptability**: Design pretraining objectives targeting properties like MSR rather than only pretraining loss.
- **Broader applicability**: The lesson likely extends beyond Muon to any pretraining choice where loss is used as the selection criterion.
- **Fine-tuning distribution dependence**: The authors note they don't necessarily expect the reversal for every fine-tuning distribution—fine-tuning on data close to pretraining data approximates continued pretraining, where Muon's advantage may persist.

### Limitations

- The crossover MSR and fine-tuning crossovers need not coincide exactly, as the learning-forgetting tradeoff also depends on pretraining loss.
- The theoretical analysis uses exact orthogonalization Muon and full-batch GD, which are simplifications of the practical AdamW vs. Muon comparison.
- Fine-tuning datasets studied (math, instruction following, music) may not represent all downstream tasks.

---

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