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 / (jnp.linalg.norm(X) + 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).

def scratch_muon(learning_rate, momentum=0.95):
    def init(params):
        return jax.tree.map(jnp.zeros_like, params)
    def update(grads, bufs, params=None):
        bufs = jax.tree.map(lambda b, g: momentum * b + g, bufs, grads)
        def step(g, b):
            G = g + momentum * b            # Nesterov momentum
            M = G.reshape(-1, G.shape[-1])  # flattens conv kernels
            O = newton_schulz(M).reshape(G.shape)
            return -learning_rate * 0.2 * math.sqrt(max(M.shape)) * O
        return jax.tree.map(step, grads, bufs), bufs
    return optax.GradientTransformation(init, update)

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.69, Muon+AdamW 1.85

• 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.921, Muon+AdamW 0.922

• 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.