Muon Optimizer：梯度白化与 SVD

引言：为什么需要 Muon？

Muon 是 K2 训练背后的关键优化器，由 Keller Jordan 提出。它相比 AdamW 或带动量的 SGD 具有更快的训练效率：达到相同 loss 所需的数据量或算力更少。Keller Jordan 正是凭借这篇工作拿到了 OpenAI 的 offer，可见业界对该优化器的高度认可。

Muon 的核心洞察

在用 SGD 或 AdamW 优化神经网络时，经验观察发现：大部分 2D 参数矩阵（隐藏层权重）的条件数很高，即矩阵是低秩的。这导致梯度更新在少数方向上幅度很大，而在大量“稀有方向”上信号极弱。

Muon 的核心：梯度白化

对梯度矩阵进行白化处理，抹去梯度的大小/尺度，只保留方向。这使得：

白化操作通过 SVD（奇异值分解）实现：

\[ G = U \Sigma V^T \quad \xrightarrow{\text{白化}} \quad \hat{G} = U V^T \]

即去掉奇异值矩阵 \(\Sigma\)，只保留旋转部分 \(U\) 和 \(V^T\)。

对于梯度矩阵 \(G\)：

去掉 \(\Sigma\) 后，\(\hat{G} = UV^T\) 是一个正交矩阵，保证了所有方向的更新幅度一致。

梯度白化是 Muon 的核心操作：通过 SVD 去掉奇异值，将梯度矩阵正交化。这使得原本被主导方向压制的稀有信号得以放大，加速了模型学习。

Muon 不是全局替代品

Muon 只用于优化隐藏层的 2D 参数矩阵。以下参数仍使用标准优化器（AdamW 或带动量的 SGD）：

实际训练中，Muon 与 AdamW/SGD 组合使用：隐藏层矩阵走 Muon，其余走标准优化器。

对于输出层（分类层），我们希望学到的不同特征对最终分类贡献不同的权重。如果对输出层也做白化，会抹掉这种有意义的尺度差异。因此输出层不适用 Muon。

Muon 的适用范围严格限定在隐藏层的 2D 矩阵参数上，与标准优化器互补使用。

矩阵的条件数定义为最大奇异值与最小奇异值之比：

\[ \kappa(A) = \frac{\sigma_{\max}}{\sigma_{\min}} \]

条件数越大，矩阵越“病态”（ill-conditioned），意味着信息集中在少数方向上。

高条件数导致：

Muon 通过白化将所有方向的“重要程度”拉平，从而加速收敛。

高条件数是神经网络参数矩阵的普遍特征，它导致标准优化器在稀有方向上收敛缓慢。Muon 的白化操作本质上是在解决这个问题。