Muon Converges under Heavy-Tailed Noise: Nonconvex H\"{o}lder-Smooth Empirical Risk Minimization

TL;DR

This paper demonstrates that Muon, an optimizer enforcing orthogonality via Stiefel manifold projections, effectively converges under heavy-tailed noise in nonconvex H"{o}lder-smooth empirical risk minimization, outperforming mini-batch SGD.

Contribution

It proves convergence of Muon in heavy-tailed noise settings for nonconvex H"{o}lder-smooth problems, a novel analysis for this optimizer.

Findings

Muon converges to a stationary point under heavy-tailed noise.

Muon converges faster than mini-batch SGD.

Theoretical guarantees extend to practical deep learning scenarios.

Abstract

Muon is a recently proposed optimizer that enforces orthogonality in parameter updates by projecting gradients onto the Stiefel manifold, leading to stable and efficient training in large-scale deep neural networks. Meanwhile, the previously reported results indicated that stochastic noise in practical machine learning may exhibit heavy-tailed behavior, violating the bounded-variance assumption. In this paper, we consider the problem of minimizing a nonconvex H\"{o}lder-smooth empirical risk that works well with the heavy-tailed stochastic noise. We then show that Muon converges to a stationary point of the empirical risk under the boundedness condition accounting for heavy-tailed stochastic noise. In addition, we show that Muon converges faster than mini-batch SGD.