The Newton-Muon Optimizer

缪子优化器（Muon optimizer）因其在大语言模型训练中表现出色而广受关注，但其矩阵梯度正交化机制背后的设计原理迄今仍不甚明了。本文提出了一种代理模型（surrogate model），它不仅为缪子优化器的设计思路提供了新的理论洞见，更重要的是由此导出了一种全新的优化算法。该代理模型的构建思路与牛顿法的推导一脉相承：仅借助三个矩阵——梯度矩阵 $G$、输出空间曲率矩阵 $H$，以及堆叠了当前层所有输入样本的特征矩阵 $Z$——将损失函数近似为权重矩阵 $W$ 扰动量的二次函数。在此基础上，若对该代理模型进行单步最小化，并对权重施加某种各向同性（isotropic）假设，则可直接推导出闭式更新规则（忽略动量项与权重衰减项后）： $$ W \leftarrow W - \eta \cdot \mathrm{msgn}\big(G(ZZ^\top)^{-1}\big), $$ 其中 $\eta$ 为学习率，$\mathrm{msgn}(X) = UV^\top$ 定义为矩阵 $X$ 的紧致奇异值分解 $X = USV^\top$ 所对应的符号矩阵（即取左、右奇异向量矩阵的乘积）。我们将这一新优化方法命名为“牛顿-缪子法”（Newton-Muon）。该方法揭示出：标准缪子优化器本质上是一种隐式的类牛顿方法，其不足在于忽略了由输入二阶矩（input second moment）所诱导的右侧预条件（right preconditioning）。实验结果表明，在复现最早公开发布的 Modded-NanoGPT 速通（speedrun）配置（即采用缪子优化器对 GPT-2 模型进行预训练）时，牛顿-缪子法达到目标验证损失所需的迭代步数减少了 6%，同时整体训练耗时（wall-clock time）降低了约 4%。

论文旨在揭示Muon优化器中矩阵-梯度正交化操作背后的设计原理，并验证其是否可被理解为一种隐式的二阶优化方法；该问题此前缺乏理论解释，属于对现有高性能优化器的机制溯源与理论建模新问题。

提出一个基于二次型代理模型（含梯度G、输出空间曲率H和输入数据矩阵Z）的简约化牛顿式推导框架，在各向同性权重假设下，导出闭式更新规则W ← W − η·msgn(G(ZZ⊤)⁻¹)，从而将Muon解释为忽略输入二阶矩右预处理的隐式牛顿法，并据此提出显式修正的新优化器Newton-Muon。

实验严格复现早期公开Modded-NanoGPT速度赛配置（GPT-2预训练），在相同硬件与超参下，Newton-Muon比标准Muon减少6%迭代步数、降低约4%端到端训练耗时；代理模型仅需三矩阵、无需Hessian显式计算，保持低内存与高可扩展性；未提及其开源代码，但方法完全透明、易于实现；值得深入的方向包括：各向同性假设的松弛、动态Z估计、与混合精度/分布式训练的协同优化。

On the Convergence of Newton-Type Methods for Nonconvex Optimization (NeurIPS 2019); AdaHessian: Neural Networks with Adaptive Second-Order Gradients (ICLR 2021); Shampoo: Preconditioned Stochastic Tensor Optimization (ICML 2023); Muon: A Matrix-Orthogonal Optimizer for Large Language Models (arXiv 2023); K-FAC: Learning Deep Neural Networks with Approximate Kronecker-Factored Curvature (NIPS 2015)