Muon

Dive into Deep Learning · §9.9

The norm decides the direction
steepest descent under three balls · Newton–Schulz · Muon vs. AdamW · the preconditioning family

Steepest descent depends on the norm

The chapter’s frame, paid off

Linearize and ask for the best step of size \eta:

\mathbf{d}^\star = \operatorname*{argmin}_{\|\mathbf{d}\| \leq \eta} \langle \mathbf{g}, \mathbf{d} \rangle.

Not fully posed until the ball is chosen (Bernstein & Newhouse, 2024):

  • Euclidean ball → -\eta\,\mathbf{g}/\|\mathbf{g}\|_2: the SGD direction — following the gradient was a choice, made silently; the step’s length is a separate dial.
  • Box (\ell_\infty) → -\eta\,\mathrm{sign}(\mathbf{g}): every coordinate moves the same distance — the equalization that is Adam’s real work (§9.6).

Matrices want the spectral norm

A hidden matrix transforms activations: \mathbf{y} = \mathbf{W}\mathbf{x}. The size of an update that matters is what it does to activations:

\|\Delta\mathbf{W}\mathbf{x}\|_2 \leq \|\Delta\mathbf{W}\|_2\, \|\mathbf{x}\|_2.

Steepest descent under \|\Delta\mathbf{W}\|_2 \leq \eta, with \mathbf{G} = \mathbf{U}\boldsymbol{\Sigma}\mathbf{V}^\top:

\Delta\mathbf{W}^\star = -\eta\, \mathbf{U}\mathbf{V}^\top

— keep the gradient’s directions, erase its weights. Per-direction equalization where Adam is per-coordinate.

Embeddings see one-hot inputs — the induced norm is the max row norm, not the spectral norm. They stay with AdamW.

Orthogonalization without an SVD

Odd matrix polynomials act on singular values alone: p(\mathbf{X}) = \mathbf{U}\,p(\boldsymbol{\Sigma})\,\mathbf{V}^\top. Iterate a polynomial with fixed point 1; the tuned quintic (Jordan et al., 2024)

p(x) = 3.4445\,x - 4.7750\,x^3 + 2.0315\,x^5

has slope 3.44 at 0 → five iterations suffice, in bfloat16, all matmuls.

def newton_schulz(M, num_iters=5, eps=1e-7):
    a, b, c = 3.4445, -4.7750, 2.0315
    tall = M.shape[0] > M.shape[1]
    X = M.T if tall else M  # keep the Gram factor X @ X.T small
    X = X / (X.norm() + eps)
    for _ in range(num_iters):
        A = X @ X.T
        X = a * X + (b * A + c * A @ A) @ X
    return X.T if tall else X

Watching the spectrum collapse

singular values after 5 iterations: [0.68, 1.13]
  • 0 iterations: an order of magnitude of spread.
  • 5 iterations: a band around 1 (~0.7–1.2) — \mathbf{U}\mathbf{V}^\top for an optimizer’s purposes, at the cost of 15 matmuls.

Muon in fifteen lines

Momentum buffer → Nesterov blend → Newton–Schulz → shape-scaled step:

\mathbf{M}_t = \mu \mathbf{M}_{t-1} + \mathbf{G}_t, \qquad \mathbf{W}_{t+1} = \mathbf{W}_t - \eta\cdot 0.2\sqrt{\max(m,n)}\;\mathrm{NS}_5(\mathbf{G}_t + \mu \mathbf{M}_t)

0.2\sqrt{\max(m,n)} sets every update’s RMS to 0.2\,\eta — AdamW’s typical RMS, so AdamW-tuned \eta transfers (Moonlight, 2025).

class Muon(torch.optim.Optimizer):
    """Steepest descent under the spectral norm: orthogonalized momentum."""
    def __init__(self, params, lr, momentum=0.95):
        super().__init__(params, dict(lr=lr, momentum=momentum))

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            for p in group['params']:
                buf = self.state[p].setdefault('buf', torch.zeros_like(p))
                buf.mul_(group['momentum']).add_(p.grad)
                G = p.grad + group['momentum'] * buf  # Nesterov momentum
                M = G.reshape(len(G), -1)             # flattens conv kernels
                O = newton_schulz(M).reshape(p.shape)
                p.add_(O, alpha=-group['lr'] * 0.2 * math.sqrt(max(M.shape)))

Dividing the census

Hidden matrices → Muon; embeddings, head, vectors → AdamW (the param-group pattern of §9.7):

Muon:   8 tensors,  393216 parameters
AdamW: 22 tensors,   18972 parameters
  • 8 hidden matrices ≈ 95% of parameters.
  • One buffer each vs. AdamW’s two: state memory nearly halves.

The race, same protocol as §9.6

Same init, 2,000 steps, constant lr, four-point grid each, weight decay off — the only difference is the direction:

final perplexity: AdamW 2.49, Muon+AdamW 2.30
  • The hybrid beat AdamW in both single-seed runs — carrying ~half the optimizer state.
  • The margin ranged from modest (PyTorch) to substantial (JAX) under the identical protocol: small races are protocol-sensitive. Mechanism, not benchmark.

On a CNN: two scoreboards

test accuracy: AdamW 0.916, Muon+AdamW 0.914
  • Training loss: the hybrid reaches the memorization regime much faster.
  • Test accuracy: within about a point — optimizing faster ≠ predicting better on a small, saturated task.
  • Same compression Adam-vs-SGD showed here (§9.6): the verdict depends on workload and metric.

The family tree

  • K-FAC (2015): layer-wise Fisher ≈ Kronecker product — two small inverses.
  • Shampoo (2018): AdaGrad-style two-sided factors; won AlgoPerf’s external-tuning track (~30% faster than tuned AdamW).
  • SOAP (2024): Adam inside Shampoo’s eigenbasis.
  • Muon = Shampoo without the memory:

(\mathbf{G}\mathbf{G}^\top)^{-1/4}\mathbf{G}(\mathbf{G}^\top\mathbf{G})^{-1/4} = \mathbf{U}\mathbf{V}^\top.

  • Lion (2023): the sign branch’s lean member — one buffer, six lines (exercise).

Adoption — and the fair-tuning accounting

In production: Moonlight (≈½ the compute of its AdamW baseline); Kimi K2 — 15.5T tokens with MuonClip, zero loss spikes; GLM-4.5; torch.optim.Muon in core.

Fair tuning deflates headlines (Wen et al., 2025): matrix methods are genuinely fastest, but ~1.4× at 100M params, → ~1.1× at 1B. Sophia’s 2× did not replicate — reported by its own authors’ group.

AdamW is still the default. Muon is the one challenger with a clean derivation and frontier mileage. Gains: tens of percent, not multiples.