

2025
推进Muon优化器
月之暗面团队
开启
AI训练
新篇章


在人工智能领域,训练优化算法的每一次突破都可能带来巨大的变革。OpenAI技术人员提出的训练优化算法Muon,自诞生以来就备受关注。而最近,月之暗面团队的一项研究成果,再次将Muon推到了聚光灯下,为其在更广泛领域的应用奠定了坚实基础。
Muon作为一种神经网络隐藏层的2D参数优化器,其核心思想是通过正交化梯度更新矩阵,避免参数更新陷入局部极小,从而使模型能够学习到更加多样化的特征表示。在其最初的成果展示中,Muon在94%的精度下,将CIFAR-10在A100上的训练时间从3.3秒缩短至2.6秒,展现出了强大的潜力。然而,当时Muon仅在小型模型和数据集上得到了验证,对于大型模型和大规模训练的适用性尚不明朗。
月之暗面团队的研究正是针对这些未解决的问题展开的。他们首先发现了Muon方法的Scaling Law,并对其进行了改进。团队吸收了AdamW中的一些特点,引入了权重衰减机制,解决了模型权重和层输出幅度增长超出高精度表示范围的问题,验证了权重衰减在大规模训练中的必要性。同时,调整了Muon的参数更新尺度,使不同形状矩阵参数的更新幅度保持一致,并与AdamW的更新幅度匹配,解决了超参数设置的难题。
在分布式训练方面,月之暗面团队也取得了重要进展。由于现有的分布式训练框架无法直接支持Muon,团队提出了分布式Muon的并行化策略,在ZeRO-1的基础上引入梯度聚合通信和基于聚合梯度矩阵的正交化更新量计算,既最小化了内存占用和通信开销,又保留了原始Muon算法的数学性质。
通过一系列实验,月之暗面团队证明了改进后的Muon在大规模训练中的效率优势。在Llama架构的稠密模型上,Muon的样本效率是AdamW的1.92倍,训练FLOPS只需AdamW的52%就能达到相当的性能。基于DeepSeek-V3-Small架构训练的Moonlight模型,更是在各类任务上取得了优于相同规模和数据量模型的性能,甚至与使用更大数据集训练的稠密模型相比也极具竞争力。
此外,团队还对Muon更新矩阵参数的内在机制进行了深入分析,对比发现Muon优化的矩阵在各层各类参数上具有更高的奇异值熵,验证了其通过正交化学习更多样化表示的直觉。在微调阶段的实验也表明,预训练和微调均使用Muon效果最佳。
月之暗面团队的这一成果,不仅为Muon在更大规模训练中的应用提供了有力支持,也为AI训练优化算法的发展开辟了新的方向。随着技术的不断进步,我们有理由相信,Muon将在未来的AI领域发挥更加重要的作用,推动人工智能技术迈向新的高度。
【END】


