Muon

Context. In 2024 the optimizers that beat AdamW per step, such as Shampoo and SOAP, lost on wall-clock because their matrix inverse roots are expensive and need float32. Muon (Jordan et al., 2024) changed this in the NanoGPT speedrun, where only time counts. Bernstein and Newhouse (2024) and Bernstein's Deriving Muon (2025) derive it as steepest descent under a norm, and the Polar Express (Amsel et al., 2025) makes its polynomial iteration optimal. For a hidden weight matrix, Muon keeps the singular vectors of the momentum and discards its singular values, so the update $UV^\top$ moves every direction by the same amount. That is steepest descent under a rescaled spectral norm, and five matmul-only polynomial steps in bfloat16 get close enough to $UV^\top$ for under 1% extra FLOPs.

What Muon does. For each 2-D hidden parameter with stochastic gradient $G_t$, Muon runs momentum and orthogonalizes the buffer: $$B_t=\mu B_{t-1}+G_t,\qquad O_t=\mathrm{NS}_5(B_t),\qquad W_t=W_{t-1}-\eta\,O_t,$$ where $\mathrm{NS}_5$ is five Newton-Schulz steps approximating $\mathrm{Ortho}(B)=UV^\top$, the Frobenius-nearest semi-orthogonal matrix to $B=USV^\top$. The reference code uses Nesterov momentum with $\mu=0.95$ and multiplies $O_t$ by $\sqrt{\max(1,\text{rows}/\text{cols})}$; the scale of the buffer is irrelevant because $\mathrm{Ortho}(cB)=\mathrm{Ortho}(B)$. Transformer updates are nearly low rank, so a raw step is almost entirely along a few top singular directions; orthogonalizing raises the directions with small singular values to the same scale. The order matters: Orthogonal-SGDM, which orthogonalizes each gradient and then averages, lost to tuned SGD-momentum, and the post finds averaging first better. Which parameters get Muon is part of the method. Scalars and vectors go to AdamW, and so do the embedding and the output head although they are 2-D, the embedding for a reason given below, the head on empirical grounds only. Q, K and V work better orthogonalized separately than as one fused matrix.

Why orthogonalize. Optimizers as steepest descent under a norm. Steepest descent under a norm splits into a direction, the dual map $\arg\max_{\|t\|=1}g^\top t$, which depends only on the norm, and a length. With the EMAs off, known optimizers become such methods. Adam with $\beta_1=\beta_2=0$ is sign descent, the $\ell_\infty$ case: the unit ball is a cube, its vertex most aligned with $g$ is $\operatorname{sign}(g)$, and the dual norm is $\|g\|_1$. Shampoo with accumulation off is $(GG^\top)^{-1/4}G(G^\top G)^{-1/4}=UV^\top$, the spectral-norm case, by von Neumann's trace bound $$\langle A,G\rangle\le\sum_i\sigma_i=\operatorname{tr}\Sigma\qquad\text{whenever }\|A\|_{\ell_2\to\ell_2}\le1,$$ with equality at $A=UV^\top$: the dual of the spectral norm is the nuclear norm, and $UV^\top$ attains it. So Muon is momentum followed by accumulation-free Shampoo, at Newton-Schulz cost. It is the constrained form, minimizing $\langle G,\Delta W\rangle$ over a spectral ball of radius $\eta$, which gives the fixed-length step $-\eta UV^\top$; the penalized form would scale each step by the nuclear norms of the gradients (our comparison of the two sources). For a linear predictor with square loss averaged over outputs and inputs of norm $\sqrt{d_{in}}$, the spectral norm is also what bounds the loss: the second-order term is at most $\frac12\frac{d_{in}}{d_{out}}\|\Delta W\|_{\ell_2\to\ell_2}^2$, because the bound must hold for every input direction.

Which spectral norm: RMS to RMS, and the width factor. Deriving Muon measures a weight change of a dense layer $y=Wx$ by how far it can move the output. With $\|v\|_{RMS}=\|v\|_2/\sqrt d$ on both sides, $$\|\Delta Wx\|_{RMS}\le\|\Delta W\|_{RMS\to RMS}\,\|x\|_{RMS},\qquad\|W\|_{RMS\to RMS}=\sqrt{\frac{d_{in}}{d_{out}}}\;\|W\|_{\ell_2\to\ell_2},$$ so an RMS-to-RMS ball of radius $\eta$ is a spectral ball of radius $\eta\sqrt{d_{out}/d_{in}}$, and the dualized update is $\Delta W=-\eta\sqrt{d_{out}/d_{in}}\,UV^\top$. When $d_{in}\le d_{out}$ and the gradient has full rank, $UV^\top$ is an isometry, so every unit-RMS input has its output moved by exactly $\eta$ whatever the width: a fixed learning rate transfers across width, which is what $\mu$P (the maximal-update parametrization) achieves, here derived from a norm. The reference factor $\sqrt{\max(1,\text{rows}/\text{cols})}$ agrees on a $4d\times d$ up-projection (both give 2) and differs on the $d\times4d$ down-projection ($\tfrac12$ against 1); which scales better is not shown. The same logic explains the embedding: it maps one-hot vectors to dense ones and gets the $\ell_1\to RMS$ norm, whose dual step normalizes each token's column separately (our derivation), whereas orthogonalization couples the tokens through shared singular vectors.

def newton_schulz5(G, steps=5):
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16() / (G.norm() + 1e-7)   # all singular values now at most 1
    if G.size(0) > G.size(1):
        X = X.T                            # use the smaller Gram matrix X X^T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X                  # each singular value s becomes a s + b s^3 + c s^5
    if G.size(0) > G.size(1):
        X = X.T
    return X

The optimal coefficients (Polar Express), and why practical Muon is a spectral method. The worst-case spectral error over matrices with singular values in $[\ell,u]$ is the scalar error $\max_{[\ell,u]}|1-p|$ of the composed polynomial. Rescaling the input preserves odd polynomials, so the only state of the iteration is the ratio $\ell_t/u_t$ of the current interval, and choosing each step greedily as the minimax polynomial on that interval (a small Remez solve, done once offline) is worst-case optimal. Each optimal step maps $[\ell,u]$ onto a narrower $[p(\ell),2-p(\ell)]$; Jordan's fixed quintic cannot shrink its own output interval and stalls at an error of about 0.3. With $u=1$ and $\ell=10^{-3}$, near bfloat16 precision, five steps map every singular value in $[10^{-3},1]$ into about $[0.85,1.12]$, against $[0.47,1.20]$ for Jordan's five steps, and the bfloat16 safety fixes cost about 0.04 of worst-case error relative to pure greedy (our re-implementation). More precision does not help: in the paper's ablations an exact SVD trains no better than five or six steps, and flipping or zeroing the update on singular directions below $10^{-4}\sigma_{\max}$ changes nothing (below $10^{-3}\sigma_{\max}$, very little). So practical Muon is a spectral method. It keeps the singular vectors exactly and moves the singular values that matter into a band around 1, and the ablations show that this is enough.

The speedrun evidence and the cost. Muon's introduction (record #3) and first improvements (#4) took the speedrun from 31.4 to 22.3 minutes; the post credits it with a 1.35x speedup, and at 1.5B parameters Muon matched tuned AdamW's final loss with 25% fewer tokens. The benchmark makes this credible: Muon costs more per step, so if AdamW could be tuned to Muon's per-token efficiency, putting it back would set a record, and Muon had survived 12 further records by 7 researchers when the post was written. For an $n\times m$ matrix with $m\le n$, $T$ Newton-Schulz steps cost at most $6Tnm^2$ FLOPs against $6nmB$ for the layer's forward and backward pass on $B$ tokens, so the overhead is at most $Tm/B$: the orthogonalization cost is fixed by the shape while the training cost grows with the batch. That is 0.7% at the speedrun's $m=768$, $B=524{,}288$ and 0.5% at Llama-405B's sizes; the state is one momentum buffer, half of Adam's. Polar Express replaced Newton-Schulz in record #38 (2.483 to 2.476 minutes); the README states about 1 s from 10 fewer steps, with a safety factor of 1.02 and a higher Muon learning rate, and says its author could not reproduce gains as large as the paper's (0.05 to 0.06 in GPT-2 validation loss over Jordan's coefficients at 1B tokens, under 0.01 at 10B).