Decomposing Muon
Does Muon's spectral design actually improve generalization, or only convergence speed?
- Training runs
- 100+
- T1 volumes
- 3,079
- Epochs to 85%
- 8 vs 28.5
Overview
Muon orthogonalizes gradient updates through Newton-Schulz iteration. It trains faster than SGD and Adam on standard benchmarks, and the claim attached to it is that this spectral design also generalizes better. This project tested that claim on 3D medical imaging under class imbalance — 3,079 ADNI T1 volumes from 1,125 subjects across 63 sites.
The stack was built from scratch in PyTorch: an 11.2M-parameter 3D DenseNet-121 with gradient checkpointing, subject-level splits frozen before any run, and W&B tracking across every configuration. SGD, Adam, Shampoo, and Muon were benchmarked over more than 100 runs.
Challenges
The reference Muon implementation had a defect: 5D Conv3d weights reached Newton-Schulz orthogonalization without being matricized first, across all 116 routed tensors. Correcting it was a prerequisite to any result meaning anything.
Every Muon arm paired the optimizer with an auxiliary AdamW on non-dense parameters — a confound capable of explaining the whole gap on its own. A matched SGD + AdamW control across 5 seeds ruled out routing and precision as explanations.
Separating the two mechanisms took a 5-seed 2×2 ablation — 60 runs covering momentum with Newton-Schulz, Newton-Schulz alone, momentum alone, and neither — evaluated with subject-clustered bootstrap confidence intervals, permutation tests, and decision rules pre-registered against a stated detection floor.
Decisions
The calls that shaped this project, the alternatives they beat, and what each one cost.
Fix the reference implementation before benchmarking anything
over run the comparison on the published code as-is
WhyThe reference passed 5D Conv3d weights to Newton-Schulz without matricizing them, across all 116 routed tensors — so orthogonalization was operating on the wrong shape. Benchmarking on top of that measures the bug, not the optimizer.
Trade-offResults are no longer directly comparable to numbers published using the unfixed path.
A 5-seed 2×2 ablation of momentum and orthogonalization
over a head-to-head benchmark against SGD and Adam
WhyA head-to-head win tells you the optimizer is better; it doesn't tell you which mechanism is responsible. Only isolating the two components makes the finding transferable to any other optimizer.
Trade-off60 runs for a result that resists being summarized in a single number.
Run a matched SGD + AdamW control
over accept the Muon-versus-SGD gap at face value
WhyEvery Muon arm paired the optimizer with an auxiliary AdamW on non-dense parameters. That confound alone could produce the entire observed gap, so it had to be ruled out before the gap could be attributed to Muon at all.
Trade-offFive more seeds of compute spent eliminating an explanation rather than pursuing one.
Subject-clustered bootstrap CIs and permutation tests
over per-volume confidence intervals
Why3,079 volumes come from 1,125 subjects, so volumes are not independent samples. Treating them as independent produces confidence intervals that are far too narrow and an effect that looks real when it isn't.
Trade-offWider intervals and less apparent statistical power.
Pre-register decision rules against a stated detection floor
over report whatever the runs showed
WhyThe hypothesis going in was that Muon generalizes better. Fixing the decision rules and the detection floor beforehand is what makes a null result a finding rather than a failure to find something.
Trade-offCommits you in advance to an answer you may not like — here, that the effect isn't there.
Outcomes
The spectral mechanism is an interaction, not a main effect. The leading singular value grows only when momentum and orthogonalization are both active, and it scales monotonically with learning rate across all 5 seeds.
That mechanism yields no discriminative benefit. Test AUROC was flat across all four optimizers against a 0.05 detection floor. Muon's advantage is confined to convergence speed — the 85% threshold reached in 5/5 seeds at a median of 8 epochs against SGD's 28.5 — and to operating point.
A pre-registered negative result, reported as one. The practical read: reach for Muon when training time is the constraint, not when you need a better model.