AdamW Overtakes Muon After Fine-Tuning as Data Grows

Muon has been the optimizer everyone’s pretraining runs have been chasing, and the paper’s related-work section notes that frontier models such as Kimi and GLM-4.5 were pretrained with Muon or its variants. But nearly every argument about which optimizer is better has stopped at one measure: pretraining loss. Whoever posts the lower loss is treated as the better optimizer.

A new paper from Carnegie Mellon turns that premise into the research question. The result the authors report is counter-intuitive: once the pretraining data budget grows large enough, the optimizer with the lower pretraining loss ends up worse after fine-tuning. The experiment scale is nearly 600 pretraining runs and 29,000 fine-tuning runs, the reversal point sits at 2.4B tokens, and the reversal arrives earlier as models get bigger.

Defining what winning means

The paper does not compare single fine-tuning scores. Instead, for each base model it collects every reachable outcome into a frontier: fine-tuned pretraining loss on one axis (how much was forgotten), fine-tuning loss on the other (how much was learned). Since learning more costs forgetting, only the lower-left Pareto frontier is kept. The two frontiers are compared with a single signed quantity, positive meaning the AdamW frontier dominates and negative meaning Muon’s does.

The data budget is measured as TPR, the token-parameter ratio, or how many pretraining tokens land on each parameter. The key metric is the crossover TPR: the ratio at which the AdamW frontier first overtakes the Muon one.

The models are OLMo-style decoder-only transformers at 30M, 60M, 100M, 300M and 600M parameters, trained on DCLM with the OLMo 2 tokenizer. The fine-tuning side covers four datasets, GSM8K, StackMathQA, Alpaca and MusicPile, with every base model fine-tuned for 20M tokens.

Where the reversal lands

On StackMathQA at the 2.4B-token budget, Muon Pareto-dominates AdamW and the gap is negative. Past that point the direction flips: AdamW reaches a lower fine-tuning loss while forgetting less, and the gap turns positive.

The authors split the reversal into two changes that happen at once. Muon’s pretraining-loss advantage narrows as model size grows, and at a fixed TPR and fixed fine-tuning loss, Muon forgets progressively more than AdamW. The two push in the same direction, which is what lets AdamW overtake at a smaller data budget.

The scale effect shows up in the crossover TPR: on StackMathQA, 30M models need 320 TPR to flip, while 600M models need only 10 TPR. The bigger the model, the earlier AdamW takes the lead.

Swapping the fine-tuning optimizer does not rescue it

A natural suspicion is that the difference comes from the fine-tuning recipe rather than the starting checkpoint. The paper runs that control. On GSM8K, the AdamW-pretrained base has its AdamW-fine-tuned frontier above the Muon-fine-tuned one at every budget, while the two fine-tuning choices on a Muon-pretrained base sit close together with no consistent winner. More telling: at 1.2B tokens both Muon-pretrained frontiers are still below the AdamW-pretrained one, and from 19.2B tokens onward the AdamW-pretrained, AdamW-fine-tuned frontier dominates both. The difference lives in the starting point.

Regularizing the fine-tuning with replay or EWC reduces forgetting for both sides and can push the crossover outward, from 80 TPR to 320 TPR on Alpaca and StackMathQA, but AdamW’s frontier is still better at large budgets on every dataset under both regularizers, with one exception: under EWC on StackMathQA the gap turns negative again at 2,560 TPR, the top of the range.

The mechanism, expressed as a measurable quantity

The explanatory route runs through perturbation. Fine-tuning is treated as perturbing the output distribution, so what matters is how robust the base model is to such perturbations.

The metric is the margin-to-spread ratio (MSR): the gap between the correct logit and the average incorrect logit, normalized by the spread of the incorrect logits, with the standard deviation of the incorrect logits standing in for that spread. A larger MSR means the correct answer is held further clear of its competitors and the prediction tolerates perturbation better. In real language models the authors compute the MSR token by token on held-out DCLM text and take the median over correctly predicted positions, across the three model sizes spanning 30M to 100M: Muon holds the larger MSR at small budgets, AdamW at large ones, the same direction as the fine-tuning reversal.

The theory uses a linear associative memory, with each fact a key-value pair and a single weight matrix trained to associate them. The authors show that from zero initialization, after one update, the ratio of gradient descent’s MSR to Muon’s grows with the data budget, with leading term 3π/8 times (1 minus 9 over 10γ). Since 3π/8 is about 1.178, the leading term crosses one at γ near 5.95, and the paper gives a proof for the range γ at or above 16. Simulations over many steps keep the trend.

One scope point has to be stated plainly: the theory compares gradient descent against Muon, not AdamW, and the linear associative memory is a toy model. The authors do not present it as a direct proof about AdamW.

Limits and a conflicting result

The paper does not line up with one line of work. Ruan et al., in “Muon Learns More Robust and Transferable Features than Adam”, report better transfer for Muon-pretrained models after full fine-tuning on Alpaca, Dolly and WizardLM. The CMU authors read their own low-budget result as consistent with that work while showing the comparison changes as pretraining data grows, and they note that the two studies use different models and downstream evaluations, so this is not a matched replication.

The authors draw their own boundaries. A fine-tuning distribution close to the pretraining distribution amounts to continued pretraining, where the picture may differ. The MSR is a candidate objective rather than a validated training target, and the paper states plainly that it does not establish MSR as a validated training objective or show that MSR causes the reversal. The experiments top out at 600M parameters, the specific crossover TPR depends on the dataset, and the reliable part is the direction of the trend.

So what

The takeaway is not that Muon is bad or that AdamW is coming back. It is methodological: pretraining loss measures how well a model predicts its training distribution, not how easily it can be adapted. And as the data budget grows, that proxy gets worse.

The conclusion section spells it out: hyperparameters, data and architecture are mostly selected by pretraining loss, yet a lower pretraining loss does not by itself yield a model that adapts better. The suggested direction is to treat adaptability as a target alongside loss, with the MSR as one candidate. Whether that indicator can actually be optimized against is left open, and the paper supplies the question rather than a recipe.

(Full text and figures: arXiv paper page)