Muon+: Towards Better Muon via One Additional Normalization Step

本文提出了 Muon+，一种在 Muon 优化器正交化后增加归一化步骤的简单有效改进方案，通过在多种模型规模和架构上的大规模预训练实验，证明了其能显著提升训练和验证困惑度。

要点

这篇论文介绍了一个名为 MUON+ 的新方法，它是为了改进大语言模型（LLM）训练中的一种优化器——Muon。

为了让你更容易理解，我们可以把训练一个大型 AI 模型想象成指挥一支庞大的交响乐团演奏一首复杂的交响曲。

1. 背景：为什么要改进？

• 现状：现在的 AI 模型（如 GPT、LLaMA）就像超级庞大的乐团，有几十亿甚至上百亿个乐手（参数）。训练它们非常烧钱、烧算力。
• 旧工具（Muon）：之前，研究人员发现 Muon 这个“指挥家”很不错。它有一个绝招叫“正交化”（Orthogonalization）。
  • 比喻：想象乐团里的乐手们（梯度）一开始乱成一团，有的声音太大，有的太小，甚至有人互相干扰（梯度秩坍塌）。Muon 的“正交化”就像是一个严厉的指挥，强行让所有乐手站成整齐的方阵，确保每个人都在自己的轨道上，互不干扰，这样演奏（训练）效率就高了。
• 问题：虽然 Muon 已经很棒了，但作者发现，在让乐手站好队形（正交化）之后，如果直接开始演奏，声音的音量（更新幅度）可能还是忽大忽小，不够稳定。

2. 核心创新：MUON+ 做了什么？

作者提出了 MUON+，它只在 Muon 的基础上加了一个非常简单但极其有效的步骤：“额外归一化”（Additional Normalization Step）。

• 比喻：
  • Muon 的做法：指挥让乐手站好队形（正交化），然后直接喊“开始演奏！”。
  • MUON+ 的做法：指挥让乐手站好队形后，多做一个动作——给每个乐手发一个“音量调节器”（归一化）。
  • 具体操作：这个“音量调节器”会检查每个乐手（或者每一行/每一列的乐手组）的声音大小，把太响的调小，把太轻的调大，确保大家出来的声音力度均匀、平衡。

3. 这个“音量调节”有什么用？

论文通过大量的实验（在从 1 亿参数到 10 亿参数不等的模型上测试）发现，加上这个“音量调节”后：

1. 训练更稳：就像乐团演奏时，不会因为某个声部突然太吵而盖过其他声部，整个训练过程更加平稳，不容易“翻车”。
2. 效果更好：最终训练出来的模型（交响曲的成品）质量更高，理解能力和生成能力更强（困惑度 Perplexity 更低）。
3. 适应性强：
   • 不管模型是“小乐团”（1 亿参数）还是“超级乐团”（10 亿参数），这个方法都管用。
   • 不管训练时间是“短跑”（计算最优，数据量刚好）还是“马拉松”（数据量极大，训练很久），它都能保持优势。
   • 它对“指挥”（学习率）的敏感度降低了，也就是说，即使你选的学习率不是最完美的，加了 MUON+ 也能练出好结果。

4. 实验结果：真的有效吗？

作者做了很多测试，结果非常漂亮：

• GPT 风格模型：在 1.3 亿到 7.7 亿参数的模型上，MUON+ 总是比原来的 Muon 表现更好，错误率明显降低。
• LLaMA 风格模型：在 6000 万到 10 亿参数的模型上，同样全面胜出。
• 超大规模训练：即使把训练数据量增加到原来的 10 倍（从“计算最优”变成“工业级”），MUON+ 依然能保持领先，说明它非常耐用。

5. 总结：为什么这很重要？

这就好比在造火箭。

• Muon 已经是一个很好的火箭推进系统了。
• MUON+ 并没有发明新的燃料，也没有把火箭造得更大，它只是在推进器后面加了一个小小的“稳定翼”（归一化步骤）。
• 结果就是：火箭飞得更稳、更准，而且能飞得更远，成本还几乎没增加。

一句话总结：
这篇论文告诉我们，在训练大 AI 模型时，在让参数“站好队”（正交化）之后，再给它们“调好音量”（归一化），就能让模型学得更快、更稳、更好。这是一个简单、低成本但回报极高的改进。

技术摘要

这篇论文提出了一种名为 MUON+ 的新型优化器，旨在改进现有的 Muon 优化器，以提升大语言模型（LLM）预训练的效率和质量。以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• 背景：随着大语言模型（如 GPT、LLaMA 等）参数规模和训练数据量的急剧增加，预训练的计算成本变得极高。优化器的选择对于提高训练效率和模型性能至关重要。
• 现状：虽然 Adam/AdamW 仍是主流，但 Muon 优化器因其通过牛顿 - 舒尔茨（Newton-Schulz）迭代对动量矩阵进行正交化，能有效对抗梯度秩坍塌（Rank Collapse），在大规模预训练中表现出优异性能，已被 Kimi、GLM 等领先模型采用。
• 痛点：尽管 Muon 表现良好，但作者发现其性能仍有提升空间。现有的改进方案（如 NorMuon, Maon）往往引入了复杂的机制（如神经元自适应缩放或流形优化），但作者通过消融实验发现，这些改进中的核心收益其实主要来自于正交化后的归一化步骤，而非其他复杂组件。

2. 方法论 (Methodology)

MUON+ 的核心思想极其简单：在 Muon 的正交化步骤之后，增加一个额外的归一化步骤。

• 基本流程：

  1. 动量更新：$M_t = \mu M_{t-1} + (1-\mu)G_t$
  2. 正交化：$O_t = \text{Ortho}(M_t)$（使用牛顿 - 舒尔茨迭代近似 SVD，得到 $U V^T$）
  3. 新增归一化：$O'_t = \text{Norm}^{(d)}(O_t)$（这是 MUON+ 的关键创新）
  4. 权重更新：$W_t = W_{t-1} - \eta \sqrt{m/n} \cdot O'_t$

• 归一化方向：
  论文探讨了多种归一化方向 $d$：

  • 列归一化 (Column-wise)：对矩阵的每一列进行 $\ell_2$ 归一化。
  • 行归一化 (Row-wise)：对矩阵的每一行进行 $\ell_2$ 归一化。
  • 组合归一化：如 col_row（先列后行）或 row_col（先行后列）。
  • 实验表明，行归一化通常优于列归一化，而组合归一化（特别是 col_row 和 row_col）通常能取得最佳效果。

• 设计动机：
  作者认为，正交化后的更新步长如果缺乏结构化的归一化，可能导致优化不稳定。通过引入归一化，可以稳定优化过程，使模型在更宽的超参数范围内（特别是学习率）表现更稳健。

3. 关键贡献 (Key Contributions)

1. 提出 MUON+：在 Muon 基础上增加一步简单的归一化操作，显著提升了预训练性能。
2. 广泛的实验验证：
   • 模型规模：涵盖了从 60M 到 1B 参数的 LLaMA 风格模型，以及 130M 到 774M 参数的 GPT 风格模型。
   • 训练 regime：不仅验证了计算最优（Compute-optimal，T2P 比约为 20）的训练，还扩展到了工业级的大数据量训练（T2P 比高达 200，即过训练场景）。
3. 系统性消融研究：
   • 证明了归一化本身是性能提升的主要驱动力，而非其他复杂的自适应缩放机制。
   • 分析了不同归一化方向（行/列/组合）的影响。
   • 验证了 MUON+ 对不同正交化近似方法（如 You, Jordan, PolarExpress）的鲁棒性。
4. 开源代码：提供了完整的实现代码。

4. 实验结果 (Results)

• 计算最优训练 (Compute-Optimal)：
  • 在 GPT 模型上，MUON+ 相比 Muon 在验证困惑度（Perplexity）上降低了 0.91 到 2.02 不等。
  • 在 LLaMA 模型上，MUON+ 同样在所有规模（60M - 1B）上优于 Muon 和 AdamW 基线，困惑度降低 0.37 到 0.61。
• 过训练场景 (Overtraining, T2P ≈ 200)：
  • 在 GPT-Base (362M) 和 LLaMA-350M 上使用 720 亿 Token 进行训练。
  • 结果显示 MUON+ 在长周期训练中依然保持更低的验证困惑度（GPT-Base 降低 1.13，LLaMA-350M 降低 0.45），且优化过程更加稳定，没有出现性能退化。
• 超参数敏感性：
  • MUON+ 对学习率的敏感度更低。在次优（过大）的学习率下，Muon 性能下降明显，而 MUON+ 仍能保持较好的性能，表明其具有更强的鲁棒性。
• 归一化方向选择：
  • 组合归一化（col_row 或 row_col）通常表现最好。
  • 行归一化（Row-wise）单独使用时通常优于列归一化。

5. 意义与影响 (Significance)

• 简单即有效：MUON+ 证明了在复杂的优化器设计中，一个极简单的修改（增加一步归一化）就能带来显著且一致的性能提升。这为优化器设计提供了新的视角：关注正交更新的结构化归一化可能比引入复杂的二阶矩估计或流形约束更重要。
• 工业级适用性：实验覆盖了从计算最优到工业级大数据量（T2P=200）的训练场景，证明了 MUON+ 不仅适用于小规模实验，也适用于大规模生产环境的预训练。
• 稳定性提升：MUON+ 降低了模型对超参数（特别是学习率）的敏感性，使得大规模训练更加稳健，减少了调参成本。
• 理论启示：该工作暗示了梯度正交化后的“尺度控制”是优化稳定性的关键，为理解 Muon 类优化器的内在机制提供了新的解释。

总结：MUON+ 通过引入正交化后的归一化步骤，以极低的计算代价显著提升了 Muon 优化器的性能、稳定性和适用范围，是目前大语言模型预训练中一个极具实用价值的优化器改进方案。