Muon is Scalable for LLM Training

最近，基于矩阵正交化的Muon优化器在训练小型语言模型方面展现了强大的效果，但其扩展到更大模型的能力尚未得到证实。我们确定了两种关键的技术来扩大Muon的应用规模：（1）添加权重衰减和（2）仔细调整每个参数的更新比例。这些技术使得Muon可以在大规模训练中直接使用，而无需进行超参数调整。根据扩展定律实验的结果，与采用计算最优训练的AdamW相比，Muon实现了大约两倍的计算效率。 基于这些改进，我们推出了Moonlight，这是一个拥有30亿或160亿参数的专家混合（MoE）模型，使用Muon进行了5.7万亿个令牌的训练。我们的模型改进了当前的帕累托前沿，在训练FLOPs显著减少的情况下实现了更好的性能，超越了之前的模型。 我们开源了我们的分布式Muon实现，该实现具有内存最优和通信高效的特点。我们还发布了预训练、指令调优以及中间检查点，以支持未来的研究。

该论文试图解决如何将Muon优化器从训练小型语言模型扩展到大型语言模型的问题，并验证通过特定技术改进，Muon优化器可以在大规模训练中实现更高的计算效率。这是一个相对较新的问题，因为之前的研究主要集中在小型模型上。

关键思路在于引入两个核心技术来提升Muon优化器的可扩展性：1) 添加权重衰减（weight decay），2) 精心调整每个参数的更新比例。这些改进使得Muon优化器无需额外的超参数调优即可应用于大规模训练。此外，作者还提出了Moonlight模型，一个基于Mixture-of-Expert架构、参数量为3B/16B的大规模语言模型，展示了更好的性能和更低的训练FLOPs需求。

论文的亮点包括：1) Muon优化器在大规模训练中的显著效率提升（相比AdamW提高了约2倍）；2) Moonlight模型不仅在性能上优于现有模型，而且训练所需的计算资源更少；3) 研究团队开源了分布式Muon实现，该实现内存优化且通信高效；4) 提供了预训练、指令调优和中间检查点，支持未来研究。

最近在这个领域中，相关研究还包括：1) 使用其他优化器（如LAMB、AdaFactor）进行大规模语言模型训练的工作；2) 针对MoE架构的优化与扩展研究；3) 探索不同优化策略对大规模模型训练效率的影响。一些相关的论文标题包括《Scaling Memory-Efficient Optimizers to Trillion Parameter Models》、《Efficient Training of BERT with Large Batch Sizes》等。