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

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 LLMs support our theory. Surprisingly, our results suggest that other Schatten geometries beyond the spectral geometry of Muon can perform competitively in certain settings.

• The paper provides a theoretical justification for Muon’s optimal performance using normalized steepest descent methods under heavy-tailed noise.
• It establishes dimension-sensitive sample complexity and lower bounds for nuclear-norm stationarity, outperforming classical Euclidean methods.
• The analysis confirms that choosing appropriate non-Euclidean norms improves optimization trajectories in large-scale Transformer models.

Theoretical Justification for Heavy-Tailed Optimization in Muon

Background and Motivation

The paper "Free Heavy-Tailed Lunch for Muon: A Theoretical Justification of Empirical Success" (2606.14560) addresses a key open question in non-Euclidean optimization: whether matrix-valued methods such as Muon can achieve provable theoretical superiority over classical Euclidean methods when training large-scale Transformer models, especially in stochastic regimes featuring heavy-tailed gradient noise. Historically, adaptive optimizers like Adam have dominated Transformer training, but alternatives based on steepest descent in general Banach spaces (such as Muon, using spectral norms) have demonstrated strong empirical improvements. The theoretical underpinnings for these non-Euclidean approaches remained underexplored, particularly regarding sample complexity and stationarity measures under realistic, heavy-tailed noise distributions typical in modern deep learning.

Main Theoretical Contributions

The paper establishes rigorous results showing dimension- and geometry-sensitive convergence guarantees for normalized steepest descent (NSD) methods with momentum, including Muon, under heavy-tailed stochastic noise. The technical innovation lies in stating assumptions and deriving bounds directly in non-Euclidean norms, thereby avoiding the dimension-dependent norm equivalence reduction to Euclidean geometry typically found in previous analyses.

For optimization in $m \times n$ matrices, Muon achieves nuclear-norm stationarity with sample complexity

$$\mathcal{O}\left(\min\{m, n\} \frac{\Delta_1 L}{\epsilon^2} \left(\frac{\sigma}{\epsilon}\right)^{\frac{p}{p-1}}\right)$$

absorbing heavy-tailed gradient noise without incurring extra dimension-dependent penalties. In sharp contrast, Euclidean algorithms, when stationarity is properly adjusted, incur a strictly worse $p$-dependent dimension factor.

Figure 1: Mean final perplexity and standard deviation for a Schatten-$r$ sweep, marking optimal learning rates and overall best results.

Additionally, the paper delivers a first-order sample complexity lower bound for nuclear-norm stationarity matching the Muon upper-bound structure, formally establishing Muon's optimality in this regime. The results extend to arbitrary Banach spaces by leveraging martingale deviation bounds for uniformly convex/smooth norms and exploiting a functional-analytic duality between convexity and smoothness constants.

Analysis of Heavy-Tailed Noise in Practice

Experiments focus on pre-training a 70M-parameter Pythia architecture on the FineWeb-Edu dataset, examining the heavy-tailed nature of stochastic gradients across weight matrices. The ratio $\sigma_{S_1,p}/\sigma_{F,p}$ (with $p=1.5$) is computed to quantify how much dimension dependence arises when the noise assumption is stated in Muon's native geometry versus the Euclidean geometry. Empirical ratios remain well below the worst-case theoretical bound $\sqrt{\min\{m,n\}}$, particularly at initialization, validating the practical severity of dimension dependence as fundamentally suboptimal for Euclidean approaches.

Figure 2: Empirical ratios $\sigma_{S_1,p}/\sigma_{F,p}$ at initialization across all Muon-trained weight matrices, stratified by layer family and depth.

Figure 3: The same ratios measured after training, showing increased values yet still far below theoretical maxima, explaining Muon's improved early-stage performance.

Additionally, the study reveals norm choice impacts optimization trajectory. A systematic sweep across Schatten-$r$ geometries shows perplexity improvement as $r$ increases from 1 up to about 8, plateauing thereafter with spectral geometry ($$\mathcal{O}\left(\min\{m, n\} \frac{\Delta_1 L}{\epsilon^2} \left(\frac{\sigma}{\epsilon}\right)^{\frac{p}{p-1}}\right)$$0) delivering competitive results without incurring the computational expense of full SVDs.

Figure 4: Plain noise values $$\mathcal{O}\left(\min\{m, n\} \frac{\Delta_1 L}{\epsilon^2} \left(\frac{\sigma}{\epsilon}\right)^{\frac{p}{p-1}}\right)$$1 at various checkpoints, illustrating geometry-aware noise magnitude.

These findings empirically reinforce that Muon's spectral geometry is preferable for matrix weights, while embedding layers align better with $$\mathcal{O}\left(\min\{m, n\} \frac{\Delta_1 L}{\epsilon^2} \left(\frac{\sigma}{\epsilon}\right)^{\frac{p}{p-1}}\right)$$2-type geometry, potentially explaining hybrid optimizer use in practice.

Dimension Dependence: Lower Bounds and Optimality

The paper also constructs theoretical lower bounds for nuclear norm stationarity, employing the coupon-collector argument and high-dimensional embeddings. The result rigorously proves that the remaining $$\mathcal{O}\left(\min\{m, n\} \frac{\Delta_1 L}{\epsilon^2} \left(\frac{\sigma}{\epsilon}\right)^{\frac{p}{p-1}}\right)$$3 dependence in Muon's sample complexity cannot be improved by any first-order algorithm, ruling out the possibility of dimension-free guarantees for nuclear norm stationarity. Thus, Muon is first-order optimal with respect to heavy-tailed sample complexity and geometry-aware stationarity.

Implications and Future Directions

Theoretical results decisively demonstrate that geometry-aware optimization, instantiated by Muon and its spectral norm, provides significant and provable reductions in sample complexity under realistic, heavy-tailed stochastic noise, relative to classical Euclidean approaches. This elevates Muon's methodology from empirical success to theoretical optimality for large-scale, matrix-valued deep architectures.

Practically, the findings encourage broader adoption of Muon-type algorithms in LLM and deep sequence modeling, especially where heavy-tailed gradient phenomena are pronounced. The framework also motivates future work on generalizing high-probability bounds, extending guarantees to full deep networks beyond idealized updates, and exploring geometry-adaptive methods for hybrid layer structures.

From a theoretical standpoint, the paper's functional analysis toolkit enables sharper analysis of optimization dynamics in arbitrary Banach spaces, opening avenues for robust, geometry-aware algorithm design beyond deep learning. Further study into high-dimensional lower bounds, probabilistic martingale-type behavior, and optimal stationarity conditions in stochastic regimes is warranted.

Conclusion

The paper rigorously closes the gap between empirical and theoretical understanding for Muon and related non-Euclidean optimization methods in transformer training. By establishing dimension-optimal sample complexity bounds under heavy-tailed noise and validating these claims both theoretically and empirically, the authors provide a formal foundation for Muon's methodology and its superiority over classical approaches. The implications for large-scale AI training, geometry-aware optimization theory, and practical algorithm design are immediate, informing both theoretical research and applied machine learning practice.