Muon在重尾非凸优化中的理论优势:最优样本复杂度

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

精选理由

Muon为何能训练Transformer?

AI 摘要

Muon和Scion等非欧几里得优化方法在训练Transformer时表现优异,但其理论优势一直未明确。本研究证明在重尾非凸场景(随机梯度p阶中心矩有界,p∈(1,2])下,非欧几里得方法在更强的平稳性度量下达到最优样本复杂度,而欧几里得方法有额外维度依赖。对于m×n矩阵,Muon在核范数下找到ε-稳定点仅需O(min{m,n}Δ1L/ε^2(σ/ε)^{p/(p-1)})个样本,可吸收重尾噪声而无额外维度开销。实验在大型语言模型上验证了理论,并表明其他Schatten几何在某些设置下也可与Muon竞争。

原文 · arXiv cs.LG

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

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.

Muon在重尾非凸优化中的理论优势:最优样本复杂度 · AI 热点