Muon Learns More Robust and Transferable Features than Adam
Tianyu Ruan, Fengzhuo Zhang, Shuche Wang, Shihua Zhang
cs.LG, cs.AI
2026-06-08
Under matched budgets, Muon beats Adam and SGD on corrupted images and typos and transfers better, with ViT-S effective rank at 6.88 versus 1.86 for Adam and 3.39 for SGD.
Muon already has a speed story. Cited earlier work puts its wall-clock time at nearly 2× Adam and SGD, and production runs such as DeepSeek and GLM-5 use it. Faster only means the loss drops sooner. Whether the same epoch and token budget produces features that survive corrupted inputs, and that move to new tasks, was still open.
The paper splits feature quality into robustness and transfer. Robustness is performance after image corruption and character-level typos. Transfer is a linear head on a frozen backbone, or a full fine-tune under one shared recipe. Both are read off hidden states. Logit margin is the correct-class logit minus the strongest rival; a wider gap is harder for noise to flip. Effective rank compresses how widely singular-value energy spreads across orthogonal directions.
Muon, Momentum Orthogonalized by Newton-Schulz, treats each weight as a matrix. It momentum-averages the gradient, sets every nonzero singular value of that matrix to 1, and steps along the remaining orthogonal factors. Singular values are the stretch along each orthogonal direction. Flattening them stops the loudest direction from taking the whole step. A full decomposition is expensive, so training uses a few Newton-Schulz iterations. Adam divides each coordinate by a running root-mean-square of that coordinate's squared gradient. SGD follows the raw gradient.
Architecture, data, budget, and evaluation stay fixed. Learning rates and schedules are tuned per optimizer. Vision pretraining is ImageNet-1K: ResNet-18 (11M parameters, 100 epochs) and ViT-S (22M, 300 epochs, DeiT recipe). Language pretraining is FineWeb: GPT-2 (12 layers, 124M) and GPT-2 Medium (24 layers, 354M), at least 2× the Chinchilla token-to-parameter ratio, with the discussion kept under 10× that ratio. SGD does not train these GPTs, so language runs compare Muon only with Adam. Three seeds each.
Corruption eval does not train further. Images use ImageNet-C: 15 corruptions at 5 severities, grouped as noise, blur, weather, and digital. Text uses a new split, FineWeb10B-C. Before tokenization, characters are inserted, deleted, replaced by a QWERTY neighbor, or swapped with a neighbor.
Transfer splits in two. Vision freezes the backbone and trains a linear classifier on penultimate features, on EuroSAT, Food-101, Flowers-102, and Stanford Cars. Language fully fine-tunes both pretrained checkpoints with the same Adam, on Alpaca, Dolly 15K, and WizardLM. The fine-tune optimizer is then not the confound.
Hidden-state probes cover transformers only. Each layer gets a linear probe that maps the state to classes or to the vocabulary, then a margin. CNN feature maps are too wide, so ResNet is skipped. Stacking a layer's hidden vectors yields effective rank and Top-10 energy, the share of spectral energy in the ten largest singular values.
A one-layer linear classifier ties the two measurements together. Classes form m coarse blocks of n fine classes. Each feature is a shared block component plus a private one, with strengths a ≤ b: a dog direction with a beagle direction on top. From a zero matrix, with infinitesimal steps and momentum off, Adam collapses to a coordinate-wise sign update. The representation stays in two complementary directions. Muon flattens singular values, so the imbalance ratio ρ is fixed at √(a²n+b²)/b. Adam's ρ is 1+an/b, strictly larger. Until the loss enters a tiny terminal regime, GD's ρ is larger too. At a matched loss, both margin and effective rank fall as ρ rises, so the smaller ratio wins both.
Muon has the best average on all four corruption suites. ViT-S leads every image-corruption family. GPT-2 and GPT-2 Medium post lower perplexity on every typo type. ResNet-18 has the best mean accuracy and leads on noise and weather, but sits slightly behind SGD on blur and digital. SGD beats Adam on the CNN; Adam beats SGD on the transformers. The main text never lists ImageNet-C accuracies or FineWeb10B-C perplexity gaps.
On corrupted data, ViT-S's final-layer margin is 3.29 for Muon against 3.13 for Adam, 5% higher, and 2.74 for SGD, where Muon is 20% higher. The gap opens in the last two layers. GPT-2 is above Adam on most layers; per-layer numbers are not printed. The vocabulary has 50,257 tokens, so the true next token is often not the probe's top guess and the margin can be negative.
| Model | Metric | Muon | Adam | SGD |
| ViT-S final layer | logit margin | 3.29 | 3.13 | 2.74 |
| ViT-S depth average | effective rank | 6.88 | 1.86 | 3.39 |
| ViT-S depth average | Top-10 energy | 0.890 | 0.964 | 0.920 |
| GPT-2 depth average | effective rank | 16.00 | 11.12 | not compared |
| GPT-2 depth average | Top-10 energy | 0.827 | 0.863 | not compared |
ViT-S effective rank is more than 3× Adam and about 2× SGD. GPT-2 is 16.00 versus 11.12, with lower Top-10 energy. Layers 7 to 10 of ViT-S are the exception: SGD has the higher effective rank there, and Muon leads again in deeper layers.
Across 8 backbone-dataset pairs, Muon matches or beats Adam and SGD. On Alpaca, Dolly, and WizardLM, both model sizes get lower perplexity when the Muon checkpoint is fine-tuned with the same Adam. Point gaps are not in the main text.
With m and n growing at the same rate, at every loss threshold ε in (0, log C), Muon's average margin and effective rank are strictly above Adam's. Against GD the claim excludes an over-training tail, losses as small as exp(−c√(n log n)). Absolute cosine similarity of embeddings, after hierarchical reordering, is blocky for ViT-S classes and for frequent GPT-2 tokens. That is the sketch offered for the shared-plus-private feature assumption.
Most of the Muon conversation is a wall-clock bill. This paper spends the same epochs and tokens and looks at the features. Up to GPT-2 Medium (354M) and ViT-S, corruption robustness, transfer, margin, and effective rank move together. The vision test freezes the backbone, so it reads pretraining features before fine-tuning rewrites them.
Teams already running Muon can track margin and effective rank as a cheap side check. Teams still on Adam or AdamW should not switch a large run on this evidence alone. The biggest model here is 354M, and the downstream point gaps are not stated in the main text.
The paper's own limit is the model family: language models and vision classifiers only. Diffusion is untouched.
The proof is narrower. Momentum is zero, the step size is infinitesimal, and Adam's β1 and β2 are zero, leaving a sign update. The momentum buffer and the Newton-Schulz approximation used in training never enter the theorems. Features are two orthogonal one-hot blocks, and m and n must grow together. The blocky similarity plot appears after hierarchical clustering, which pulls neighbors together, so it is a weak check.
Language has no SGD arm. Full fine-tuning rewrites the weights, so a perplexity gap is not pure evidence about pretrained features. ResNet-18 still trails SGD on blur and digital corruptions. Adversarial robustness, a direct AdamW comparison, and the same budget at production scale are all missing. The edge over GD is also stated only before the loss becomes extremely small.