Muon

Who gets Muon vs AdamW?

Muon is not a drop-in for every tensor. Scalars, vectors, and input / output layers stay on AdamW (or similar). Hidden 2D weights go to Muon. Conv filters: flatten the last three dims.

Hidden 2D weights use Muon; embeddings, heads, and 1D parameters stay on AdamW.

Ill-conditioned matrix updates

Orthogonalization, in plain terms: keep the directions of a momentum update but rebalance their scale so no single singular direction dominates. Rare directions that still matter for learning get scaled up; the step becomes closer to full rank.

Manual inspection in the post: SGD-momentum and Adam updates on transformer 2D weights are almost low-rank — a few singular directions carry most of the energy until Ortho spreads it.

Theory flavours: steepest descent under spectral norm / accumulation-free Shampoo (Bernstein & Newhouse). Empirics: raise small singular values.

Gradient → momentum → Newton–Schulz → apply

Momentum runs before orthogonalization (opposite order from Orthogonal-SGDM). That ordering wins in Jordan’s tests.

Source: Muon design section — momentum update, then NS post-process, then parameter step.

Why NS orthogonalizes

Newton–Schulz works on singular values one at a time. Write the momentum update as G = USVᵀ. Each NS step applies a small polynomial φ to every singular value; after scaling and repeating about five times, they land near 1 — which is what orthogonalization needs.

Formally, one NS step with coefficients (a,b,c) maps singular values by φ(x) = ax + bx³ + cx⁵. After normalize + N compositions: U φᴺ(S) Vᵀ → UVᵀ when φᴺ(x) → 1 for x ∈ (0,1].

Blog reference vs our implementation

Left: Keller’s blog kernel (as cited on the post). Right: the function we actually train with — same coeffs, same φ matmul loop; batched .mT / last-two-dim norms, plus our muon_update wrapper.

def newtonschulz5(G, steps=5, eps=1e-7):
    assert G.ndim == 2
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G.bfloat16()
    X /= (X.norm() + eps)
    if G.size(0) > G.size(1):
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    if G.size(0) > G.size(1):
        X = X.T
    return X

# zeropower ≈ Ortho(G) = UVᵀ (drop singular values; Keller blog: newtonschulz5)
def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    assert G.ndim >= 2
    a, b, c = (3.4445, -4.7750, 2.0315)  # tuned NS polynomial coeffs (same as blog)
    X = G.bfloat16()
    # Tall matrix: transpose so the cheap matmul is X @ X.mT (min side squared).
    if G.size(-2) > G.size(-1):
        X = X.mT
    # Frobenius normalize → singular values land in [0, 1] (batched over last two dims).
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    for _ in range(steps):
        A = X @ X.mT
        B = b * A + c * A @ A   # degree-5 polynomial in A via matmuls
        X = a * X + B @ X       # φ(X) = aX + (bA + cA²)X
    if G.size(-2) > G.size(-1):
        X = X.mT                # undo tall-matrix transpose
    return X


# Momentum first, then Newton–Schulz (Muon order). Default: Nesterov lookahead.
def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True):
    momentum.lerp_(grad, 1 - beta)
    update = grad.lerp_(momentum, beta) if nesterov else momentum  # Nesterov vs plain momentum
    if update.ndim == 4:
        update = update.view(len(update), -1)  # 4D conv weights: flatten spatial dims
    update = zeropower_via_newtonschulz5(update, steps=ns_steps)
    update *= max(1, update.size(-2) / update.size(-1)) ** 0.5  # aspect-ratio scale max(1, n/m)^0.5
    return update

One-to-one correspondence

Our commentary Blog step → our symbol. No sealed NoCap ΔT% here.

Full map: reviews/MUON_SOURCE_MAP.md (local-only). Split-QKV variants live in kit/muon_opt_split_qkv_cheap_v1.py (sealed) and related kit/muon_opt_split_qkv_*.py files; they call the same zeropower_via_newtonschulz5 — that is a separate campaign mechanism, not the NS kernel.

Our commentary These five Newton–Schulz steps are Keller Jordan’s Ortho kernel (blog newtonschulz5 / our zeropower_via_newtonschulz5). Any sealed wall-clock claim in this package is for the full candidate stack (Muon + timing-cheap split-QKV) vs F1_seal on the report — not a Muon-alone speed percentage invented on this page.