Polar Lab

running

◷train_gpt.py▸polar_express◎learner＄no GPU

Blocks

Quick summary

5-iteration polynomial approximation of the polar decomposition. Input gradient G → orthogonal matrix U (from G =…

Full explanation below the code →

fba-lab — train_gpt.py · polar_expressexecuting

// block: polar_express · lines 151–174$ study train_gpt.py --block polar-express-fn5-iteration polynomial approximation of the polar decomposition. Input gradient G → orthogonal matri… ✓

Explanation

5-iteration polynomial approximation of the polar decomposition. Input gradient G → orthogonal matrix U (from G = U·Σ·V^T). dynamic=False required — Triton recompiles on shape change, killing speed.

Think about

What does it mean for the update matrix to be orthogonal, and why does Muon need this?

// architecture

Live diagram

02 · train_gpt.py line 151

polar_express

Turning gradient matrices into orthogonal update directions

The NorMuon optimizer doesn't apply raw gradients. Before updating weights, it orthogonalizes the gradient — turns it into a matrix where every row is perpendicular to every other row. This is called the polar decomposition. polar_express computes it fast using 5 iterations of a polynomial approximation.

@torch.compile(dynamic=False, fullgraph=True) # dynamic=False required — re-tracing kills speed
def polar_express(G: torch.Tensor, split_baddbmm: bool = False):
X = G.bfloat16() # run in bf16 for speed
if G.size(-2) > G.size(-1): X = X.mT # algorithm needs tall/square matrix
X = X / (X.norm(dim=(-2,-1), keepdim=True) * 1.02 + 1e-6) # normalize spectral norm ≤ 1
for a, b, c in polar_express_coeffs:
A = XXT(X)
B = b*A + c*(A @ A)
X, C = a*X + B@X, X
return X # ≈ orthogonal: U from polar decomposition G = U·Σ·V^T

Where it fits in the optimizer

# NorMuonAndAdam._normuon_update()
grad = grad.float()
momentum_buf.lerp_(grad, 1 - momentum)
v = polar_express(momentum_buf) # ← orthogonalize
v = _apply_normuon_variance_reduction(v)
p_slice.add_(v, alpha=-lr)

Why orthogonalize at all?

Raw gradient updates can couple different feature directions — one dimension's update interferes with another's. Orthogonalizing the update matrix ensures each direction is updated independently. This is the core insight from Muon (MomentUm Orthogonalized by Newton-Schulz). polar_express uses precomputed polynomial coefficients for speed.

Turning gradient matrices into orthogonal update directions

Where it fits in the optimizer

Why orthogonalize at all?