Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition
Abstract
Most first-order optimizers treat matrix-valued parameters as vectors, ignoring the intrinsic geometry of hidden-layer weights in neural networks. Muon addresses this mismatch by updating along the polar factor of a momentum matrix, but its theoretical understanding has lagged behind practice. In particular, practical implementations incorporate Nesterov momentum, compute the polar factor only approximately, and operate with stochastic gradients that may be heavy-tailed. We close this gap by developing a convergence theory for Muon with Nesterov momentum and inexact polar decomposition in non-convex matrix optimization under heavy-tailed noise. Our analysis builds on a unified framework for inexact polar decomposition that captures practical iterative approximations such as Newton-Schulz and quantifies how their errors propagate through the optimization dynamics. Under uniform bounds on polar errors, we establish an optimal iteration and sample complexity of O (ε(−(3α−2))/((α−1)) ) for finding an ε-stationary point, where α∈(1,2]~denotes the heavy-tail index. For the inexact-polar setting with σ1=0, we also provide guarantees that do not require prior knowledge of α. We analyze a randomized low-rank polar decomposition that is substantially more efficient than full-space methods while remaining compatible with our theory. Numerical experiments further demonstrate the effectiveness of the proposed inexact and randomized~variants.
Related
- Continuized Nesterov Momentum Achieves the O(ε−7/4) Complexity in Smooth Nonconvex Optimization
- Momentum in large-batch training: Polyak enlarges the critical batch size, Nesterov improves data efficiency
- Analysis of an inexact Schwarz domain decomposition method in 3D
- Randomized inexact block triangular preconditioners for double saddle-point systems in PDE-constrained optimization
- Randomized Branch Methods with Inexact Subproblems for Bouligand Stationarity in Linear and Quadratic Programs with Complementarity Constraints
BibTeX
@misc{choudhury2026muon,
title = {{Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition}},
author = {Sayantan Choudhury and Xiaoran Cheng and Martin Takáč and Sen Na and Mladen Kolar},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2605.06884},
eprint = {2605.06884},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.06884},
}
RIS
TY - UNPB
TI - Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition
AU - Sayantan Choudhury
AU - Xiaoran Cheng
AU - Martin Takáč
AU - Sen Na
AU - Mladen Kolar
PY - 2026
DA - 2026/11/01
PB - arXiv
DO - 10.48550/arxiv.2605.06884
UR - https://arxiv.org/abs/2605.06884
AB - Most first-order optimizers treat matrix-valued parameters as vectors, ignoring the intrinsic geometry of hidden-layer weights in neural networks. Muon addresses this mismatch by updating along the polar factor of a momentum matrix, but its theoretical understanding has lagged behind practice. In particular, practical implementations incorporate Nesterov momentum, compute the polar factor only approximately, and operate with stochastic gradients that may be heavy-tailed. We close this gap by developing a convergence theory for Muon with Nesterov momentum and inexact polar decomposition in non-convex matrix optimization under heavy-tailed noise. Our analysis builds on a unified framework for inexact polar decomposition that captures practical iterative approximations such as Newton-Schulz and quantifies how their errors propagate through the optimization dynamics. Under uniform bounds on polar errors, we establish an optimal iteration and sample complexity of O (ε^(−(3α−2))/((α−1)) ) for finding an ε-stationary point, where α∈(1,2]~denotes the heavy-tail index. For the inexact-polar setting with σ_1=0, we also provide guarantees that do not require prior knowledge of α. We analyze a randomized low-rank polar decomposition that is substantially more efficient than full-space methods while remaining compatible with our theory. Numerical experiments further demonstrate the effectiveness of the proposed inexact and randomized~variants.
ER -