Pretraining Data Scale Reversal: Why Muon's Fine-Tuning Advantage Is Overtaken by AdamW

Introduction: A New Twist in the Optimizer Debate

Muon optimizer has recently emerged in large language model pretraining, reducing pretraining loss faster than AdamW. However, a new study reveals a surprising phenomenon: when the pretraining data scale is sufficiently large, AdamW actually outperforms Muon after fine-tuning. This finding challenges the intuition that "Muon is universally superior to AdamW" and prompts us to re-examine the relationship between optimizer choice and data budget.

Muon optimizer is a recently proposed optimization algorithm for neural network training. Its core idea is to approximately orthogonalize gradient matrices (typically via Newton-Schulz iteration), thereby maintaining the "orthogonality" of update directions during parameter updates and avoiding interference between updates across different dimensions. Unlike AdamW, which relies on element-wise adaptive learning rates, Muon focuses more on leveraging the spectral structure of gradient matrices. In pretraining GPT-style models, Muon has been reported to reduce training loss faster under the same compute budget, making it a strong competitor to AdamW. However, Muon's design primarily targets loss minimization during pretraining, and its behavior in downstream fine-tuning has lacked systematic study.

Core Finding: How Data Budget Reverses Fine-Tuning Performance

Through systematic experiments, the paper finds that Muon's advantage in pretraining does not simply carry over to fine-tuning. Specifically, there exists a critical point in the pretraining data budget: beyond this point, AdamW begins to lead in downstream performance after fine-tuning. More interestingly, as model size increases, the token-parameter ratio required for AdamW to overtake Muon gradually decreases. This means that larger models may exhibit the reversal at relatively smaller data budgets.

This trend is counterintuitive, since Muon typically performs better on pretraining loss. The researchers treat fine-tuning as a perturbation to model outputs and introduce the margin-to-spread ratio to explain this phenomenon. This metric measures how much the correct logit leads the incorrect logits, as well as the dispersion of the incorrect logits. In language models, Muon has a larger margin-to-spread ratio under small data budgets, while AdamW has a larger one under large data budgets, which precisely aligns with the fine-tuning trends.

Token-parameter ratio is a commonly used metric for measuring the pretraining data budget relative to model size, typically expressed as "total training tokens ÷ model parameters." For example, a 1-billion-parameter model trained on 20 billion tokens has a token-parameter ratio of 20. The Chinchilla scaling laws indicate that compute-optimal training requires approximately 20 tokens per parameter, but in practice this ratio can range from single digits to hundreds. A larger ratio means the model has seen more sufficient data during pretraining. This paper finds that the token-parameter ratio required for AdamW to overtake Muon decreases as model size increases, indicating that larger models exhibit the reversal at relatively smaller data budgets—a finding with direct implications for optimizer selection in large model training.

Theoretical Explanation: Margin-to-Spread Ratio and Linear Associative Memory

To gain deeper insight into the reversal mechanism, the authors conduct a theoretical analysis in a linear associative memory model. The proof shows that after a single update, as the data budget grows, the margin-to-spread ratio of gradient descent (GD) gradually exceeds that of Muon. Simulation experiments further show that this trend persists with extended training. This provides mathematical support for "data scale reversing the fine-tuning gap" and suggests essential differences between optimizers under different data regimes.

This theoretical framework links fine-tuning performance to the geometric properties of model outputs, offering a new perspective for optimizer design. It suggests that optimizer selection should not focus solely on pretraining loss but should also consider the fine-tuning dynamics of downstream tasks.

Linear Associative Memory is a classic class of neural network models used to store and retrieve input-output associations. Its basic form is: given input vector x and output vector y, the model learns a mapping Wx≈y through a weight matrix W, typically trained via least squares or gradient descent. Because of its analytical tractability, researchers often use it as a "minimal reproducible model" for theoretical analysis of optimizer behavior. This paper leverages this framework to rigorously compare the output geometric properties of gradient descent (GD) and Muon updates in a simplified setting, proving that as the data budget grows, GD's margin-to-spread ratio exceeds that of Muon. Although this theoretical result is based on a linear model, it provides an interpretable mathematical starting point for understanding optimizer differences in real language models.

Practical Implications: How to Choose an Optimizer

Based on these findings, the research team recommends that practitioners weigh optimizer choice according to pretraining data budget and model size. For small models with limited data budgets, Muon may still be the better choice after fine-tuning; for large models with large data budgets, AdamW may be more advantageous during fine-tuning. This conclusion has guiding significance for practical training pipelines, especially when resources are limited and trade-offs must be made between pretraining and fine-tuning investment.

Furthermore, the margin-to-spread ratio can serve as a diagnostic tool to predict fine-tuning performance, helping researchers make more informed decisions early in training. Future work could further explore the universality of this metric across other architectures and tasks.

Conclusion

The fine-tuning gap between Muon and AdamW is not static; pretraining data scale is a key moderating variable. AdamW's comeback under large data budgets, together with the theoretical explanation via the margin-to-spread ratio, jointly reveals the complexity of optimizer behavior. This research reminds us that while pursuing pretraining efficiency, we must also pay attention to actual downstream fine-tuning performance, avoiding the neglect of overall pipeline optimization due to a single metric.

One-line Summary

When pretraining data scale is sufficiently large, AdamW outperforms Muon in fine-tuning, and the margin-to-spread ratio explains this reversal.

Paragraph Summary

This article challenges the intuition that Muon optimizer is universally superior to AdamW, revealing the key moderating role of pretraining data budget on fine-tuning performance. The study finds a critical data budget point: beyond it, AdamW begins to lead in downstream performance after fine-tuning; and the larger the model, the smaller the token-parameter ratio required for AdamW to overtake. The authors introduce the margin-to-spread ratio as an explanatory metric, which measures the lead of correct logits and the dispersion of incorrect logits. Under small data budgets Muon has a larger margin-to-spread ratio, while under large data budgets AdamW does, consistent with fine-tuning trends. Theoretical analysis in a linear associative memory model proves that gradient descent's margin-to-spread ratio surpasses Muon's as data budget grows, providing mathematical support. In practice, it is recommended to weigh optimizer choice based on data budget and model size, and the margin-to-spread ratio can be used as a diagnostic tool to predict fine-tuning performance.