Muon orthogonalization limits in noisy gradient optimization near optima

Muon Under Gradient Noise and the Limits of Orthogonalization Near Optima

Machine Learning

Summary

Training deep learning models involves updating many numbers called weights using gradients estimated from small data batches, which can be noisy. The paper studies a method called Muon that keeps some parts of these weights orthogonal to stabilize training. The authors find that near an optimal solution where noise is significant, Muon's orthogonalization behaves like a scaled version of simpler momentum methods but also introduces extra complicated noise effects. These extra effects never fully go away and raise a baseline level of loss during training. The work suggests simpler momentum methods might be better when noise is dominant during optimization.

What this means in practice

  • For deep learning engineers: Assess when to replace orthogonal weight updates with momentum methods during noise-dominated training to optimize convergence speed and stability.
  • For machine learning infrastructure teams: Optimize training loop implementations by selecting the right momentum or orthogonalization techniques based on noise conditions near convergence.

A theory result. No direct application yet.

Authors

Xiaohui Xie

Abstract

Muon replaces the momentum buffer of each weight matrix by its orthogonal polar factor. We ask what this orthogonalization does near an optimum, where minibatch noise dominates the gradient. Under a Gaussian noise model, the expected Muon update becomes a scaled gradient step, so a linear method with a suitably matched learning rate reproduces Muon's first-order mean response. The stochastic update is a different matter: after the response is matched, Muon retains a nonlinear residual that is uncorrelated with the input noise and contributes additional covariance. A Hermite expansion shows how momentum acts on this residual. Its higher-order components decorrelate faster than the linear component, so momentum suppresses the residual's accumulated covariance relative to the linear part, but never eliminates it. In a local quadratic surrogate that evaluates the residual on the stationary noise buffer, the residual adds stationary covariance and raises the stationary loss floor at every stable step size, while leaving the contraction dynamics unchanged. Simulations of the full nonlinear recursion on quadratics and measurements on frozen transformer gradients support each step of this picture. Together, the results make a theoretical case for replacing orthogonalization by response-matched momentum SGD once optimization becomes noise-dominated.
View PDFOpen arXiv