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.