← 研踪 PaperTrailarXiv:2606.30509 · v2
Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics
首次提交 2026-06-30 00:17
版本更新 2026-11-07 03:09
雷达观测 2026-11-10 10:11
摘要
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 of GD from small initialization. Muon instead learns all the top singular vectors of $\mathbf{M}^\star$ at the same rate, and they saturate in order of increasing (rather than decreasing) singular value. (2) Muon does not diverge 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 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.
观测到的版本
相邻记录可能跳过中间版本;首次发现 v2 及以上也会标为版本更新。
v2 · 2026-11-07 03:09 · 最新观测版本
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 of GD from small initialization. Muon instead learns all the top singular vectors of $\mathbf{M}^\star$ at the same rate, and they saturate in order of increasing (rather than decreasing) singular value. (2) Muon does not diverge 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 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.
查看此版本解读 →查看此版本原文 ↗