Muon

Momentum-based optimizers take a step in whatever direction the gradient points, and for a weight matrix that direction is usually dominated by a few strong components. Muon flattens the update before applying it, so every direction gets a comparable share of the step. The matrix moves in a more balanced way than the raw gradient asks for.

Take the momentum buffer for a 2-D parameter and orthogonalize it before the update, approximating the nearest semi-orthogonal matrix with a handful of Newton–Schulz iterations rather than an SVD. Five steps is typical and runs in bf16. It applies only to matrix-shaped parameters — embeddings, biases, and norms stay on AdamW — so it is a supplement to an existing optimizer, not a wholesale replacement.

For momentum buffer M with SVD M = UΣVᵀ, the orthogonalized update is O = UVᵀ, discarding Σ entirely. Newton–Schulz approximates this by iterating X ← aX + b(XXᵀ)X + c(XXᵀ)²X on a spectrally normalized X₀ = M/‖M‖ sub _F , with coefficients chosen to push singular values toward 1. The iteration need not converge tightly: the update direction is what matters, and a rough orthogonalization already equalizes the spectrum enough to change the step.