Muon优化器改进了矩阵参数更新方式,支持重尾噪声下的高效稳定训练。
Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition
- 基于近似极分解与随机化低秩计算,实现矩阵参数的几何感知更新
- 在重尾噪声下达到最优复杂度 $O(\varepsilon^{-(3α-2)/(α-1)})$ 的收敛性能
- 适用于深度学习中大规模矩阵参数优化,尤其适合高噪声场景
大多数一阶优化器将矩阵参数视为向量处理,忽略了神经网络隐层权重的内在几何结构。Muon通过沿动量矩阵的极因子进行更新,解决了这一不匹配问题,但其理论分析滞后于实际应用。实践中,它结合Nesterov动量,仅近似计算极因子,并使用可能具有重尾分布的随机梯度。本文建立了带有Nesterov动量和不精确极分解的Muon在非凸矩阵优化中的收敛理论,适用于重尾噪声环境。分析基于统一的不精确极分解框架,涵盖牛顿-舒尔茨等迭代近似方法,并量化误差在优化动态中的传播。我们证明,找到ε-驻点的最优迭代与样本复杂度为 $O \left(\varepsilon^{\frac{-(3α-2)}{(α-1)}} \right)$,其中 $α\in(1,2]$ 表示重尾指数。对于 $σ_1=0$ 的不精确极分解情形,还提供了无需先验已知 $α$ 的保证。进一步分析了随机化低秩极分解,其效率显著高于全空间方法,且与理论兼容。数值实验验证了所提近似与随机化变体的有效性。
原文摘要 · Abstract (English)
Most first-order optimizers treat matrix-valued parameters as vectors, ignoring the intrinsic geometry of hidden-layer weights in neural networks. Muon addresses this mismatch by updating along the polar factor of a momentum matrix, but its theoretical understanding has lagged behind practice. In particular, practical implementations incorporate Nesterov momentum, compute the polar factor only approximately, and operate with stochastic gradients that may be heavy-tailed. We close this gap by developing a convergence theory for Muon with Nesterov momentum and inexact polar decomposition in non-convex matrix optimization under heavy-tailed noise. Our analysis builds on a unified framework for inexact polar decomposition that captures practical iterative approximations such as Newton-Schulz and quantifies how their errors propagate through the optimization dynamics. Under this framework, we establish an optimal iteration and sample complexity of $O \left(\varepsilon^{\frac{-(3α-2)}{(α-1)}} \right)$ for finding an $\varepsilon$-stationary point, where $α\in(1,2]$ denotes the heavy-tail index. For the inexact-polar setting with $σ_1=0$, we also provide guarantees that do not require prior knowledge of $α$. We analyze a randomized low-rank polar decomposition that is substantially more efficient than full-space methods while remaining compatible with our theory. Numerical experiments further demonstrate the effectiveness of the proposed inexact and randomized variants.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。