解释了非欧优化方法为何在大模型训练中表现优异。
Free Heavy-Tailed Lunch for Muon: A Theoretical Justification of Empirical Success
- 在重尾噪声下，非欧方法比欧氏方法更高效。
- Muon找到ε-平稳点所需样本数与维度无关。
- 理论证明其性能最优，适合大模型优化研究者。
具有矩阵更新的非欧优化方法（如 Muon、Scion）在训练 Transformer 模型时表现出强劲的实证性能，但其相对于欧氏方法的理论优势尚不明确。本文在重尾非凸场景下填补这一空白，其中随机梯度的 p 阶中心矩有界，p ∈ (1,2]。我们证明某些非欧方法在更强的平稳性度量下可达到最优样本复杂度，而欧氏方法则引入额外的维度依赖成本。对于 m×n 矩阵，Muon 在核范数下找到 ε-平稳点仅需 𝒪(min{m,n} Δ₁L/ε² (σ/ε)^{p/(p-1)}) 次采样，能吸收重尾噪声且无维度依赖，而欧氏方法不具备此特性。进一步证明该复杂度（含维度依赖）在核范数平稳性下对所有一阶方法均为最优。大规模语言模型实验支持理论结果。令人意外的是，其他施瓦茨几何（如非谱几何）在特定场景下也可表现良好。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。