Muon optimizer updates improve training by orthogonalizing directions
A Quadratic Lens on Muon's Polar Update: Benefits, Pitfalls, and the Role of Newton-Schulz
Machine Learning
Summary
Training large language models involves adjusting many parameters to make the model better. The authors studied a method called Muon's Polar updates, which change parameters in special directions that are nearly orthogonal or independent. They found that this method helps when the data has certain complex structures and can speed up training in some cases. However, when some parts of the data are very small or weak, this method can be slower than normal approaches. The authors also explain how a mathematical technique called Newton-Schulz helps improve speed by handling these small parts differently.
What this means in practice
- •For machine learning engineers: Optimize training of transformer models by using PolarGD updates to speed up convergence in regimes with complex parameter structures.
- •For numerical optimization developers: Improve matrix-valued optimization algorithms by integrating orthogonalized updates and Newton-Schulz iterations to handle poorly conditioned cases.
Authors
Egor Shulgin, Sam Laing, Antonio Orvieto, Peter Richtárik
Abstract
Muon updates matrix-valued parameters along approximately orthogonalized directions and often trains large language models faster than AdamW, yet it remains unclear when and why orthogonalization helps. We study PolarGD, the orthogonalized update at the core of Muon, on matrix quadratics, where optimizers can be analyzed in closed form, and compare update directions under line search to remove the effect of learning-rate tuning. Orthogonalization pays off when the gradient is far from low rank and its largest singular directions see high curvature. PolarGD reaches the minimizer in one step whenever the rows of the error are orthogonal with equal norms, for any Hessian, whereas gradient descent needs all Hessian eigenvalues to be equal; and unlike sign descent, PolarGD is unaffected by rotations of the problem. Over many steps, PolarGD acts as gradient descent preconditioned by the current iterate, so its speed is governed by the conditioning of the iterate rather than of the Hessian. This explains fast regimes as well as a slow one, caused by tiny singular values of the iterate, in which PolarGD is much slower than gradient descent. Finally, the few Newton-Schulz iterations used in Muon treat these tiny singular values as gradient descent does, so fewer iterations can converge faster than exact orthogonalization; in a transformer trained with Muon, about a quarter of the singular values of the attention matrices fall in this regime.