Momentum method with few Newton iterations converges for nonconvex optimization
Convergence of Practical Muon with Finite Newton-Schulz Iterations and Nesterov Momentum
Machine Learning
Summary
When optimizing complex models with many layers, maintaining speed and accuracy is tough. This paper studies a method called Practical Muon that uses a few special update steps for each layer, combined with a momentum strategy, to improve training. The authors prove that this approach converges to a solution at a predictable rate, without needing assumptions that are often used but hard to meet in practice. Their analysis holds under broad conditions and includes how the momentum error behaves over time.
What this means in practice
- •For machine learning engineers: Use Practical Muon updates with Nesterov momentum for layer-wise optimization schemes that guarantee convergence under realistic noisy gradient scenarios.
- •For numerical optimization developers: Incorporate finite-step Newton-Schulz iterations with momentum into matrix parameter optimization routines to improve stability and convergence guarantees in nonconvex settings.
A theory result. No direct application yet.
Authors
Hanyng Peng, Hui Wang, Yue Yu
Abstract
Practical Muon maintains momentum and performs a small, fixed number of Newton--Schulz iterations separately for each parameter matrix, often with a Nesterov correction. We analyze these layer-wise finite-step updates jointly on a coupled nonconvex objective, rather than replacing them by exact polar factors or one global orthogonalization. Under gradient-dependent $(\mathcal L_0,\mathcal L_1,q)$-smoothness and conditionally unbiased stochastic gradients with bounded layer-wise variance, we establish an $\mathcal O(T^{-1/4})$ bound on the expected average Frobenius gradient norm. The analysis retains the Nesterov recursion and requires neither bounded stochastic gradients, symmetric noise, nor a uniform positive lower bound on the nonzero output singular values. Its constants contain no explicit matrix-dimension or rank factors when the number of blocks and problem constants are fixed. The proof follows a descent inequality and a decomposition of the momentum tracking error into initialization, noise, and drift. For the original five-step quintic, we verify the required scalar-map bounds analytically; the result also allows step-dependent coefficients satisfying the same bounds. A complementary nuclear-norm result quantifies rank dependence under a stronger spectral condition. The vanishing rate uses coupled learning-rate and momentum schedules, including the standard single-coefficient Nesterov rule.