Last-Iterate Stationarity of Muon with Finitely Many Newton-Schulz Iterations
Abstract Overview
This paper analyzes the last-iterate convergence behavior of practical Muon, where the update direction is produced by a fixed number of quintic Newton–Schulz iterations rather than an exact polar-factor computation. The setting is smooth, lower-bounded nonconvex optimization, and the analysis distinguishes deterministic gradients from stochastic gradients with persistent noise. The authors first show that finite-step Muon without input stabilization can cycle even on a smooth convex quadratic, so last-iterate stationarity is not automatic. They then prove that positive stabilization enables last-iterate gradient convergence with a sufficiently small constant step size in the deterministic case, while the stochastic case requires decreasing learning rates together with increasing momentum. The paper also shows that general pointwise polynomial rates need not hold, but density-one rates can still be established for almost all late iterations.
Novelty
The main novelty is a last-iterate stationarity theory for practical finite-step Muon, rather than average-iterate or best-iterate guarantees that dominate prior analyses. The work is also distinctive in identifying two concrete failure modes—cycling without stabilization and a persistent-noise floor under fixed stochastic hyperparameters—and in establishing density-one convergence rates when pointwise polynomial rates fail.
Results
For exact gradients, the paper proves that positive input stabilization and a sufficiently small constant learning rate yield square-summable gradients and convergence of the full gradient sequence to zero for any fixed momentum coefficient below one. For stochastic gradients, it shows that fixed learning rates and momentum can fail under bounded unbiased noise, but almost-sure and mean-square last-iterate convergence are recovered using decreasing step sizes and increasing momentum. It further proves density-one gradient decay rates, while giving a smooth counterexample showing that no general pointwise polynomial rate is guaranteed.
Key Points
- Without input stabilization, practical finite-step Muon can exhibit a two-cycle even on a smooth convex quadratic, so last-iterate convergence can fail in a simple deterministic setting.
- With positive stabilization, the authors derive a coupled descent-tracking Lyapunov analysis that yields full-sequence gradient convergence for deterministic smooth nonconvex objectives under a small constant step size.
- In the stochastic setting, persistent unbiased bounded-variance noise prevents convergence with fixed hyperparameters, but coordinated decreasing learning rates and increasing momentum restore almost-sure and mean-square last-iterate stationarity and support density-one rates.