Delving into Muon and Beyond: Deep Analysis and Extensions

技术摘要：深入探讨 Muon 及其延伸：深度分析与扩展

问题陈述
Muon 优化器因其在训练大型语言模型时利用矩阵级正交化更新而展现出的经验性能，近期受到了广泛关注。然而，Muon 的底层机制及其与 Adam 等既有自适应优化器的关系仍未得到充分理解。现有文献往往将 Muon 的效果与辅助技术（如 QK-Norm、QK-Clip）混为一谈，或缺乏与标准基准的系统性比较。此外，目前尚不清楚 Muon 的谱正交化（spectral orthogonalization）究竟是一种普遍优越的优化策略，还是一种特定的稳定化手段，以及它是否可以被改进或泛化。

方法论
作者提出了一个统一的谱框架来分析 Muon 及其变体。他们将 Muon 定义为形式为 $U\Sigma^p V^\top$ 的一系列谱变换中的 $p=0$ 端点，其中 $U\Sigma V^\top$ 是更新矩阵的奇异值分解（SVD）。该研究研究了中间指数 $p \in \{1, 1/2, 1/4, 0\}$：

• $p=1$：恒等变换（标准更新）。
• $p=1/2, 1/4$：部分谱压缩。
• $p=0$：完全正交化（Muon）。

这些变换应用于两种不同的输入类型：

1. 一阶矩动量 ($O_{mom}$)：等效于带动量的 SGD 输入。
2. RMS 归一化更新 ($O_{rms}$)：等效于 Adam 输入，即梯度通过二阶矩进行归一化。

为了在不使用显式 SVD 的情况下高效计算分数幂（$p=1/2, 1/4$），作者开发了一种耦合牛顿-舒尔茨迭代（coupled Newton-Schulz iteration）。该算法仅通过矩阵乘法同时计算矩阵平方根及其逆，确保了数值稳定性和可扩展性。

实验评估在 OpenWebText 数据集上训练的 124M 参数 GPT-2 模型上进行。研究采用了高度受控的设置：

• 解耦学习率：对矩阵参数和向量参数使用独立的学习率，以隔离优化器的影响。
• 无辅助技术：禁用权重衰减（weight decay）、QK-Norm 和 QK-Clip，以确保公平比较。
• 系统化调优：对所有八种优化器变体（4 种谱指数 $\times$ 2 种输入类型）进行了两阶段学习率搜索（先进行粗粒度对数扫描，随后进行细粒度精细化调整）。

核心贡献

1. 统一谱框架：论文引入了一个形式化框架 $U\Sigma^p V^\top$，将 Muon 定位为特定的情况（$p=0$），并提出了中间变体（$p=1/2, 1/4$）以在标准更新与完全正交化之间进行插值。
2. 高效算法：开发了耦合牛顿-舒尔茨迭代，用于高效计算分数谱更新（$p=1/2, 1/4$），避免了显式 SVD 的计算开销。
3. 严谨的经验比较：研究提供了 Muon、Adam 及其谱变体之间的受控比较，将谱效应从权重衰减或辅助归一化等干扰因素中分离出来。

关键结果
实验得出了三个主要发现：

1. 一阶矩更新的稳定性：在处理一阶矩输入（mSGD）时，Muon ($p=0$) 提供了显著的稳定性，允许使用更高的学习率并具有更强的鲁棒性，优于标准 mSGD。然而，一旦稳定性建立，部分压缩（$p=1/4$）通常能获得更好的最终性能。
2. 在 RMS 归一化输入上的收益有限：当应用于 RMS 归一化输入（Adam 式）时，谱压缩带来的收益递减。应用于 $O_{rms}$ 的完全正交化（AdamZ/Muon）并不能超越标准 Adam，甚至可能逊于部分压缩。这表明，元素级的 RMS 归一化已经缓解了谱压缩所针对的大部分各向异性问题。
3. 二阶矩方法的优越性：在所有谱指数下，利用二阶矩归一化的优化器（基于 RMS）在稳定性及最终损失方面均一致优于其一阶矩（仅动量）对应的优化器。

意义与主张
论文得出结论，Muon 最好的理解方式是作为一种有效的谱归一化形式，而非一种普遍优越的优化方法。

• 稳定性 vs. 性能：Muon 的主要贡献在于通过丢弃奇异值的大小来稳定一阶矩更新，从而防止具有大奇异值的方向占据主导地位。然而，这种“平坦化”是一种粗糙的机制；它移除了基于幅度的滤波，可能会放大噪声梯度方向（小奇异值），并在信息丰富的优势模态上步长不足。
• 上下文效用：虽然谱压缩提供了真实的稳定性收益（尤其是在输入更新尚未归一化的情况下），但完全正交化 ($p=0$) 并不是现代基于 RMS 的二阶矩优化器（如 Adam）的替代品。作者认为，此前工作中观察到的 Muon 优势可能归功于特定技术的结合（如 QK-Norm），而非正交化本身。

该工作强调，尽管谱变换是控制各向异性的强大工具，但它们本质上并不能超越在大规模训练体系中自适应二阶矩缩放的优势。