Tight Complexity of Muon under Heavy-Tailed Noise

Abstract

A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized. The layerwise input-output structure of neural networks motivates Muon and Scion, whose updates also support hyperparameter transfer. At the same time, stochastic gradient noise in deep learning is often far from sub-Gaussian and may exhibit heavy tails. These observations have shaped recent algorithmic principles for training neural networks, yet their joint theoretical consequences are underexplored. In particular, it is unclear what dimension dependence is unavoidable for gradient-based methods when the problem class is defined using an input-output norm and the stochastic gradient noise is heavy-tailed, and if higher-order smoothness can accelerate optimization. We study these questions through nonconvex smooth stochastic optimization over ${\mathbb{R}}^{m\times n}$ equipped with general norms and under $p^{\mathrm{th}}$-moment heavy-tailed noise, where the goal is to achieve an $\epsilon$-stationary point in the dual norm. Our first contribution is a dimension-dependent lower bound: when $\max \{m,n\}$ is large enough, any method using only the stochastic gradients requires $\Omega (\min \{m,n\}{\epsilon }^{-\frac{3p-2}{p-1}})$ oracle calls for the problem class defined by the spectral norm, which is a common input-output matrix norm. We prove that a batched Muon method with the spectral norm can achieve the matching upper bound of $O(\min \{m,n\}{\epsilon }^{-\frac{3p-2}{p-1}})$. To exploit higher-order smoothness, we propose a transported Muon method and improve the bound to $O(\min \{m,n\}{\epsilon }^{-\frac{5p-3}{2p-2}})$ when the norm is spectral and the Hessian is Lipschitz. Finally, we incorporate practical heuristics into our transported method and evaluate it across multiple architectures and model sizes, showing its compatibility with neural network training.

arXiv comment46 pages, 3 figures