[arXiv 2025] MUON+：多一步归一化，大模型预训练效率再突破

TL;DR

Transformer 预训练优化器迎来了“极简主义”的又一次胜利。加州大学圣塔芭芭拉分校（UCSB）的研究团队通过在 Muon 优化器中添加一个极其简单的**更新方向归一化（Normalization）**步骤，提出了 MUON+。该方法在不显著增加计算开销的前提下，在 GPT 和 LLaMA 架构上全面超越了 Muon 和 AdamW，尤其在长周期训练中展现了极强的稳定性。

背景定位：正交优化的兴起

在 LLM 预训练领域，AdamW 长期占据统治地位。然而，近年来的研究（如 Muon）发现，对动量矩阵进行**正交化（Orthogonalization）**可以有效对抗梯度秩崩塌（Rank Collapse），从而加速收敛。Muon 已经成为 Kimi、GLM 等主流模型预训练中的重要技术组件。

痛点与动机：复杂背后的简单真相

尽管 Muon 表现优异，但研究者们在后续的改进（如 NorMuon, Mano）中引入了复杂的神经元自适应缩放或流形投影。 作者的直觉（Insight）是：这些复杂改进带来的增益，是否仅仅源于其中的归一化操作？ 通过消融实验，作者发现，去掉复杂的二阶矩估计（Second-moment），仅保留归一化，性能几乎没有损失。这促使了 MUON+ 的诞生——直接回归本质，给正交后的矩阵再套一个 L2 Normalization。

方法论详解：MUON+ 的核心架构

MUON+ 的算法流程非常直观，其核心公式如下：

1. 动量更新：$M_t=\mu M_{t-1}+(1-\mu )G_t$
2. 正交化：$O_t=extOrtho(M_t)$（通常使用 Newton-Schulz 迭代近似 SVD）
3. 归一化（MUON+ 核心）：$O_t=ext{Norm}_{(d)}(O_t)$
4. 权重更新：$W_t=W_t-\eta \cdot \sqrt{m/n}\cdot O_t$

归一化方向的选择

作者探讨了四种方向：行归一化 (Row)、列归一化 (Col)、以及它们的组合 (Col-Row / Row-Col)。实验发现：

• 行归一化效果普遍优于列归一化。
• **混合归一化（Col-Row/Row-Col）**表现最为稳健，能够对更新矩阵施加更强的结构约束。

实验与结果：全线 SOTA

作者在 GPT 和 LLaMA 两种主流架构上，覆盖了从 60M 到 1.3B 的参数规模进行了广泛实验。

1. Compute-Optimal 预训练

在符合缩放法则（Scaling Laws）的计算最优设置下，MUON+ 在所有规模上均优于原始 Muon。

2. 极致的训练稳定性（Overtraining）

在 T2P 比率达到 200 的“过度训练”场景下，MUON+ 的优势不仅没有随训练进行而缩小，反而保持了稳定的领先差距。这意味着它不仅适用于快速实验，更适用于生产环境中的大规模长周期预训练。

3. 对学习率的鲁棒性

消融实验揭示，MUON+ 显著拓宽了最优学习率的搜索区间。即便使用了过大的学习率，MUON+ 带来的性能衰减也远小于 Muon，大大降低了超参数调优的难度。

深度洞察与总结

MUON+ 为什么有效？

从几何角度看，Muon 的正交化确保了更新矩阵的奇异值全部为 1，而 MUON+ 的归一化则进一步约束了矩阵每一行或每一列的 L2 范数。这种“双重约束”起到了强效的正则化作用，使每一层权重的更新步长在结构上更加均衡，从而带来了更好的优化动力学特征。

局限性与展望

目前实验主要集中在 1.3B 以下规模。虽然已有研究证明 Muon 在千亿级模型中是可扩展的，但 MUON+ 在更大规模（如 70B+）上的数值稳定性仍需进一步在大规模算力集群上验证。

总结：MUON+ 证明了在深度学习优化中，简单的结构约束往往优于复杂的启发式算法。它为在大规模语言模型训练中进一步压榨算力效率提供了新的基准。