Pretraining data size changes which optimizer works best after tuning
Pretraining Data Scale Reverses the Muon-AdamW Fine-Tuning Gap
Machine Learning
Summary
Choosing the best optimizer helps train language models faster and better. This paper finds that Muon, a newer optimizer, works better when the amount of training data is small. However, when the training data is large, the older optimizer AdamW performs better after the model is fine-tuned. The authors explain this switch by looking at how confidently the model separates correct and incorrect predictions, which shifts in favor of AdamW with more data.
What this means in practice
- •For machine learning engineers: Choose the right optimizer depending on pretraining data size to improve fine-tuning results of language models.
- •For natural language processing teams: Adjust optimization strategies for large-scale language modeling tasks based on data and model size to enhance downstream performance.
Authors
Catherine Li, Jingchu Gai, Andrej Risteski, Jacob Mitchell Springer, Aditi Raghunathan
Abstract
Muon has recently emerged as a promising optimizer for pretraining large language models, offering a speedup over AdamW in minimizing pretraining loss. However, it remains unclear whether Muon's advantage persists after fine-tuning. In this paper, we study the downstream behavior of Muon and AdamW and uncover a surprising trend: beyond a sufficiently large pretraining data budget, AdamW outperforms Muon after fine-tuning. Moreover, as model size increases, AdamW overtakes Muon at progressively smaller pretraining token-parameter ratios. To explain this reversal, we view fine-tuning as a perturbation of the outputs and introduce the margin-to-spread ratio, which measures how far the correct logit lies above the wrong ones relative to how widely the wrong logits are spread. In language models, Muon has the larger margin-to-spread ratio at small data budgets and AdamW at large ones, mirroring the fine-tuning trend. In a linear associative memory, we prove that after one update the margin-to-spread ratio of gradient descent gains on that of Muon as the data budget grows, and simulations show that this trend persists after extended training.