Gram Newton-Schulz: A Fast, Hardware-aware Muon Optimizer

Abstract

Muon is quickly becoming the optimizer of choice for training deep neural networks such as large language models. Compared to the popular optimizer AdamW, Muon requires fewer optimizer steps to reach the same loss, but each optimizer step takes longer to compute. On net, it’s preferable, since it reaches the same loss in a shorter amount of wall clock time, but its expensive per-step computation is still a significant drawback. The expensive optimizer step runtime comes from the polar decomposition routine in Muon, which uses the Newton-Schulz method to orthogonalize 2D weight gradient updates, stepping in the steepest direction with respect to the spectral norm. The Newton-Schulz method approximates a polynomial that is typically computed as a series of 15 matrix ultiplications and additions (GEMMs), which are not present in AdamW. Specifically, Newton-Schulz’s runtime has two key shortcomings: 1) it has a strong dependence on the gradient’s longer dimension, which is detrimental as architectures become increasingly rectangular and 2) it computationally does not exploit the fact that many of its intermediate matrices are symmetric. We introduce a mathematically equivalent, hardware-aware algorithmic reimplementation of Newton-Schulz: Gram Newton-Schulz. Instead of iterating on the rectangular input matrix X ∈ Rn×m, we iterate on the small, square, symmetric Gram matrix XXT ∈ Rn×n. Gram Newton- Schulz addresses both of standard Newton-Schulz’s problems: 1) it rewrites the iteration such that a greater share of GEMMs are symmetric and square along the smaller dimension, reducing the runtime dependence on the gradient’s longer dimension, and 2) it uses fast symmetric GEMM kernels for the Hopper and Blackwell GPU architectures to exploit the greater share of symmetric matrices. However, naively iterating on the Gram matrix can introduce numerical instability - spurious negative eigen-values in XXT introduced by half-precision arithmetic can cause future matrices to diverge. We stabilize Gram Newton-Schulz with a simple restarting strategy - we first iii run a few iterations, then recompute XXT to reset the growing negative eigenvalues, before running the remaining iterations. We demonstrate that Gram Newton-Schulz matches standard Newton-Schulz’s training quality and derive the optimal restart strategy. Ultimately, we provide:

A drop-in replacement for Muon’s Newton-Schulz routine that is mathematically equivalent, numerically stable, and up to 2× faster, via https://github.com/Dao-AILab/gram-newton-schulz.

Fast GPU kernels for symmetric matrix multiplication (AB, αAB +βC) written in CuTeDSL for Hopper and Blackwell, via https://github.com/Dao-AILab/quack/blob/main/quack/gemm_symmetric.py.