Five matrix multiplies that turn a gradient into a rotation

01What we are actually computing

Muon is SGD with momentum, plus one extra operation. For a 2-D weight matrix W:

If the momentum matrix has the singular value decomposition M = UΣV^⊤, then

That is it. Every singular value gets set to 1; the singular directions are untouched. This is the orthogonal factor of the polar decomposition M = QP, and it is also the closest semi-orthogonal matrix to M in Frobenius norm.

Why you would want that

Two ways to see it. The geometric one: a real momentum matrix has singular values spread over orders of magnitude, so a raw SGD step moves the network almost entirely along two or three dominant directions and barely touches the rest. Orthogonalizing gives every direction the same step size.

The clean one: UV^⊤ is exactly the steepest-descent direction when you measure step size in the spectral norm. Among all updates with ‖ΔW‖₂ ≤ 1, the one most aligned with the gradient is UV^⊤:

Drag the slider below to morph a 2×2 matrix into its orthogonalized version. The faint dots are the input unit circle; the bright dots are their images. Orthogonalizing squashes the ellipse back into a circle without spinning the dots around it.

02Where the name comes from

Newton's method can be aimed at the polar factor directly. The iteration is

Inverting a matrix every step, in low precision, for every weight matrix in the network, is a non-starter. This is where Schulz enters. Note that

and that near the answer X^⊤X is close to the identity, so its inverse can be approximated by the first two terms of a Neumann series — the trick Schulz published in 1933 for inverting matrices iteratively:

Substitute that into Newton's step and the inverse vanishes:

Newton gives the iteration; Schulz makes it inverse-free. That is the whole naming story, and it is also the whole reason it belongs in a training loop: matmuls are the one thing a GPU is spectacularly good at.

03The one fact that makes it all work

Here is the piece that usually clicks last, and it is the piece worth internalizing. Write X = UΣV^⊤ and expand the cubic term, using V^⊤V = I and U^⊤U = I:

So one Newton–Schulz step gives

This holds for any odd polynomial built from XX^⊤, because (XX^⊤)^jX = UΣ^2j+1V^⊤. That freedom is exactly what Muon exploits in §5.

04So — does σ go to 1?

The scalar map is

Its fixed points solve f(s) = s, i.e. s(1 − s²) = 0, so s ∈ {0, 1} for positive values. On (0, 1) we have f(s) > s and f(s) < 1, so the sequence climbs monotonically to 1. Near the top, set s = 1 − ε:

The basin is (0, √3). Past √3 the map goes negative and blows up. Which explains a line in the code that otherwise looks arbitrary:

Reading the cobweb

A cobweb diagram makes a scalar iteration visible. From the current value on the horizontal axis, go vertically to the curve (that applies the map), then horizontally to the diagonal (that feeds the output back in as the next input). Repeat. The staircase is the iteration.

Set σ₀ near 1 and the classic cubic snaps home in two or three steps. Now drag it down toward 0.05 and watch the staircase turn into a crawl. That is the whole problem.

05Why Muon changes the numbers

Near zero, f(s) ≈ 1.5s. Small singular values grow by a factor of 1.5 per step, and nothing faster. Real momentum matrices are badly conditioned; after Frobenius normalization the smallest singular value can easily sit around 10⁻³. Getting from there to 1 at 1.5× per step takes on the order of seventeen iterations — every optimizer step, on every 2-D parameter. Far too expensive.

Muon's answer is to give up on exactness. Use a quintic, and hand-tune the coefficients for speed rather than for correctness:

The slope at the origin is now 3.4445 instead of 1.5 — small singular values are hauled upward more than twice as fast per step. The price is real, and worth stating plainly:

Flip the toggle in Fig 2 back to quintic and set σ₀ near 1: you can watch the staircase refuse to settle. Here is the same starting spectrum run through both maps, on a log axis, so that constant multiplicative growth appears as constant vertical movement:

The shaded band is 0.7 to 1.3 — close enough to orthogonal for Muon's purposes. After five steps the classic iteration has landed its top handful of values and left the tail in the basement. The quintic has everything in the band.

Concretely, here is how many iterations each map needs to drag a singular value up to 0.7:

06The code, annotated

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

1 — bfloat16. The output is deliberately approximate to within ±30% anyway, so roughly three significant digits is plenty. Halves the memory traffic and lands the whole thing on tensor cores.

2 — Frobenius normalization. The basin guard from §4. Since σ_max ≤ ‖G‖_F, this puts every singular value inside (0, 1]. The eps covers a zero gradient.

3 — The transpose. The iteration only ever forms XX^⊤, so you want the short side on the outside. Flipping a tall matrix to wide means the Gram matrix is m × m with m the smaller dimension. For a 4096×1024 weight that is a 1024² product instead of a 4096² one.

4 — The factoring. Written naively, aX + bXX^⊤X + c(XX^⊤)²X costs four matmuls. Setting A = XX^⊤ and pulling out X on the right gives aX + (bA + cA²)X — three. Combined with the transpose, that took the cost from 6nm² to 4nm² + 2m³ FLOPs.

Fifteen matmuls total on a matrix the size of one weight tensor. Against a forward and backward pass over a full batch, it disappears — reported overhead is under 1% of training time.

07Watch it run on a real matrix

Below is an actual 12×12 matrix constructed with singular values log-spaced from 1 down to 0.01, then Frobenius-normalized — a plausible stand-in for a badly conditioned momentum buffer. The left panel is X itself; the right is XX^⊤, which converges to the identity exactly when X becomes orthogonal. Step through with each polynomial and compare.

Watch the diagonal of the right panel light up. With the cubic, five steps illuminate only the corner corresponding to the large singular values, and the rest of the diagonal stays dark — those directions are still being suppressed. With the quintic the whole diagonal comes up together, unevenly, hovering around 1 rather than pinned at it. That unevenness is the ±30% band, and it is what Muon ships.

08Fine print

• 2-D parameters only. Muon is for hidden-layer weight matrices. Embeddings, the output head, biases and normalization scales are conventionally left on AdamW — "orthogonalize the singular values" is meaningless for a vector.
• The coefficients are not sacred. They were tuned, not derived, and there is active work on better polynomial schedules — including per-iteration coefficients and Chebyshev-style constructions that hit the same band in fewer steps.
• Production variants add more. Nesterov momentum, weight decay, and an update-scale factor depending on the matrix aspect ratio are common. The Newton–Schulz core is unchanged.
• Convergence is not the goal. The single most counterintuitive thing here: a numerical method deliberately configured not to converge, because the application only needs the right order of magnitude and would rather have the compute back.

Primary source: Keller Jordan's Muon: An optimizer for hidden layers in neural networks (December 2024), and Bernstein & Newhouse, Old Optimizer, New Norm: An Anthology, which is where the spectral-norm framing comes from. Both go further than this page does.