Muon：加速神经网络训练的开源新工具

Muon：加速神经网络训练的新工具

神经网络需要有效地优化其“隐藏层”，才能运行得更快、更高效。一种名为 Muon 的新型优化器正是为此而设计的，它可以帮助隐藏层权重比 AdamW 等标准方法更有效地学习。让我们来详细分析一下它的工作原理及其重要性。

什么是 Muon？

Muon 是一个专门针对神经网络“隐藏层”的优化器——可以将其视为处理数据的核心部分，例如卷积滤波器或 Transformer 注意力权重。它将其他参数（例如嵌入、分类器头或偏差）留给标准优化器处理，专注于对速度最关键的部分。

如何使用 Muon

使用 Muon 很简单，但您需要将模型拆分成几个部分。以下是分步示例：

1. 将您的模型拆分成三组：

• hidden_​​weights：Muon 将优化的核心层（例如，卷积滤波器、Transformer 层）。
• hidden_​​gains_biases：较小的参数，例如偏差或缩放因子，使用 AdamW 进行优化。
• nonhidden_​​params：头（分类器）、嵌入或其他“外围”层，也使用 AdamW。

2. 设置优化器：
   Muon 可与一个包装器 MuonWithAuxAdam 配合使用，该包装器将 Muon 用于隐藏层，并将 AdamW 用于其他层。请根据模型大小调整学习率和权重衰减（Muon 在较大的模型中具有良好的扩展性，因此您无需经常调整学习率）。

为何更胜一筹

像 AdamW 这样的传统优化器有时会难以应对深度网络中“混乱”的更新，尤其是在隐藏层包含大量参数的情况下。Muon 通过以下方式解决了这个问题：

• 专注于优化最关键的部分（隐藏层权重），并采用专门的更新规则。
• 能够根据模型规模进行良好扩展——在训练更大规模的网络时，无需重新调整学习率。

实际结果

Muon 已在基准测试中展现出令人印象深刻的性能：

• 在 CIFAR-10 上训练速度更快：达到 94% 准确率所需的时间从 3.3 A100 秒缩短至 2.6 A100 秒（速度提升 20%）。
• 经济高效的 LLM 训练：仅需 175 美元的计算成本，即可将 Transformer 训练到 GPT-2 (XL) 级别。
• 支持大批量：在大批量情况下也能出色地工作，而大批量是训练效率的常见挑战。

关键要点

对于任何训练大型神经网络（无论是 Transformer、ConvNet 还是 LLM）的人来说，Muon 都是一款极具潜力的工具。它简化了关键隐藏层的优化，加快了训练速度，并能随着模型规模的扩大而扩展。如果您希望突破训练速度的极限或降低成本，Muon 值得添加到您的工具包中。

阅后请思考

• Muon优化的具体机制是什么？
• AdamW与Muon结合有何优势？
• MUON在其他数据集表现如何？