Mousse:通过曲率感知预条件化校正μon的几何
Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning
AI总结:
本文提出Mousse优化器,通过曲率感知预条件化改进谱方法的结构稳定性,实验证明其在大规模语言模型训练中显著提升效率。
摘要:
近年来,谱优化方法(特别是μon)展示了通过约束更新步骤到Stiefel流形可显著加速训练和提升泛化能力。然而,μon隐式假设优化景观等价性,强制所有特征方向的谱更新范数均匀。我们认为这种“平等”约束在深度神经网络中是次优的,因为曲率谱已知高度厚尾且病态。在这样的景观中,μon可能放大高曲率方向的不稳定性,同时限制平坦方向的必要进展。本文提出Mousse(Muon Optimization Utilizing Shampoo's Structural Estimation),一种新的优化器,协调谱方法的结构稳定性与二阶预条件化的几何适应性。Mousse不直接对动量矩阵应用牛顿-施卢茨正交化,而是在由克罗内克因子统计(源自Shampoo)诱导的白化坐标系统中操作。数学上,我们将Mousse公式化为受各向异性信任区域约束的谱最陡下降问题的解,其中最优更新通过白化梯度的极分解得到。在从160M到800M参数的语言模型上的实验证明,Mousse在训练步骤减少约12%的同时,计算开销 negligible。
英文摘要:
Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improve generalization. However, Muon implicitly assumes an isotropic optimization landscape, enforcing a uniform spectral update norm across all eigen-directions. We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned. In such landscapes, Muon risks amplifying instabilities in high-curvature directions while limiting necessary progress in flat directions. In this work, we propose \textbf{Mousse} (\textbf{M}uon \textbf{O}ptimization \textbf{U}tilizing \textbf{S}hampoo's \textbf{S}tructural \textbf{E}stimation), a novel optimizer that reconciles the structural stability of spectral methods with the geometric adaptivity of second-order preconditioning. Instead of applying Newton-Schulz orthogonalization directly to the momentum matrix, Mousse operates in a whitened coordinate system induced by Kronecker-factored statistics (derived from Shampoo). Mathematically, we formulate Mousse as the solution to a spectral steepest descent problem constrained by an anisotropic trust region, where the optimal update is derived via the polar decomposition of the whitened gradient. Empirical results across language models ranging from 160M to 800M parameters demonstrate that Mousse consistently outperforms Muon, achieving around $\sim$12\% reduction in training steps with negligible computational overhead.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。