Free Heavy-Tailed Lunch for Muon: A Theoretical Justification of Empirical Success

2026-06-12

AI Summary

This work addresses the challenge of heavy-tailed nonconvex stochastic optimization, where gradients possess only a finite $p$-th central moment for $p \in (1,2]$, rendering existing Euclidean methods ineffective in guaranteeing efficient convergence. The paper presents the first theoretical analysis of the non-Euclidean optimizer Muon under this setting, establishing its sample complexity in terms of the nuclear norm. It proves that Muon achieves the dimension-free optimal convergence rate of $O\left(\min\{m,n\} \cdot \Delta_1 L / \varepsilon^2 \cdot (\sigma / \varepsilon)^{p/(p-1)}\right)$, and shows this bound is tight for all first-order methods under nuclear norm stability. The analysis integrates the geometry of Schatten spaces with heavy-tailed optimization theory, and empirical validation on large language model training demonstrates Muon’s effectiveness while revealing the competitive potential of other Schatten geometries.

Abstract

Non-Euclidean optimisation methods with matrix-valued updates, such as Muon and Scion, have recently shown strong empirical performance for training Transformer models, yet their theoretical advantages over Euclidean methods remain poorly understood. We address this gap in the heavy-tailed non-convex regime, where stochastic gradients have bounded $p$-th central moments, $p \in (1,2]$. We show that certain non-Euclidean methods achieve optimal sample complexity under stronger stationarity measures, while Euclidean methods incur additional dimension-dependent costs. As a consequence, for $m \times n$ matrices, Muon finds an $\varepsilon$-stationary point in nuclear norm within $\mathcal{O}\left(\min\{m, n\} \frac{Δ_1 L}{\varepsilon^2} \left(\frac σ\varepsilon \right)^{\frac p {p-1}}\right)$ samples, absorbing heavy-tailed noise without extra dimension dependence, unlike Euclidean methods. We further prove this sample complexity, including its dimension dependence, is optimal for all first-order methods under nuclear-norm stationarity. Experiments on large language models support our theory. Surprisingly, our results suggest that other Schatten geometries beyond the spectral geometry of Muon can perform competitively in certain settings.

Problem

• non-Euclidean optimization
• heavy-tailed noise
• sample complexity
• dimension dependence
• non-convex optimization

Innovation

• non-Euclidean optimization
• heavy-tailed gradients
• sample complexity
• nuclear norm stationarity
• Schatten geometry