Muon算法的收敛性保证：新的参数区间与推广

中文摘要 AI 辅助

本文通过比典型矩阵符号函数更精确的Newton-Schultz迭代近似，首次为Muon算法建立了渐近收敛性保证。我们证明，在适当的超参数选择下，迭代满足$\u003cbr\u003e$\lim_{k\to\infty}\\|\nabla f(x_k)\\|=0$，并且在全局Polyak-Łojasiewicz条件下，函数值序列线性收敛。关键洞察在于，Muon的Newton-Schulz实现中隐含的正则化产生了一个有界预条件子，将Muon揭示为一种\u003cem\u003e预条件Polyak重球\u003c/em\u003e方法，从而能够进行经典的Lyapunov分析。这一观察自然促使将相同的预条件结构应用于Nesterov梯度评估。我们通过引入\u003cem\u003eMuesterov\u003c/em\u003e（Muon的基于Nesterov的变体）来形式化这一思想，并证明它享有相同的收敛性保证，将理论框架扩展到重球设置之外。在标量交叉熵问题上的数值实验证实了该理论，并阐明了学习率和Newton-Schulz正则化器在控制收敛中的共同作用。训练nanoGPT数据集的初步数值模拟为本文观察结果与实际应用的相关性提供了直觉。

英文摘要

In this paper, we establish the first asymptotic convergence guarantees for the Muon algorithm through a more accurate proxy for the Newton-Schultz iteration than the typical matrix sign function. We prove that, for appropriate choices of hyperparameters, the iterates satisfy $\lim_{k\to\infty}\|\nabla f(x_k)\|=0$, and, under a global Polyak-Łojasiewicz condition, that the sequence of function values converges linearly. The key insight is that the regularization, implicit in Muon's Newton-Schulz implementation, induces a bounded preconditioner, exposing Muon as a \emph{preconditioned Polyak heavy-ball} method and enabling a classical Lyapunov analysis. This observation naturally motivates applying the same preconditioning structure to the Nesterov gradient evaluation. We formalize this idea by introducing \emph{Muesterov}, a Nesterov-based variant of Muon, and prove that it enjoys the same convergence guarantees, extending the theoretical framework beyond the heavy-ball setting. Numerical experiments on a scalar cross-entropy problem corroborate the theory and illuminate the joint role of the learning rate and the Newton-Schulz regularizer in controlling convergence. Preliminary numerical simulations training the nanoGPT dataset provide intuition regarding the relevance of the observations in this paper to practical applications.

发表机构

• Northeastern University（东北大学）
• University of São Paulo（圣保罗大学）
• Escola Politecnica（理工学院）

机构由 AI 辅助整理，请以论文原文为准。