【论文解读】Kimi开源《Muon is Scalable for LLM Training》
Github：https://github.com/MoonshotAI/Moonlight
HF：https://huggingface.co/moonshotai/Moonlight-16B-A3B
Paper：https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
1. 摘要
-
背景与动机
随着大语言模型（LLM）的迅速发展，训练这些模型所需的计算资源和时间呈现指数级增长。当前大多数训练任务使用 Adam/AdamW 优化器，但最近的研究（例如 K. Jordan et al. 2024）提出了基于矩阵正交化思想的 Muon 优化器，在小规模语言模型上展现出较好的效果。
-
核心贡献
论文通过两个关键改进使 Muon 优化器可以扩展到大规模 LLM 训练：- 加入权重衰减：解决在大规模训练中模型权重和输出均值过大、超出 bf16 表示范围的问题；
- 调整每个参数的更新尺度：通过理论推导和实验验证，提出了基于矩阵尺寸（即
max
(
A
,
B
)
\sqrt{\max(A, B)}
max(A,B)）的更新尺度调整策略，使得不同形状的矩阵更新保持一致，并与 AdamW 的更新 RMS 相匹配。
实验表明，在计算最优训练设置下，Muon 优化器在达到相同性能时只需要约 52% 的训练 FLOPs。
-
实验证明与开源
基于上述改进，论文引入了 Moonlight 模型（一个 3B/16B 参数的 MoE 模型，使用 5.7T tokens 进行训练），在多个基准（如 MMLU、HumanEval、GSM8K 等）上取得了优于 AdamW 的结果，同时也开源了分布式 Muon 实现、预训练模型、指令调优版本以及中间检查点，方便后续研究。
2. 引言
-
大模型训练的挑战
随着 LLM 在自然语言理解、代码生成和数学推理等任务上的广泛应用，训练这些模型变得极其计算密集。现有优化器如 AdamW 在训练过程中虽较为稳定，但在训练极大模型时效率和稳定性依然面临瓶颈。 -
Muon 优化器的提出
Muon 优化器最初由 K. Jordan 等人提出，利用矩阵正交化的思想，通过对梯度动量进行正交化处理，避免参数更新过于集中于某几个主方向，从而提升优化效果。论文中指出：- 小规模实验中，Muon 已表现出超越 AdamW 的潜力；
- 但在大规模训练中，直接使用 Muon 会遇到数值范围、更新尺度不一致以及分布式实现等一系列问题。
-
论文主要问题
论文重点讨论了以下几个问题：- 如何将基于矩阵正交化的优化器有效扩展到拥有数十亿参数、数万亿 token 训练的场景；
- 如何在分布式环境下高效计算近似正交化；
- 这种优化器是否能在预训练和监督微调（SFT）等不同阶段保持一致的性能优势。
3. 方法
3.1 Muon 优化器的基本原理
-
更新公式
Muon 在每次迭代中以动量方式累积梯度，然后利用 Newton-Schulz 迭代对梯度矩阵进行正交化：M t = μ M t − 1 + ∇ L ( W t − 1 ) M_t = \mu M_{t-1} + \nabla L(W_{t-1}) Mt=μMt−1+∇L(Wt−1)
接着采用 Newton-Schulz 迭代求解：
O t = Newton-Schulz ( M t ) ≈ ( M t M t T ) − 1 / 2 M t O_t = \text{Newton-Schulz}(M_t) \approx (M_t M_t^T)^{-1/2} M_t Ot=Newton-Schulz(Mt)≈(MtMtT)−1/2Mt
最后更新权重：
W t = W t − 1 − η t O t W_t = W_{t-1} - \eta_t O_t Wt=Wt−1−ηtOt
-
直观解释
正交化确保了更新矩阵的各个方向具有等同性，从而避免模型仅在某些主方向上快速收敛，导致表示能力受限。
3.2 Newton-Schulz 迭代
-
迭代过程
为了计算 ( M t M t T ) − 1 / 2 M t (M_t M_t^T)^{-1/2} M_t (MtMtT)−1/2Mt ，论文使用 Newton-Schulz 迭代。初始设 X 0 = M t ∥ M t ∥ F X_0 = \frac{M_t}{\|M_t\|_F} X0=∥Mt∥FMt，然后通过迭代：X k = a X k − 1 + b ( X k − 1 X k − 1 T ) X k − 1 + c ( X k − 1 X k − 1 T ) 2 X k − 1 X_k = a X_{k-1} + b (X_{k-1} X_{k-1}^T) X_{k-1} + c (X_{k-1} X_{k-1}^T)^2 X_{k-1} Xk=aXk−1+b(Xk−1Xk−1T)Xk−1+c(Xk−1Xk−1T)2Xk−1
当迭代 N N N 次后， X N X_N XN 近似得到所需的正交化结果。论文中给出了特定的系数设置（例如 a = 3.4445 , b = − 4.7750 , c = 2.0315 a = 3.4445, \, b = -4.7750, \, c = 2.0315 a=3.4445,b=−4.7750,c=2.0315），以保证快速收敛。
3.3 权重衰减与更新尺度的调整
-
权重衰减的重要性
在大规模训练中，若不引入权重衰减，模型中的权重和层输出的均方根值（RMS）会不断增长，可能超出 bf16 的数值表示范围，进而影响模型性能。为此，论文将 AdamW 的权重衰减机制引入 Muon，更新公式变为：W t = W t − 1 − η t ( O t + λ W t − 1 ) W_t = W_{t-1} - \eta_t \Bigl( O_t + \lambda W_{t-1} \Bigr) Wt=Wt−1−ηt(Ot+λWt−1)
-
参数形状与更新 RMS
论文提出一个重要定理：对于形状为 [ A , B ] [A, B] [A,B] 的全秩矩阵，其理论更新 RMS 为RMS ≈ 1 max ( A , B ) \text{RMS} \approx \sqrt{\frac{1}{\max(A, B)}} RMS≈max(A,B)1
当不同层的矩阵尺寸差异较大时，更新尺度不一致会影响整体训练。因此提出两种解决方案：
- 直接归一化更新：将更新矩阵 O t O_t Ot 归一化使得其 RMS 达到预期值（例如 0.2）；
- 调整学习率：对每个矩阵的更新乘以 max ( A , B ) \sqrt{\max(A, B)} max(A,B)，使得更新 RMS 与 AdamW 的更新 RMS 0.2 ⋅ H 0.2 \cdot \sqrt{H} 0.2⋅H（H 为隐藏层大小）相匹配。
-
结果与分析
实验表明：- 直接归一化和调整学习率均优于 Baseline；
- 对于形状较大（如 MLP 矩阵）的层，调整学习率策略使得更新 RMS 提升至 Baseline 的约两倍，而对于形状为
[
H
,
H
]
[H, H]
[H,H] 的层，二者保持一致。
最终选择调整学习率策略，因其实现成本更低且性能表现良好。
4. 实验部分
4.1 更新 RMS 一致性实验
-
实验设置
在一个 800M 参数模型上，使用 100B tokens 进行训练。论文对比了三种策略：- Baseline：直接将更新矩阵乘以 0.2 ⋅ H 0.2 \cdot \sqrt{H} 0.2⋅H（H 为隐藏层大小）；
- Update Norm：直接归一化更新矩阵，使其 RMS 固定为 0.2；
- Adjusted LR：按每个矩阵的形状调整学习率，即乘以 (0.2 \cdot \sqrt{\max(A, B)})。
-
结果与分析
实验表明：- Update Norm 和 Adjusted LR 均优于 Baseline；
- 对于形状较大（如 MLP 矩阵）的层，Adjusted LR 使得更新 RMS 提升至 Baseline 的约两倍，而对于形状为 ([H, H]) 的层，二者保持一致。
最终选择 Adjusted LR 策略，因其实现成本更低且性能表现良好。
4.2 扩展定律实验
-
实验目标
在 Llama 架构的多个模型上对比 Muon 与 AdamW 的扩展定律，确保在计算最优训练设置下，Muon 能够在更低 FLOPs 下达到相似性能。 -
主要发现
根据拟合的扩展定律曲线（见图 1a 和图 3），Muon 在达到与 AdamW 相同性能时，只需大约 52% 的训练 FLOPs，证明了其计算效率优势。
4.3 预训练实验：Moonlight 模型
-
模型架构与训练设置
使用 DeepSeek-V3-Small 作为基础架构，对模型进行预训练，模型参数约为 2.24B（激活参数）/15.29B（总参数，含 embedding 部分），后续 MoE 模型扩展至 3B/16B。预训练数据包括长上下文（8K）数据，训练分阶段进行：- 初始阶段（0 至 33B tokens）：学习率从 0 线性上升到 4.2 × 1 0 − 4 4.2 \times 10^{-4} 4.2×10−4，batch size 固定；
- 衰减阶段（33B 至 5.2T tokens）：学习率按余弦衰减；
- 冷却阶段（5.2T 至 5.7T tokens）：先短暂上升再线性衰减至 0，同时使用高质量数据集（数学、代码、推理）。
4.4 奇异谱动态分析
-
目的
通过对权重矩阵进行奇异值分解（SVD），计算 SVD 熵来评估参数更新的多样性。 -
发现
与 AdamW 相比，使用 Muon 优化器的权重矩阵在各训练阶段均表现出更高的 SVD 熵，说明 Muon 能够提供更丰富的更新方向，这一现象在 MoE 模型的路由权重中尤为明显。
4.5 监督微调（SFT）实验
-
实验设置
分别采用相同预训练模型（Moonlight 或 Moonlight-A），在开放的 tulu-3-sft-mixture 数据集上进行微调，比较预训练和微调阶段优化器的组合：- 预训练和 SFT 均使用 Muon；
- 预训练使用 Muon 而 SFT 使用 AdamW，反之亦然。
5. 讨论
-
全参数整合问题
当前 Muon 与 AdamW 混用，部分参数（如 RMSNorm、Embedding）仍使用 AdamW。未来的研究可以探讨如何将所有参数完全纳入 Muon 框架中。 -
扩展到 Schatten 范数
Muon 本质上是在谱范数下进行最速下降，考虑到 Schatten 范数的广泛应用，将 Muon 扩展到一般 Schatten 范数下可能获得更强的优化能力。 -
预训练-微调不匹配问题
论文中观察到预训练和微调阶段使用不同优化器时，模型性能不如统一使用 Muon 的情况。理解这种不匹配的根本原因，并提出解决方案，是未来研究的重要方向。
6. 结论
- 本文系统地探讨了如何将基于矩阵正交化的 Muon 优化器扩展到大规模 LLM 训练中。
- 通过引入权重衰减、调整每个参数的更新尺度，并设计分布式实现方案，Muon 成功在预训练和监督微调阶段均展示了优异性能。
- 实验结果表明，在计算资源上，Muon 具有大约 2 倍的效率提升，同时在多个评测任务上超越了基于 AdamW 的模型。
- 论文的开源工作（实现、模型和检查点）为后续在大规模优化器设计及训练方法改进提供了宝贵资源。
7. 附录与参考文献
论文中还附有以下补充内容：
- 附录 A： 关于 Lemma 1 的证明，展示了对于形状为 [ A , B ] [A, B] [A,B] 的矩阵，其 Muon 更新 RMS 理论上为 1 / max ( A , B ) \sqrt{1/\max(A,B)} 1/max(A,B)；
- 附录 B： 系统的超参数搜索实验细节，包括模型大小、token 数、学习率和 batch size 与 FLOPs 之间的关系；
- 附录 C： 模型架构的修改细节（例如去除多 token 预测层、auxfree bias 更新规则的调整、门控缩放因子计算等）。