今天小编分享的科学经验：月之暗面开源改进版Muon优化器，算力需求比AdamW锐减48%，DeepSeek也适用，欢迎阅读。
算力需求比 AdamW 直降 48%，OpenAI 技术人员提出的训练优化算法 Muon，被月之暗面团队又推进了一步！
团队发现了 Muon 方法的 Scaling Law，做出改进并证明了 Muon 对更大的模型同样适用。
在参数量最高 1.5B 的不同 Llama 架构模型上，改进后的 Muon 算力需求仅为 AdamW 的 52%。
同时团队还基于 DeepSeek 架构训练出了一个 16B 的 MoE 模型，与改进后的优化算法一同开源。
Muon 技术博客发布当时主要适用于较小的模型和数据集，作者留下了三个悬而未决的问题：
Muon 能否用于更大规模的训练？
Muon 能否在更大规模的 GPU 集群上使用？
Muon 是否同样适用于微调和强化学习？
现在月暗团队用实验给出了回答——全部都是 Yes。
消息一出，当时 Muon 的作者也都很激动，主要作者Keller Jordan表示这是 Muon 规模化的首个成功报告，为团队送上了祝贺。
另一名贡献者，当时负责 Muon 规模化实验的 Hyperbolic Labs 联创兼 CTO Yuchen Jin也表示，月暗团队的这项成果，是 Muon 的一次胜利。
将 AdamW 特点引入 Muon
在介绍月暗团队的工作之前，先来了解一下 Muon 是个什么样的技术。
这是一种神经网络隐藏层的 2D 参数优化器，主要作者是 OpenAI 深度学习团队的 Keller Jordan。
这项成果发表于去年的 12 月 8 日，而 Keller 也是去年 12 月加入的 OpenAI。
Muon 的核心思想是通过正交化梯度更新矩阵，避免参数更新陷入局部极小，使模型能够学习到更加多样化的特征表示。
在 94% 的精度下，Muon 把 CIFAR-10 在 A100 上的训练时间从 3.3 秒缩短至 2.6 秒。
不过当时 Muon 团队只证明了其在小型模型和数据集上的可行性，对于较大的模型能否适用则是个未知数。
现在经过月暗团队的改进之后，Muon被证明对于更大的模型和数据集同样适用。
针对模型本身，团队吸收了 AdamW 中的一些特点，移植到了 Muon 当中，具体包括两个方面。
一是引入了权重衰减机制，在权重更新公式中添加了一个带有衰减系数的项。
这样做的原因是作者发现直接将 Muon 应用到大规模训练时，模型权重和层输出的幅度会持续增长，最终超出 bf16 的高精度表示范围，损害模型性能。
在训练一个 8 亿参数模型至 100B tokens（约 5 倍计算预算最优）的过程中，团队对比了 AdamW、无权重衰减的 Muon 和带权重衰减的 Muon。
结果显示，带权重衰减的 Muon 在过拟合阶段取得了最佳效果，验证了权重衰减的必要性。
第二项改进，是调整了 Muon 的参数更新尺度，使不同形状矩阵参数的更新幅度保持一致，并与 AdamW 的更新幅度匹配。
Muon 的一个特点是，对于形状为 [ A,B ] 的矩阵参数，其理论更新幅度为 sqrt ( 1/max ( A,B ) ) 。
这导致不同形状矩阵参数的更新幅度差异很大，比如对于 MLP 这种宽矩阵，更新会过小，而将每个 head 看作独立矩阵时，更新又会过大。
此外，这个幅度也与 AdamW 不一致，给超参数的設定带来困难。
为了让不同矩阵参数的更新幅度匹配，并与 AdamW 保持一致，作者尝试了几种改进方案，最终选择直接基于形状调整每个参数的学习率。
其中 0.2 是通过实验确定的一个常数，用于将 Muon 的更新尺度与 AdamW 对齐。
除了对 Muon 本身的改进，要想将 Muon 用于更大规模的训练，还需要将其扩展到分布式训练环境中。
由于 Muon 需要完整的梯度矩阵来计算正交化的更新量，而现有的分布式训练框架（如 ZeRO-1、Megatron-LM 等）都假设优化器状态可以独立地按元素切分到不同设备上，所以它们无法直接支持 Muon。
为了解决这个问题，论文作者提出了分布式 Muon 的并行化策略。
它在 ZeRO-1 的基础上引入了两个额外的操作 :
一是在每个数据并行组内做梯度聚合通信，将分散的梯度切片合并成完整的矩阵；
二是基于聚合后的梯度矩阵并行计算正交化的更新量，然后只保留与本地参数对应的那一部分。
这种实现方式在最小化内存占用和通信开销的同时，最大限度地保留了原始 Muon 算法的数学性质。
证明 Muon 扩展可行性
基于上述 Muon 改进 , 作者取得了以下成果，作者在 Llama 架构的一系列稠密模型上，进行了 Muon 和 AdamW 的模型缩放对比实验。
结果表明 , 在计算预算最优的情况下，Muon 的样本效率是 AdamW 的 1.92 倍，即训练 FLOPS 只需 AdamW 的 52%，就能达到相当的性能。
这一发现证实了 Muon 在大规模训练中的效率优势。
在此基础之上，作者以 DeepSeek-V3-Small 架构作为基础，用改进的 Muon 训练了Moonlight模型。
Moonlight 是一个 MoE 模型，具有 15.29B 的总参数和 2.24B 激活参数，训练 token 量为 5.7T。
与相同规模和数据量的模型相比，Moonlight 在英语理解与推理（MMLU、TriviaQA、BBH）、代码生成（HumanEval、MBPP）、数学推理（GSM8K、MATH、CMATH）、中文理解（C-Eval、CMMLU）等各类任务上都取得了明显更好的性能。
即使与使用更大数据集训练的稠密模型相比，Moonlight 也展现了极强的竞争力。
与多个知名语言模型的对比表明，Moonlight在性能 - 训练预算平面上推进了帕累托前沿（Pareto Frontier）。
（注：帕累托前沿是一个经济学和管理学中的概念，描述的是在多目标决策问题中所有可能的最优解的集合，这些解在多个目标之间取得了最佳平衡。在帕累托前沿上的每一个点，都意味着一个目标的改善必然以牺牲另一个目标为代价，因此它代表了在多个目标之间实现的最佳权衡。）
为了进一步分析 Muon 更新矩阵参数的内在机制，作者对比了 Muon 和 AdamW 训练得到的模型在不同训练阶段的参数矩阵奇异值谱。
结果发现，Muon 优化的矩阵在各层各类参数上，总是比 AdamW 有更高的奇异值熵。这从经验上验证了 Muon 通过正交化来学习更多样化表示的直觉。
最后，在 Moonlight 模型的基础上，作者还探索了 Muon 在指导微调阶段的效果，结果表明，在预训练和微调阶段均使用 Muon 的效果是最佳的。
技术报告：
https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
Code：
https://github.com/MoonshotAI/Moonlight
Moonlight 模型：
https://huggingface.co/moonshotai/Moonlight-16B-A3B
