The Truing Stand — Muon, Newton-Schulz & Adam

A bicycle wheel is "true" when every spoke pulls the rim into a perfect circle. Muon's orthogonalization step does the same thing to a weight update — and you can watch it happen below, with the real coefficients, on a matrix you control.

Primary sources

Adam treats parameters one at a time. Muon treats a whole matrix as one shape.

Same goal — a well-scaled update — from two different units of analysis.

Adam — elementwise
m_t = β1·m_(t-1) + (1-β1)·g_t v_t = β2·v_(t-1) + (1-β2)·g_t² update = m_t / (√v_t + ε) ← every entry rescaled by its OWN history, independently
Muon — whole-matrix
buf = β·buf + (1-β)·g (SGD-style momentum) update = orthogonalize(buf) ← nearest matrix whose singular values are all ≈1 param -= lr · update
AdamMuon
Unit of analysiseach scalar weight, independentlythe whole 2D weight matrix, as one geometric object
What "well-scaled" meanseach entry's step size matches its own gradient historyevery direction the matrix can stretch gets an equal-size step — no direction dominates
Optimizer state / param2 buffers (m, v)1 buffer (momentum only) — ≈45% less optimizer memory for hidden weights
Where it's usedembeddings, output head, norms, biases — everything, in a hybrid setuponly 2D hidden weight matrices (attention/MLP projections); everything else still uses AdamW
Learning rate meansa per-entry step size"spectral norm per update" — how far the matrix moves as a linear operator

Reported results: 1.35× faster NanoGPT speedrun wallclock, a 1.5B-parameter model hitting its target loss in 10h vs 13.3h for AdamW, at 0.7% computational overhead for the orthogonalization step itself.

The truing stand: watch orthogonalization happen

A matrix maps the unit circle to an ellipse — its two axis lengths are its singular values. "Orthogonalize" means: push both axis lengths to 1, turning the ellipse back into a circle. That's a wheel going "true."

The wheel (a 2×2 update matrix)

Tip: "Already true" keeps both spokes exactly equal at every step — but even it doesn't land exactly on the dashed rim at step 5. Rescaling by the matrix's overall size is part of what the iteration has to fix, even for a perfectly round wheel.

Newton-Schulz passes
step 0 of 12
0 (raw)5 (Muon stops here)12
Dashed circle = "true" (every spoke length = 1). Solid shape = the wheel's current rim. Red spoke = σ₁, blue spoke = σ₂.
σ=1 g=1
The curve is g(σ)=3.4445σ−4.7750σ³+2.0315σ⁵ — every singular value bounces along this SAME fixed curve, independently, once per step. Dashed diagonal = "no change" (g(σ)=σ).
σ₁ (red spoke)
—
σ₂ (blue spoke)
—
target
1.000
Before you drag the step slider past 5 — predict: Muon always stops at exactly 5 Newton-Schulz passes. If it kept going to 20 or 50 passes instead, what would σ₁ and σ₂ do?

Try it: drag the step slider to 12 on any preset. The coefficients (3.4445, -4.7750, 2.0315) were tuned to maximize how fast values approach 1 early on — not to converge exactly, ever. Keller Jordan's own code comment: this produces "something like US'V^T where S' is diagonal with S'ᵢᵢ ~ Uniform(0.5, 1.5)" — a deliberately loose approximation, traded for finishing in 5 cheap matmuls instead of running a slow, exact SVD.

The full recipe

What actually runs inside Muon.step(), in order — verbatim from muon.py.

muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True)
momentum.lerp_(grad, 1 - beta) # buf = β·buf + (1-β)·grad update = grad.lerp_(momentum, beta) if nesterov # Nesterov look-ahead: else momentum # update = (1-β)·grad + β·momentum update = zeropower_via_newtonschulz5(update, ns_steps) # ← the truing stand, above update *= max(1, rows/cols) ** 0.5 # rescale for non-square matrices param -= lr * update

Everything above the orthogonalization line is ordinary SGD-with-momentum. The only new ingredient is the line in the middle — and that line is exactly what the two panels above just showed you happening, step by step.

Quick self-check