提出Newton-Muon优化器,显著提升大模型训练效率。
The Newton-Muon Optimizer
- 基于权重扰动的二次逼近,结合梯度、输出曲率与输入数据矩阵设计新方法。
- 在相同配置下比Muon少6%迭代步数,训练时间缩短约4%。
- 适合追求高效训练的大模型研究者,尤其关注优化器改进者。
Muon优化器在训练大语言模型中表现优异,但其矩阵梯度正交化的设计原理尚不清晰。本文提出一个代理模型,不仅揭示了Muon的设计机制,更导出一种新优化器。受牛顿法推导启发,该代理模型将损失近似为权重矩阵 $W$ 扰动的二次函数,仅依赖三个矩阵:梯度 $G$、输出空间曲率矩阵 $H$ 以及堆叠层输入的数据矩阵 $Z$。通过最小化该代理函数并假设权重各向同性,得到闭式更新规则(含动量与权重衰减):$W ightarrow W - η\cdot \mathrm{msgn}(G(ZZ^\top)^{-1})$,其中 $η$ 为学习率,$\mathrm{msgn}(X)=UV^\top$ 当 $X=USV^\top$ 为紧凑SVD。新方法称作Newton-Muon,表明标准Muon可视为忽略输入二阶矩右预处理的隐式牛顿型方法。实验表明,在复现最早公开的Modded-NanoGPT速度跑配置(使用Muon训练GPT-2)中,Newton-Muon以6%更少迭代步数达到目标验证损失,并将实际训练时间减少约4%。
原文摘要 · Abstract (English)
The Muon optimizer has received considerable attention for its strong performance in training large language models, yet the design principle behind its matrix-gradient orthogonalization remains largely elusive. In this paper, we introduce a surrogate model that not only sheds new light on the design of Muon, but more importantly leads to a new optimizer. In the same spirit as the derivation of Newton's method, the surrogate approximates the loss as a quadratic function of the perturbation to a weight matrix $W$ using only three matrices: the gradient $G$, an output-space curvature matrix $H$, and the data matrix $Z$ that stacks the layer inputs. By minimizing this surrogate in one step and adopting a certain isotropic assumption on the weights, we obtain the closed-form update rule (up to momentum and weight decay) $W \leftarrow W - η\cdot \mathrm{msgn}(G(ZZ^\top)^{-1})$, where $η$ is the learning rate and $\mathrm{msgn}(X)=UV^\top$ if $X=USV^\top$ is a compact singular value decomposition. This new optimization method, which we refer to as Newton-Muon, shows that standard Muon can be interpreted as an implicit Newton-type method that neglects the right preconditioning induced by the input second moment. Empirically, on a reproduction of the earliest publicly released Modded-NanoGPT speedrun configuration using Muon for GPT-2 pretraining, Newton-Muon reaches the target validation loss in 6\% fewer iteration steps and reduces wall-clock training time by about 4\%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。