Muon优化器在矩阵分解中避免慢鞍点动力学,实现平衡解

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

精选理由

如果你关注优化器理论,这篇论文揭示了Muon比梯度下降更快的机制,并且给出了一个只需两步就对齐参数的学习率调度,很实用。

AI 摘要

该论文研究Muon优化器在矩阵分解问题中的参数动力学,发现其与梯度下降有三点关键差异:1)Muon避免从小初始化开始的慢鞍点动力学,以相同速率学习所有顶模,较小模先收敛;2)即使学习率超过局部损失尖锐度的临界阈值,Muon仍保持稳定,允许通过指数学习率退火实现快速收敛;3)Muon流守恒矩阵√(P^T P)-√(Q^T Q),而梯度流守恒P^T P - Q^T Q。从零初始化时,两者都能找到平衡解。论文还推导了简单设定下的对齐速率,并利用Muon结构属性设计了仅需两步达到近完美对齐的学习率调度。

原文 · arXiv cs.LG

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

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.

Muon优化器在矩阵分解中避免慢鞍点动力学,实现平衡解 · AI 热点