Keywords: Muon, Newton–Schulz, Orthogonalization, Nonconvex Optimization
TL;DR: Muon with Newton–Schulz as used in practice converges at the SVD-polar rate up to a constant that shrinks doubly-exponentially, which explains why using 2–3 steps is both fast and accurate and improves rank dependence over SGD with momentum.
Abstract: We analyze Muon as originally proposed and used in practice---using the momentum orthogonalization with a few Newton-Schulz steps. The prior theoretical results replace this key step in Muon with an exact SVD-based polar factor. We prove that Muon with Newton-Schulz converges to a stationary point with the same rate as the SVD-polar idealization, up to a constant factor for given the number of Newton-Schulz steps $q$. We further analyze this constant factor, and prove that it converges to 1 doubly exponentially in $q$ and improves with $\kappa$, which is the degree of a polynomial used in Newton-Schulz required when approximating the orthogonalization direction.
We also prove that Muon removes the typical square-root-of-rank loss compared to its vector-based counterpart, SGD with momentum. Our results explain why Muon with a few low-degree Newton-Schulz steps matches exact-polar (SVD) behavior at much faster wall-clock time, and explain how much momentum matrix orthogonalization via Newton-Schulz benefits over the vector-based optimizer. Overall, our theory justifies the practical Newton-Schulz design of Muon, narrowing its practice–theory gap.
Primary Area: optimization
Submission Number: 23375
Loading