Muon优化器：AI模型训练算法的下一个里程碑？| 目前还不是业界焦点，但有潜力是重大基础创新 精华

人工智能（AI）快速发展，模型训练是核心环节，优化器扮演着至关重要的角色，它负责调整模型的参数，让模型在数据上表现得更好。多年来，AdamW优化器一直是优化器的标杆，因其稳定性和高效性深受研究者和工程师的喜爱。然而，随着AI模型规模的不断扩大，训练成本和时间的需求也在激增，这让人们开始寻找更高效的优化方法。近期，一种名为Muon的优化器算法悄然出现（源代码 ​​​​https://github.com/KellerJordan/Muon ​​​​），尽管它还未成为业界焦点，但其独特的设计和卓越的性能表明，它可能是AI模型训练领域的一次重大基础创新。

优化器算法：AI训练的“幕后推手”

为什么优化器如此重要？

在深度学习中，模型训练的目标是通过调整参数，让模型的预测结果尽可能接近真实数据。这个过程通常是通过定义一个损失函数来实现的，损失函数衡量了模型预测与真实值之间的差距。而优化器的任务，就是根据损失函数的梯度（gradient），一步步调整模型的参数，找到损失最小的“最佳状态”。

想象一下，我们在一个崎岖的山谷中寻找最低点。优化器就像我们的导航仪，告诉我们每一步该往哪个方向走、走多远。一个好的优化器不仅能更快地带我们到达谷底（收敛），还能避免在陡峭的坡道上上上下下（训练不稳定）。优化器的效率直接决定了训练速度、计算资源需求，以及最终模型的性能。

过去几年，AdamW（Adam with Weight Decay）一直是训练大型语言模型的首选（如Qwen、DeepSeek、LLaMA等，闭源的大模型不清楚，但大概率也是AdamW）。它结合了自适应学习率和权重衰减（weight decay）的优点，能够在复杂的参数空间中稳定、高效地工作。然而，随着模型参数从几亿增加到几千亿，训练时间从几天变成几周甚至几个月，AdamW的局限性开始显现——它在超大规模场景下的效率开始受到挑战。进一步提升AI能力，我们需要更大的模型和更多的训练资源。但计算资源的成本高昂，训练时间过长也会拖慢研究和应用的进度。因此，开发更高效的优化器，不仅是技术上的追求，更是经济和实践上的迫切需求。

Muon：从动量到正交化的革新

Muon的基本原理

Muon的全称是MomentUm Orthogonalized by Newton-Schulz（动量正交化Newton-Schulz），它是一种专为神经网络隐藏层设计的优化器。它的核心思想并不复杂，但却非常巧妙：先用经典的SGD（随机梯度下降）动量法生成参数更新，然后通过一个特殊的“后处理”步骤——Newton-Schulz正交化，让这些更新更高效。

让我们一步步拆解这个过程：

1. 动量法生成更新Muon首先沿用了SGD动量的思路。动量法就像给梯度加了一个“惯性”，让参数更新不仅依赖当前梯度，还参考之前的更新方向。这样可以加速收敛，避免在参数空间中“左摇右晃”。在Muon中，这一过程生成一个更新矩阵（记为(B_t)）。
2. 正交化：让更新更“聪明”接下来，Muon引入了关键创新：通过Newton-Schulz迭代，将更新矩阵(B_t)“正交化”，生成一个新的更新矩阵(O_t)。所谓正交化，简单来说，就是让更新矩阵的列（或行）彼此垂直，变成一个“半正交矩阵”（满足( O^TO = I )或( O O^T= I )）。为什么这样做？传统的动量更新往往会被某些“主导方向”牵着走，而其他“稀有方向”（对学习也很重要，但幅度小）的贡献被掩盖。正交化就像重新分配了这些方向的“发言权”，让更新更全面、更高效地探索参数空间。
3. 参数更新最后，Muon用正交化后的更新矩阵( O_t)来调整参数：
   
   

这里，η是学习率，θ是模型参数。

Newton-Schulz迭代：高效的正交化工具

正交化听起来很高级，计算上也很复杂。如果用传统的SVD（奇异值分解）来正交化，计算量太大，速度太慢，完全不适合现代GPU加速的训练环境。而Muon采用的Newton-Schulz迭代，则是一个高效的替代方案。

这个迭代过程的核心是，通过反复计算一个多项式函数（比如五次多项式），逐步将更新矩阵调整到接近正交的状态。它的优点在于：

• 低计算开销：每次迭代只需几次矩阵乘法，远比SVD快。
• 数值稳定性：可以用bfloat16（一种低精度浮点格式）运行，非常适合现代GPU。

例如，Muon的实现中，Newton-Schulz迭代的公式是：

其中，(a = 3.4445, b = -4.7750, c = 2.0315)是精心调优的系数。经过5次迭代，就能得到一个“足够正交”的更新矩阵，既高效又实用。

Muon的优势：效率与潜力的结合

Muon并不是凭空出现的“黑魔法”，它的优势经过了多次实验验证。以下是Muon相比AdamW的几个突出优点：

1. 更快的训练速度在多个基准测试中，Muon展现了惊人的加速能力。例如：

• 在CIFAR-10图像分类任务中，Muon将达到94%准确率的时间从3.3 A100-seconds降到2.6 A100-seconds。
• 在NanoGPT speedrunning任务（训练一个小规模GPT模型）中，Muon将训练速度提升了1.35倍。 这些结果表明，Muon能显著缩短训练时间，尤其是在资源有限的场景下。

2. 卓越的可扩展性Muon在大规模模型上同样表现出色。例如，在训练一个1.5亿参数的语言模型时，Muon仅用10个8xH100-hours就达到了GPT-2 XL的性能水平，而AdamW需要13.3 hours。随着模型规模继续扩大，这种差距可能会更明显。
3. 低计算开销尽管正交化听起来很复杂，但Muon的额外计算开销非常小。在典型语言模型训练中，Muon的FLOP（浮点运算）开销低于1%。这得益于Newton-Schulz迭代的高效性，以及对现代硬件的优化。
4. 与现有框架兼容Muon并非完全替代AdamW，而是与之互补。实际应用中，Muon通常优化网络的隐藏层参数，而嵌入层（embedding）和分类器头（head）仍由AdamW处理。这种“分工合作”的方式，让Muon可以无缝集成到现有训练流程中。

如果Muon替代AdamW

假设Muon的潜力被充分验证，并逐步取代AdamW成为新的标准优化器，AI领域可能会迎来以下重大变化：

1. 训练成本大幅降低Muon的高效性意味着，训练一个大规模模型所需的计算资源和时间将显著减少。这不仅能节省数千万RMB的算力成本，还能让更多中小型研究团队参与到前沿AI研究中，降低技术门槛。
2. 模型规模的进一步突破更高的训练效率，让研究者有能力尝试更大规模的模型。比如，现在训练一个1000亿参数模型可能需要数月，而有了Muon，或许几周就能完成。这将推动AI能力的上限不断提升。
3. AI应用加速落地更快的训练速度意味着模型从研究到应用的时间缩短。无论是工业大模型（创新奇智在做的事情），还是其他行业大模型，AI技术的普及速度都可能因此加快。
4. 优化器研究的复兴Muon的成功可能会重新点燃对优化器算法的兴趣。过去几年，AdamW几乎“一统江湖”，新优化器的研究相对沉寂。如果Muon证明了创新优化器的价值，研究者可能会投入更多精力，探索其他潜在的突破。

OpenAI：Muon的未来舞台？

2024年，Muon的开发者Jordan在个人博客中详细介绍了Muon的设计理念和实验结果，随后不久，他宣布加入OpenAI。这一动向并非巧合，很可能OpenAI看中了Muon的潜力。

月之暗面的实践：Muon的验证

Moonshot AI（月之暗面）近期在Muon的基础上进行了大规模实践，验证了它在真实场景中的潜力。他们不仅改进了Muon，还训练了一个名为Moonlight的3B/16B参数Mixture-of-Expert（MoE）模型，用5.7万亿tokens的数据进行了测试。

Moonshot AI的改进

Moonshot AI发现，原始Muon在小规模任务上表现出色，但在超大规模训练中会遇到问题，比如模型权重增长过大，影响稳定性。为此，他们提出了两个关键改进：

1. 引入权重衰减他们将AdamW的权重衰减机制融入Muon，更新公式变为：这有效控制了权重大小，提升了长期训练的性能。
2. 调整更新尺度Muon的更新幅度（RMS）会因参数矩阵的形状而变化，可能导致训练不稳定。Moonshot AI提出按矩阵最大维度缩放更新，比如：

 这样可以保持更新幅度一致，并与AdamW兼容。

Moonlight的惊艳表现

基于这些改进，Moonshot AI用Muon训练了Moonlight模型，并在多个基准测试中取得了优异成绩。例如：

• 在MMLU（英语理解）上，Moonlight得分70.0，超越了同规模的Llama3.2-3B（54.7）和Deepseek-v2-Lite（58.3）。
• 在GSM8K（数学推理）上，得分77.4，接近Qwen2.5-3B（79.1），但训练tokens仅为后者的三分之一。
• 训练效率上，Moonlight只需约52%的FLOPs，就能达到AdamW的性能水平。

这些结果表明，Muon不仅能加速训练，还能提升模型性能，尤其在数学和代码任务上表现突出。

结语：Muon的潜力与未来

Muon作为一种新兴的优化器算法，以其独特的设计和卓越的性能，展现了巨大的潜力。它通过正交化更新矩阵，打破了传统优化器的局限，在训练速度、可扩展性和计算效率上都超越了AdamW。尽管目前它还未引起全球关注，但Keller Jordan的加入OpenAI，以及Moonshot AI的成功实践，都预示着Muon可能成为AI训练领域的下一个里程碑。

​ 图片 ​

本文转载自​​后向传播​​，作者：张发恩