#38 · Polar Express

running

◷train_gpt.py▸polar_express function◎learner＄no GPU

Blocks

Quick summary

Polar Express computes the orthogonal polar factor of G using Newton-Schulz iterations — without ever computing a full…

Full explanation below the code →

fba-lab — train_gpt.py · polar_express functionexecuting

// block: polar_express function · lines 376–400$ study train_gpt.py --block polar-fnPolar Express computes the orthogonal polar factor of G using Newton-Schulz iterations — without eve… ✓

Explanation

Polar Express computes the orthogonal polar factor of G using Newton-Schulz iterations — without ever computing a full SVD. It normalizes the spectral norm to ≤1 first, then iterates a cubically-convergent update until G·Gᵀ ≈ I.

Think about

Why is this cheaper than computing SVD for an optimizer step?

// architecture

Live diagram

Record #38 · PR #134 · Polar Express

Muon optimizer

Orthogonal gradient updates — without the expensive SVD

The Muon optimizer needs to compute the orthogonal polar factor of each gradient matrix — the rotation that strips away scaling, leaving only direction. The original approach called full Singular Value Decomposition (SVD). SVD is exact but slow: it scales as O(mn²) and runs poorly on GPU tensor cores because it can't be expressed as a simple batched matmul.

Polar Express replaces SVD with 5 Newton-Schulz iterations — a cubically-convergent loop that needs only matmuls. Each iteration is G ← 1.5G − 0.5·G·Gᵀ·G. Five steps is enough for the precision needed. On H100 tensor cores, five matmuls run roughly 10× faster than one SVD at these matrix sizes.

# lines 376–400 — Newton-Schulz orthogonalization (no SVD)
def polar_express(G, steps=5):
    G = G / (G.norm() + 1e-7) # normalise spectral norm
    for _ in range(steps):
        G = 1.5 * G - 0.5 * G @ G.T @ G
    return G

Why orthogonal updates?

Orthogonal matrices preserve vector norms — they rotate without stretching. Applying the orthogonal polar factor to a gradient keeps every parameter update the same effective magnitude, regardless of how large or small the gradient was. It is similar to sign-SGD but smoother and more principled.

Why does Newton-Schulz converge so fast?

The update G ← 1.5G − 0.5·G·Gᵀ·G is cubically convergent once G is close to orthogonal — the error cubes each step. After normalizing the spectral norm below 1 first, 5 iterations gives machine-precision orthogonality for the matrix shapes used here.

What changed vs #29

+     Polar Express Sign Method: https://arxiv.org/pdf/2505.16932
+     by Noah Amsel, David Persson, Christopher Musco, Robert M. Gower.
+     Code adapted from https://github.com/NoahAmsel/PolarExpress/tree/main by @varunneal.
+     """
+     X = X / (X.norm(dim=(-2, -1), keepdim=True) * (1 + 2e-2) + 1e-6)
+     aX_plus_BX = torch.baddbmm if X.ndim > 2 else torch.addmm
-     for _ in range(5):
-         ns_line_1(X, out=A)  # A = X @ X.mT
-         ns_line_2(A, alpha=c, beta=b, out=B)  # B = b * A + c * A @ A
-         ns_line_3(X, B, X, beta=a, out=C)  # C = a * X + B @ X
+     # Perform the iterations
+     for a, b, c in coeffs_list:
+         XXT(X, out=A)  # A = X @ X.mT
+         ba_plus_cAA(A, alpha=c, beta=b, out=B)  # B = b * A + c * A @ A
+         aX_plus_BX(X, B, X, beta=a, out=C)  # C = a * X + B @ X

Orthogonal gradient updates — without the expensive SVD

Why orthogonal updates?

Why does Newton-Schulz converge so fast?

Lab

#38 · Polar Express

Live diagram

Orthogonal gradient updates — without the expensive SVD

What changed vs #29

polar_express function

polar_express in Muon step

#38 · Polar Express

Live diagram

Orthogonal gradient updates — without the expensive SVD

What changed vs #29

polar_express function

polar_express in Muon step