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

Choice of update geometry

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 — the direction assumes Euclidean geometry, while the step length follows a separate rule.
  • Box (\ell_\infty) → -\eta\,\mathrm{sign}(\mathbf{g}): every coordinate moves the same distance, matching Adam’s per-coordinate equalization (§9.6).

The spectral norm for matrix updates

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

The update keeps the gradient’s singular vectors and replaces its nonzero singular values by one. It equalizes directions rather than coordinates.

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

Effect of Newton–Schulz iterations on the spectrum

singular values after 5 iterations: [0.68, 1.13]
  • 0 iterations: an order of magnitude of spread.
  • After 5 iterations, the singular values lie in a band around 1 (approximately 0.7–1.2), using 15 matrix multiplications.

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)))

Assigning parameter groups

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.

Muon and AdamW on the language model

The comparison uses the same initialization, 2,000-step budget, constant learning rate, and four-point grid, with weight decay disabled. The optimizer direction is the intended difference:

final perplexity: AdamW 2.49, Muon+AdamW 2.29
  • The hybrid finished below AdamW in both single-seed runs while using about half the optimizer state.
  • The margin ranged from modest (PyTorch) to substantial (JAX) under the identical protocol, which shows that small comparisons are protocol-sensitive.

Training loss and test accuracy on a CNN

test accuracy: AdamW 0.916, Muon+AdamW 0.914
  • Training loss: the hybrid reaches the memorization regime much faster.
  • Test accuracy differs by about one point; faster optimization does not imply better prediction on this small, saturated task.
  • Same compression Adam-vs-SGD showed here (§9.6): the verdict depends on workload and metric.

Large-scale results and matched tuning

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.

With matched tuning (Wen et al., 2025), matrix methods are fastest in these experiments, with speedups of about 1.4× at 100M parameters and 1.1× at 1B. Sophia’s 2× did not replicate — reported by its own authors’ group.

AdamW remains the default. Reported Muon gains are typically tens of percent rather than multiples.