AdamW vs Muon

Two identical networks, same initial weights, same minibatches. Only the optimizer differs. Pause at any time to poke the models, then resume.

step 0
The models' view of the planecolour = predicted class, brightness = confidence
AdamW
Muon
Pause snapshotseach pause photographs both models, newest first
Press ▶ Train, then pause. A snapshot of both models lands here every time you pause, so you can compare what they knew at each stop.
Loss
AdamW train Muon train held-out (dots) ┆ pause
Shape of the last update
Optimizerschanges apply live
Data & modelchanges here reset training

What the two optimizers do

AdamW treats every weight as its own one-dimensional problem. It keeps a running mean m and a running mean square v of each weight's gradient, divides one by the square root of the other, and shrinks the weight a little (decoupled weight decay):

m ← β₁·m + (1−β₁)·g v ← β₂·v + (1−β₂)·g² W ← W·(1 − η·λ) − η · m̂ / (√v̂ + ε)

Muon (MomentUm Orthogonalized by Newton–Schulz) treats each hidden weight matrix as one object. It keeps a momentum matrix and writes it as U Σ Vᵀ. Then it throws Σ away and steps along U Vᵀ, so every direction in the update gets the same step size:

B ← μ·B + (1−μ)·G (momentum) M ← G + μ·(B − G) (Nesterov; plain: M = B) O ≈ U Vᵀ via Newton–Schulz on X = M / ‖M‖_F: X ← a·X + (b·A + c·A²)·X, A = X Xᵀ (a, b, c) = (3.4445, −4.7750, 2.0315), 5× W ← W·(1 − η·λ) − η·√max(1, out/in) · O

Gradients of hidden layers are usually dominated by a few directions (low "stable rank"). AdamW's update inherits that shape. Muon's update is close to orthogonal: rare directions get about as much step as the dominant ones. The Shape of the last update chart shows this live. The AdamW line drops right after the first direction. The Muon line stays high across many directions, then drops only where the momentum holds almost nothing (below about 1/500 of its norm, see the right column). In these tiny networks the gradient spans only a few dozen of the 64 directions, so Muon's "effective rank" is high but not 64.

Muon is meant only for hidden 2D weights. Here the Muon network still uses AdamW (same settings as the AdamW network) for biases, the character embedding, the output layer, and the 2D input layer, as Muon's authors recommend. Tick Muon on every matrix to see why.

Newton–Schulz, in one picture

After dividing by the Frobenius norm, every singular value σ of the momentum lies in (0, 1]. Each iteration maps σ to f(σ) = aσ + bσ³ + cσ⁵. The coefficients are tuned for speed, not exactness: small σ grow by about 3.4× per step, so five steps lift anything above roughly 1/500 of the norm into a band near 1 (about 0.7–1.2) instead of exactly 1. Weaker directions, which are mostly noise, stay small. That is good enough, and it costs a few matrix products per step. The opt ms column shows that cost.

Honest notes