Jason Trinh

Interactive Visualizations

Orthogonalization for Muon Optimization

In modern deep learning, optimizers like Muon achieve state-of-the-art performance by orthogonalizing gradient updates during model training. Mathematically, this orthogonalization step normalizes the update geometry, which corresponds to forcing all singular values of the gradient matrix to converge to exactly 1.

The Solution: Computing a full SVD at every training step is too computationally expensive. Instead, we use iterative algorithms like Polar Express. These methods rely purely on matrix-matrix multiplications—which are highly efficient on GPUs—to rapidly "flatten" the spectrum, pushing all singular values toward 1 without explicit decomposition.

Open Polar-Express / Muon Visualizer ↗

The Algorithms Visualized

Polar Express

A novel, highly-optimized iterative method designed for modern hardware. It uses specific polynomial scaling coefficients to accelerate the convergence of singular values to 1 far faster than standard methods.

Key Trait: Uses "cushioning" and "safety" parameters to aggressively converge without stability loss.

Newton-Schulz

The classic iteration for computing the Polar Decomposition. It converges quadratically to orthogonal form but is only locally convergent (requires singular values close to 1 initially).

Xk+1 = 0.5 * Xk(3I - XkTXk)

Jordan (Matrix Sign)

Based on the Matrix Sign Function. It is globally convergent and robustly forces singular values to 1. However, it typically requires matrix inversion, making it computationally expensive for large models.

Xk+1 = 0.5 * (Xk + Xk-T)