AdamW vs Muon · Optimization · Deep Learning · NLP
Same accuracy. Very different path there.
A controlled optimizer study on BERT-small
I started this as a fairly simple optimizer comparison.
Fine-tune BERT-small with AdamW. Fine-tune it again with Muon. Compare the results.
The final validation accuracy ended up almost identical.
That made the experiment more interesting, not less. If both optimizers arrive at roughly the same task performance, what actually changed during training?
- Model
- prajjwal1/bert-small (28.8M parameters)
- Task
- SST-2 sentiment classification
- Stack
- PyTorch · Hugging Face · NumPy
- BERT-small
- SST-2
- 3 paired seeds
- Custom PyTorch training loop
- Calibration + gradient + sharpness analysis
The question
If two optimizers reach almost the same accuracy, did they actually learn in the same way?
The point isn’t whether Muon beats AdamW. It’s that two optimizers can end up at nearly the same task accuracy while taking meaningfully different optimization paths.
I wanted the comparison to be boringly controlled.
The model was prajjwal1/bert-small, fine-tuned on the full SST-2 training split and evaluated on its 872-example validation split.
Hugging Face loads the pretrained model, the tokenizer and the dataset. The training loop itself is written directly in PyTorch.
That matters here. The experiment tracks things that high-level trainer abstractions usually hide:
- gradient norms
- clipping behavior
- actual parameter updates
- validation cross-entropy
- calibration
- generalization gaps
- sharpness
- displacement from initialization
Muon couldn’t simply replace AdamW everywhere.
Muon is designed for 2D weight matrices. So the Muon run is a hybrid: Muon updates the 24 internal attention and feed-forward weight matrices, and an auxiliary AdamW handles everything else.
The baseline uses AdamW for every parameter. Everywhere on this page, “Muon” means this hybrid setup, not a pure Muon run.
BERT-small
- all parametersAdamW
BERT-small
- 24 attention / FFN matricesMuon
- embeddings · biases · LayerNorm · pooler · classifierAdamW
The comparison only matters if the setup is fair.
Each AdamW–Muon pair shared the same seed, which gives both runs the same classifier initialization and minibatch order. Everything else was held fixed:
- Epochs
- 3 (6,315 updates)
- Batch size
- 32
- Max sequence length
- 128
- Warmup
- 6% linear
- LR decay
- linear
- Gradient clipping
- global norm 1
- Weight decay
- 0
- Seeds
- 1, 2, 3 (paired)
- AdamW learning rate
- 3 × 10⁻⁵ for every parameter
- Hybrid Muon learning rate
- 10⁻³ on the 24 matrices · 3 × 10⁻⁵ auxiliary AdamW
Why weight decay = 0
The two optimizers used different learning rates, about 33× apart on the matrices.
With decoupled weight decay, the same nonzero coefficient would not imply the same effective shrinkage.
So keeping weight decay at zero removes one avoidable confound.
Accuracy basically tied.
AdamW
88.57% ± 0.18
Hybrid Muon
88.69% ± 0.93
Validation accuracy, mean ± SD over three paired seeds.
That difference is not the interesting part.
The roughly 0.1-point gap is smaller than Muon’s own seed-to-seed variation.
Muon got to low training loss faster.
AdamW
4,743 ± 559
Hybrid Muon
2,958 ± 686
optimizer steps to rolling training loss ≤ 0.10
Muon reached rolling training loss 0.10 in roughly 38% fewer optimizer steps.
Faster optimization on the training objective did not translate into a clearly better final validation result.
optimization speed ≠ generalization (in this experiment)
AdamW looked steadier at the end.
Validation accuracy was effectively tied. But the AdamW runs ended with lower validation cross-entropy, and lower variance across seeds.
AdamW
0.418 ± 0.014
Hybrid Muon
0.481 ± 0.020
final validation cross-entropy
Accuracy only asks whether the final class was right.
Cross-entropy also cares about how much probability the model assigned to the correct answer.
Why two correct answers can differ
Model A: correct class probability = 0.51 Model B: correct class probability = 0.95 Both are correct by accuracy. They are not equally confident, or equally calibrated.
Correct isn’t the same as well-calibrated.
A classifier can have good accuracy while being badly overconfident. Calibration asks whether the model’s probabilities mean what they appear to mean.
If a model says “90% confident” on 100 examples, we’d like it to be correct on roughly 90 of them.
| Metric | AdamW | Hybrid Muon |
|---|---|---|
| Brier score | 0.100 ± 0.003 | 0.102 ± 0.007 |
| Expected calibration error (10 bins) | 0.084 ± 0.005 | 0.093 ± 0.008 |
| Mean confidence on mistakes | 0.902 ± 0.008 | 0.929 ± 0.008 |
AdamW came out slightly better on all three, and Muon was more confident when it was wrong.
The gradient isn’t the update.
For plain SGD, the update is just a scaled gradient. Adaptive and transformed optimizers don’t simply move along the raw gradient.
So the project measures the update actually applied, every 100 steps, on the same 24 matrices for both optimizers.
Plain SGD
Δθ ≈ −η g
The step is the gradient, scaled.
What’s measured
cos(gₜ, Δθₜ)
How aligned the actual move Δθₜ is with the clipped gradient gₜ.
| Metric | AdamW | Hybrid Muon |
|---|---|---|
| Relative update size ‖Δθ‖ / ‖θ‖ | 7.87 × 10⁻⁵ | 37.73 × 10⁻⁵ |
| Gradient–update cosine | 0.293 ± 0.009 | 0.109 ± 0.002 |
Muon’s updates on those matrices were about 4.8× larger and much less aligned with the instantaneous gradient: evidence that orthogonalized momentum changed the trajectory, not just its speed.
Two optimizers can see the exact same gradient and still move through parameter space differently.
Both runs had the same speed limit.
Global gradient clipping was set at τ = 1. If the gradient norm exceeds the threshold, the whole gradient is scaled down to it.
g′ = g · min(1, τ / ‖g‖), τ = 1
| Metric | AdamW | Hybrid Muon |
|---|---|---|
| Fraction of updates clipped | 0.990 ± 0.002 | 0.884 ± 0.002 |
If an optimizer repeatedly sees gradients large enough to trigger clipping, the shared clipping rule shapes its trajectory more heavily.
Here, clipping was active on most steps for both runs, so the shared threshold was itself an influential part of the setup.
Training loss is not the whole story.
The project tracks validation cross-entropy minus cross-entropy on a fixed 4,096-example training subset. It’s a simple overfitting diagnostic, not a generalization bound.
| Metric | AdamW | Hybrid Muon |
|---|---|---|
| Train–validation CE gap | 0.342 ± 0.015 | 0.449 ± 0.017 |
Muon’s larger gap suggests that, under a fixed three-epoch budget, fitting faster also meant more overconfidence or overfitting.
Then I started poking the loss landscape.
If I perturb the final parameters slightly, how quickly does the loss get worse?
Random-direction sharpness
Apply matched perturbations in ten random directions, scaled to a fixed fraction of each tensor’s norm, and measure how much the loss changes.
Adversarial sharpness
Use a short projected-gradient search to find nearby parameters that increase the loss as much as possible.
At matched training loss for seed 1, Muon was less sensitive to both matched random perturbations and the projected-gradient search.
At a 10% per-tensor radius, the worse-side loss increase was 0.0191 for AdamW and 0.00554 for Muon.
Caveat
This is evidence of a flatter neighborhood under the chosen parameterization. Sharpness changes under parameter rescaling, so it isn’t a parameterization-invariant claim, and it’s one seed.
Comparing checkpoints at the same epoch can be misleading.
If Muon reaches lower training loss earlier, then comparing epoch 2 of AdamW with epoch 2 of Muon doesn’t compare equivalent optimization states.
So training saved checkpoints when the rolling training loss crossed fixed levels, 0.15 and 0.10, and the sharpness analysis compares those.
Same epoch
AdamW @ epoch 2 vs Muon @ epoch 2
Different training losses, different optimization states.
Same training loss
AdamW @ loss 0.10 vs Muon @ loss 0.10
Checkpoints saved at 0.15 and 0.10.
When the two optimizers have achieved roughly the same training objective, how different are their local geometries?
On the 512 fixed examples used for sharpness, the matched seed-1 checkpoints had base losses of 0.07895 (AdamW) and 0.07861 (Muon).
It made a difference: at the final, unmatched-loss checkpoints, the smallest-radius ordering differed.
How far did fine-tuning move the pretrained model?
The project also measures how far fine-tuning moved the parameters from the pretrained initialization, per matrix, and the stable rank of that displacement.
‖θ_fine-tuned − θ_pretrained‖
In seed 1, Muon’s displacement on the 24 matrices had a much higher stable rank: its changes were spread across many more directions (rightmost panel of the update-geometry figure).
It’s another way of asking whether the two optimizers modify the pretrained model in structurally different ways, not proof of anything about representation quality.
Same accuracy doesn’t mean the same mistakes.
The project stores per-example predictions, so two runs can be compared example by example.
Two classifiers can both get about 89% while disagreeing on many examples. McNemar’s test asks whether the pattern of disagreements points to a real difference.
| Muon correct | Muon wrong | |
|---|---|---|
| AdamW correct | a | b |
| AdamW wrong | c | d |
Only b and c matter: the examples where the two models disagree.
seed 1
b = 22 · c = 29
p = 0.40
seed 2
b = 33 · c = 23
p = 0.23
seed 3
b = 24 · c = 30
p = 0.50
No seed showed a significant difference. The models disagreed on a few dozen examples each time, in both directions.
I kept the experiment inspectable.
Training deliberately avoids the Hugging Face Trainer, so every measurement sits in plain code:
src/
- data.pyloads SST-2, tokenizes examples, builds data loaders
- optimizers.pyAdamW and the hybrid Muon parameter split
- measurements.pycalibration and optimizer-update metrics
- train.pythe training and evaluation loop
- compare_runs.pymulti-seed aggregation and paired McNemar tests
- sharpness.pyrandom-direction sensitivity
- adversarial_sharpness.pysearches nearby parameters for high-loss perturbations
- checkpoint_analysis.pydisplacement from pretraining and matrix rank
- plot_*.pyfigures generated from the saved outputs
I wanted the runs to be reproducible, not just screenshot-able.
- Three paired seeds
- Saved configuration for every run
- Saved metrics and per-example predictions
- Saved analysis outputs
- Loss-matched checkpoints
- Completed runs are never overwritten
- Figures generated from stored outputs
Checkpoint weights are generated locally and not versioned: the 18 saved weight files total roughly 2 GB. The compact metrics, predictions, analyses and figures are in the repository.
What I found
01
Final accuracy was effectively tied
AdamW 88.57% ± 0.18 · hybrid Muon 88.69% ± 0.93.
02
Muon optimized the training objective faster
Rolling training loss 0.10 in about 38% fewer optimizer steps.
03
AdamW ended with stronger validation-loss behavior
Lower final validation cross-entropy, and lower variance across seeds.
04
The optimization geometry differed
At matched training loss (seed 1), Muon was less sensitive to matched random perturbations and to projected-gradient search: a flatter local neighborhood under the chosen parameterization.
What this experiment does not prove
- It doesn’t establish that Muon is better than AdamW, or the reverse.
- Three seeds are useful, but not enough for sweeping conclusions. Sharpness used one.
- It’s one model, one task, and one untuned learning-rate pair.
- The Muon setup is hybrid, not pure Muon.
- Sharpness depends on parameterization and measurement choices.
- Faster training-loss convergence doesn’t imply better generalization.
- Nearly equal accuracy doesn’t imply identical optimization dynamics.
Why I liked this experiment
This was one of those projects where the boring result ended up being the interesting one.
If Muon had simply beaten AdamW by five points, there would not have been much to think about.
Instead, they landed at basically the same accuracy.
That forced me to look at everything the final metric was hiding: speed, calibration, update geometry, variance, sharpness, and how far the model actually moved from pretraining.
I like experiments like that.
The answer gets less clean, but the model gets more interesting.