Mousse改进了穆翁优化器,让模型训练更稳定高效。
Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning
- 在白化坐标系中用结构估计自适应调整更新方向
- 训练步数减少约12%,参数量160M到800M均有效
- 适合对收敛稳定性要求高的大模型训练场景
近期谱优化方法(如Muon)通过将更新步骤约束在Stiefel流形上,显著加速训练并提升泛化能力。然而,Muon隐含假设优化空间各向同性,对所有特征方向施加统一的谱更新范数。我们指出,这一‘平等’约束在深度神经网络中并不理想——其曲率谱通常呈重尾分布且病态。在此类空间中,Muon可能放大高曲率方向的不稳定性,同时限制平坦方向的必要进展。本文提出新型优化器Mousse(Muon Optimization Utilizing Shampoo's Structural Estimation),融合谱方法的结构稳定性和二阶预条件的几何自适应性。不同于直接对动量矩阵应用Newton-Schulz正交化,Mousse在由克罗内克分解统计量导出的白化坐标系中运作。数学上,我们将Mousse建模为带各向异性信任区域的谱最速下降问题,最优更新通过白化梯度的极分解求解。在160M至800M参数的语言模型上,实验表明Mousse始终优于Muon,训练步数减少约12%,计算开销可忽略。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。