The optimizer that flattens every gradient, read closely

Imagine steering a ship with sixteen engines. One engine is roaring; the other fifteen barely turn. A usual optimizer scales individual coordinates, but it does not first ask whether a whole weight matrix is trying to move along only one dominant direction. Muon does.

It treats a neural-network weight as a matrix, takes the gradient plus momentum, and replaces its singular values with values near one. Singular values are the strengths of the matrix's independent directions. Flattening them means the update can move every available direction instead of spending most of its budget on the loudest one.

Suppose a 2×2 gradient has strengths 8 and 0.5 along two perpendicular directions: a 16:1 imbalance. If we write the gradient's singular value decomposition as G = UΣVᵀ, an exact polar-factor update is UVᵀ. The 8 and 0.5 both become 1; orientation survives, magnitude does not.

Computing a full decomposition for every layer would be expensive. The paper instead uses five quintic Newton–Schulz steps. Each step is only matrix multiplication, which accelerators already do exceptionally well. The TypeScript below is what draws the figures; Python and C++ are faithful translations.

def muon_direction(g, steps=5):
    x = copy_matrix(g)
    tall = len(x) > len(x[0])
    if tall:
        x = transpose(x)
    scale = sqrt(sum(v * v for row in x for v in row))
    x = scale_matrix(x, 1.0 / (scale + 1e-7))
    for _ in range(steps):
        a = matmul(x, transpose(x))
        b = add(scale_matrix(a, -4.7750), scale_matrix(matmul(a, a), 2.0315))
        x = add(scale_matrix(x, 3.4445), matmul(b, x))
    if tall:
        x = transpose(x)
    return x

Here is the part a skim misses. If the orthogonalized update is an m × n full-rank matrix, its root-mean-square entry is 1 / √max(m,n). RMS—root mean square—is a typical entry size. Make one dimension four times wider and the raw Muon update becomes half as large, even when the layer deserves the same learning rate.

That matters because Transformer matrices have very different shapes. A square attention projection and a wide feed-forward projection would receive different effective learning rates for geometric reasons alone. The paper found wide MLP updates too small and some narrow attention updates too large.

The direct fix is almost embarrassingly small: measure the update RMS and rescale it. The paper also tests an equivalent shape-aware learning-rate adjustment and chooses that version because it is cheaper at scale.

def rescale_update_rms(update, target=0.2):
    n = sum(len(row) for row in update)
    rms = sqrt(sum(v * v for row in update for v in row) / n)
    if rms == 0:
        return copy_matrix(update)
    return scale_matrix(update, target / rms)

The original small-scale Muon recipe omitted weight decay. Weight decay gently shrinks parameters at every step, independently of the loss gradient. At larger scale and longer training, the paper observed weight and layer-output RMS growing until values pressed against the useful precision range of BF16, the 16-bit format used for much of training. Vanilla Muon started faster, then lost its advantage. Adding decoupled weight decay restored it.

At 1.2T tokens, the authors compare the same 15.29B-total-parameter MoE architecture trained with AdamW and Muon. Muon wins clearly on HumanEval (37.2 versus 29.3), MBPP (52.9 versus 49.2), GSM8K (45.0 versus 43.8), and MATH (19.8 versus 16.1), while BBH moves the other way (43.2 versus 45.3). Those are the paper's reported scores. The mixed benchmark pattern is healthier evidence than a story in which every number magically rises.

The scaling-law comparison is the cleaner optimizer result: both optimizers receive a serious compute-optimal setup, and Muon reuses the AdamW-tuned learning rate only after its update RMS is matched. Still, the training data are proprietary, the report is an arXiv technical report rather than a peer-reviewed benchmark, and the tiny reconstruction here demonstrates the mechanism—not the 52% FLOP claim.

• Repeat the scaling-law comparison on a fully public token mixture so optimizer and data effects can be separated.
• Test whether the higher singular-spectrum entropy predicts better downstream performance layer by layer, rather than merely correlating with it.
• Map where the pretraining-to-fine-tuning optimizer mismatch begins: checkpoint age, learning-rate ratio, model width, or loss change.

Want the geometric intuition first? Use the linear model to watch one gradient direction move a fit, then poke at matrix-shaped layers on the neural-network page and attention projections on the Transformer page.