arXiv:2606.30509cs.LG2026-06

Muon优化器可快速学习矩阵分解中的平衡解,避免梯度下降的慢收敛问题。

Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

论文配图:Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics
图 1 · 摘自论文原文
  • Muon通过均衡各主成分的学习速率,避免了初始小值下的鞍点间缓慢迁移。
  • 即使学习率超过临界阈值,仍能稳定收敛,实现指数级退火加速训练。
  • 首次揭示其守恒量机制,适合需快速对齐权重的推荐系统与低秩建模场景。

矩阵分解(形式为 $\ ext{min}_{\ extbf{P},\ extbf{Q}} \ ext{\ extbardbl}\ extbf{M}^\ extster - \ extbf{P}^\ op\ extbf{Q}\ ext{\ extbardbl}_\ ext{F}^2$)是兼具非线性参数动态和表征学习能力的最小学习问题。本文研究了Muon优化器与梯度下降在参数轨迹上的差异。发现三类主要动力学区别:1)Muon从初始小值起避免慢速鞍点间迁移,以相同比例学习所有主模式,较小模式率先收敛;2)当学习率超过由局部损失尖锐度决定的临界阈值时仍保持稳定,摆脱问题条件数约束,支持指数学习率退火实现快速收敛;3)权重对齐后,Muon流守恒量为 $\ ext{sqrt}(\ extbf{P}^\ op\ extbf{P}) - \ ext{sqrt}(\ extbf{Q}^\ op\ extbf{Q})$,而梯度流守恒 $\ extbf{P}^\ op\ extbf{P} - \ extbf{Q}^\ op\ extbf{Q}$。尽管守恒量不同,两者均从零初始化出发获得所谓的‘平衡解’。训练从随机小值开始时,权重会早期自发对齐。我们在简单设定下推导对齐速率,并证明其可预测一般情况下的经验对齐速率。最后,利用Muon的结构特性设计学习率调度,在仅两步内实现近乎完美的权重对齐。

原文摘要 · Abstract (English)

Matrix factorization (i.e., problems of the form $\min_{\mathbf{P},\mathbf{Q}} \|\mathbf{M}^\star - \mathbf{P}^\top\mathbf{Q}\|_\mathrm{F}^2$) is a minimal learning problem that exhibits both nonlinear parameter dynamics and representation learning. In this setting, we study how parameter trajectories under the Muon optimizer differ from those of gradient descent. We identify three main dynamical differences: 1) Muon avoids the slow saddle-to-saddle dynamics from small initialization. Muon instead learns all the top modes of $\mathbf{M}^\star$ at the same rate, with the smaller modes converging first. 2) Muon remains stable even when the learning rate exceeds the critical threshold set by the local loss sharpness. This frees the learning rate from the condition number of the problem, enabling rapid convergence via exponential learning rate annealing. 3) Once the weights are aligned with each other and the target, Muon flow conserves the matrix quantity $\sqrt{\mathbf{P}^\top \mathbf{P}}-\sqrt{\mathbf{Q}^\top \mathbf{Q}}$, while gradient flow is known to conserve the matrix $\mathbf{P}^\top\mathbf{P} - \mathbf{Q}^\top\mathbf{Q}$. Despite having distinct conserved quantities, both optimizers find the so-called \textit{balanced} solution from vanishing initialization. When training from small random initialization, the weights spontaneously align early in training. We derive the alignment rates in simple settings and show that they predict the empirical alignment rates in general. Finally, we exploit structural properties of Muon to construct a learning rate schedule that achieves near-perfect alignment in only two optimization steps.

矩阵分解优化器权重对齐快速收敛

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。