🤖 machine learning

MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training

本文介绍了 MONA，这是一种新颖的优化器，它将 Nesterov 式加速集成到 Muon 框架中，以逃离尖锐的局部极小值，并在参数量从 10 亿到 680 亿的语言模型上实现最先进的收敛性和下游性能。

原作者： Jiacheng Li, Jianchao Tan, Hongtao Xu, Jiaqi Zhang, Yifan Lu, Yerui Sun, Yuchen Xie, Xunliang Cai

发布于 2026-05-27

📖 1 分钟阅读☕ 轻松阅读

原作者： Jiacheng Li, Jianchao Tan, Hongtao Xu, Jiaqi Zhang, Yifan Lu, Yerui Sun, Yuchen Xie, Xunliang Cai

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教一个巨大且超级聪明的机器人说人类语言。为此，你必须调整其大脑中数百万个微小的旋钮。寻找这些旋钮完美设置的过程被称为“训练”，而用来转动这些旋钮的工具被称为优化器。

长期以来，这项工作的标准工具被称为AdamW。它像一位稳健、谨慎的徒步者，迈着小步并不断查看地图。最近，一种名为Muon的新工具问世了。Muon 像是一位能一眼纵观整座山脉的徒步者，将成组的旋钮视为一个整体。这使得 Muon 更快、更高效，但它仍有一个缺陷：它可能会陷入“尖锐山谷”。

问题：尖锐山谷陷阱

想象机器人学习过程的景观是一片丘陵地带。

平坦山谷是理想的定居点；这意味着机器人找到了一个稳定的、通用的解决方案，能在多种情况下良好运作。
尖锐山谷则是狭窄、深邃的坑洞。如果机器人掉进其中一个，它可能看似找到了最低点，但实际上却被困在一个非常具体且脆弱的地方。如果你轻轻推它一下，它就会从那个“完美”的位置跌落，表现变差。

AdamW 和 Muon 都可能意外掉入这些尖锐山谷并被困其中。它们缺乏一种方式来“感知”前方地面的形状，从而判断某个位置是否过于尖锐。

解决方案：MONA（带 Nesterov 加速的 Muon）

本文的作者创造了一种名为MONA的新工具。可以将 MONA 想象为配备了一副声呐护目镜的 Muon。

以下是通过一个简单类比的工作方式：

旧方法（Muon）： Muon 查看它站立位置的正下方坡度并向下移动。它在高效移动方面表现出色，但它不知道前方的地面是平缓的斜坡还是悬崖边缘。
新方法（MONA）： 在 Muon 迈出一步之前，MONA 会查看坡度从上一站到当前站是如何变化的。
- 如果坡度变化非常突然（就像撞上陡峭的悬崖），MONA 会意识到：“哇，这是一个尖锐山谷！让我们用力推一把，弹出去。”
- 如果坡度变化非常缓慢（一个平缓、平坦的山谷），MONA 会说：“这看起来很安全。让我们在这里安顿下来。”

MONA 在数学公式中添加了一个特殊的“加速项”。它计算当前方向与先前方向之间的差异。这种差异就像一个传感器，用于检测地形的“尖锐度”。如果地形尖锐，它就会将机器人推出坑洞；如果地形平坦，它就让机器人休息。

他们发现了什么？

研究人员在三种不同规模的語言模型（小、中、大）上测试了这一新工具，参数量从 10 亿到 680 亿不等。他们在海量文本（多达 1 万亿个单词）上对它们进行了训练。

更好的训练： 在每一项测试中，MONA 都帮助模型学得更快，并达到比旧标准（AdamW）和新工具 Muon 更好的最终状态。
更聪明的机器人： 使用 MONA 训练的模型在通用任务、数学问题和代码编写方面表现更佳。例如，在一个 680 亿参数的模型上，MONA 在数学和代码编写的基准测试中取得了最高分。
微调： 即使在初始训练之后，当给予模型额外的特定训练（例如专门教它们编写代码）时，那些以 MONA 为起点的模型也比以 Muon 为起点的模型表现更好。

使其实用化（MONA-Lite）

作者们也意识到，添加这种“声呐”系统会占用稍多的计算机内存。为了解决这个问题，他们创建了一个更轻量级的版本，称为MONA-Lite。

他们使用了一种技巧来压缩内存数据（就像从高清视频切换到标清视频），并采用了一种方法，在不存储额外文件的情况下实时计算坡度变化。
这将所需的额外内存减少了约 75%，使其即使在内存有限的计算机上也易于使用，同时不会损失其优势。

总结

简而言之，MONA是 Muon 优化器的升级版。它保留了 Muon 的速度和效率，但增加了一个“曲率传感器”，帮助 AI 在学习过程中避免陷入糟糕的尖锐位置。这造就了更智能、更强大的语言模型，它们在数学、代码和通用推理方面表现更佳，同时效率足以在标准硬件上运行。

技术摘要：MONA——带 Nesterov 加速的 Muon 优化器

1. 问题陈述

训练大型语言模型（LLM）日益需要具有卓越样本效率的优化器。尽管 AdamW 作为标准已沿用超过十年，但近期出现了如 Muon 等替代方案。Muon 通过将权重矩阵视为单一几何单元，并对动量缓冲区应用矩阵正交化（通过 Newton-Schulz 迭代），改进了传统的一阶方法。这种方法在谱范数下与最速下降法一致，在大规模混合专家（MoE）训练中，其计算效率比 AdamW 高出约 $2\times$ 。

然而，Muon 与所有一阶梯度方法共享一个根本局限：它缺乏超越局部梯度信息来探索损失景观的显式机制。因此，Muon 容易陷入尖锐的局部极小值，特别是在大批量训练场景中。尖锐极小值与较差的泛化能力相关，而平坦的极小值往往能带来更好的下游性能。虽然 ALTO 优化器引入了基于梯度差指数移动平均（EMA）的加速项以逃离尖锐极小值，但其作为 Adam 和 LAMB 通用适配器的设计，使其与 Muon 高度非线性的正交化过程集成时面临挑战。

2. 方法论：MONA

作者提出了 MONA（带 Nesterov 加速的 Muon 优化器），将曲率感知的加速项无缝集成到 Muon 的正交化框架中。

核心算法

MONA 通过在动量累积和正交化之前引入加速项，修改了标准的 Muon 流程。矩阵参数 $W$ 在第 $k$ 次迭代的更新过程如下：

梯度差计算：计算当前随机梯度 $G_k$ 与上一梯度 $G_{k-1}$ 之间的差值：
$D_k = G_k - G_{k-1}$
加速缓冲区更新：维护这些差值的指数移动平均：
$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$
其中， $\beta_a$ 控制加速的记忆长度。
曲率感知梯度：通过将缩放后的加速项加到原始梯度上，构建增强梯度 $\tilde{G}_k$ ：
$\tilde{G}_k = G_k + \alpha A_k$
系数 $\alpha$ 对加速项进行缩放。理论动机表明， $G_k - G_{k-1}$ 近似于 $H_k(\theta_k - \theta_{k-1})$ ，其中 $H_k$ 是 Hessian 矩阵。因此，该项指向远离尖锐极小值（此处梯度变化较大）的方向，而在平坦区域保持较小。
Muon 正交化：将增强梯度输入到标准的 Muon 动量和正交化步骤中：
$M_k = \mu M_{k-1} + \tilde{G}_k$
$O_k = \text{Newton-Schulz}(M_k)$
$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$

设计原理

在正交化之前应用加速项至关重要。如果在之后应用，将会破坏 Muon 赖以进行谱范数正则化的正交结构。通过首先将原始梯度转换为曲率感知方向，动量缓冲区能够捕捉这些丰富的方向，随后再经由 Muon 的标准流程进行处理。

对于向量参数（如嵌入、偏置），MONA 默认采用 AdamW，遵循 Muon 现有的惯例。

3. 理论分析

该论文在标准假设（L-平滑、无偏梯度且方差有界、期望方向对齐）下，提供了 MONA 的收敛性分析。

收敛速率：作者证明，在非凸设置下，MONA 实现了 $O(K^{-1/2})$ 的收敛速率，与 SGD 和 AdamW 的渐近速率相匹配。加速项影响界限中的常数，但不影响渐近速率。
逃离尖锐极小值：该分析形式化了加速项如何选择性地抵抗尖锐极小值。在驻点附近，期望加速 $E[A_k]$ 与过去方向的 Hessian 加权总和成正比。对于具有大特征值的方向（尖锐极小值），加速幅度较大，推动优化器远离；对于小特征值（平坦方向），加速较小，允许稳定收敛。该机制在保留 Muon 谱范数正则化的同时，增加了曲率感知能力。

4. 实证结果

MONA 在三种规模的 MoE 预训练模型（1B、6B 和 68B 参数）以及监督微调（SFT）阶段进行了评估。

预训练性能

模型：MOE-1B-A0d2B（4000 亿 token）、MOE-6B-A0d5B（1.2 万亿 token）和 MOE-68B-A3B（7000 亿 token）。
指标：在代码、数学推理、通用英语和中国学术文本领域监测验证损失。
发现：在所有三个模型规模和所有领域，MONA 始终比 Muon 和 AdamW 实现更低的验证损失。性能差距在更大的 68B 模型中尤为明显，其中 MONA 表现出更稳定的收敛性。

下游基准测试（MOE-68B-A3B）

使用 MONA 训练的 68B 模型在综合基准测试套件上进行了评估：

通用能力：MONA 在包括 MMLU、MMLU-Pro、CMMLU、CEVAL 和 BBH 在内的基准测试中取得了最高的平均分（0.4557），优于 Muon（0.4478）和 AdamW（0.4382）。
数学推理与代码生成：MONA 在 GSM8K、GPQA 以及各种代码生成基准测试（HumanEval+、MBPP+、BigCodeBench、LiveCodeBench）中表现出卓越性能。值得注意的是，MONA 在需要可迁移代码表示的任务上取得了显著提升。
监督微调（SFT）：当 MONA 预训练的检查点在代码数据上进行 SFT 时，与 Muon 预训练的基线相比，其在整个过程中保持了更低的训练损失。BigCode 套件上的 SFT 后评估证实，与 Muon 初始化的模型相比，MONA 初始化的模型在 7 项任务中的 6 项上取得了更高的分数。

效率与部署

为解决存储上一梯度（ $G_{k-1}$ ）和加速缓冲区（ $A_k$ ）带来的内存开销，作者引入了 MONA-Lite：

量化：将辅助状态存储为 bfloat16（BF16）而非 float32。
流式计算：就地计算梯度差，消除对专用 $G_{k-1}$ 缓冲区的需求。
结果：这些策略在不牺牲训练质量的情况下，将额外内存开销减少了约 75%。计算开销分析显示，MONA 在优化器步骤层面的速度仅慢约 1%，与 Muon 相比，端到端迭代时间无实际差异。

5. 意义与主张

该论文声称，MONA 成功弥合了 Muon 的几何效率与曲率感知加速需求之间的差距。其主要意义在于：

算法集成：证明了加速项可以集成到矩阵正交化框架中，而不会破坏其几何属性。
可扩展性：在广泛的模型规模（1B 至 68B 参数）和 token 数量（高达 1 万亿 token）上验证了该方法，显示出相对于 AdamW 和 Muon 的一致改进。
泛化能力：提供了实证证据，表明在预训练期间逃离尖锐极小值能带来更好的下游任务泛化能力，特别是在复杂推理和代码生成方面。
实用性：提供了一种内存高效的变体（MONA-Lite），使该优化器适用于资源受限的环境。

作者总结道，MONA 代表了大规模 LLM 训练的最先进优化器，在保持 Muon 框架计算效率的同时，实现了卓越的收敛性和下游性能。