Prologue
最近有一些工作认为使用行归一化对梯度矩阵做 Precondition，可以替代 Muon 的正交化达到接近甚至超越的结果。例如前段时间 Yiping 做的Row/Column Normalization and Hyperparameter Transfer。虽然论文给了一些有趣的理论结果，但是当时我对他们的实验 setting 是有一些质疑的。最近又有新的工作RMNP，声称使用行归一化达到了超越 Muon 的结果。
事实果真如此吗？按照笔者的经验，行归一化的效果其实跟 Adam 更接近，没有超越 Muon。本篇文章以 RMNP 为例来讨论下。
RMNP
假设梯度的 Momentum 矩阵是 $V_t$，Muon 使用：
$$ \Delta W\propto (V_tV_t^\top)^{-1/2}V_t. $$
而 RMNP 的作者认为，因为 $V_tV_t^\top$ 大概率是一个对角主导的矩阵（对角线元素远大于非对角线元素），因此可以用如下的近似：
$$ \Delta W\propto \operatorname{diag}(V_tV_t^\top)^{-1/2}V_t. $$
这等价于对 Momentum 矩阵每一行单独归一化。
这个假设不难理解，实际上就是非常经典的维度灾难的问题，我对作者的这个假设先按下不表。看看作者的实验效果如何：
这里有个很有意思的结论，作者从“行归一化与正交化在高维下渐进等价”这个角度出发，用行归一化取代 Muon，取得了超越 Muon 的结果。这就好像「在模仿吴亦凡的比赛中吴亦凡获得了第二名」，听起来令人匪夷所思。
渐进等价的逻辑漏洞
RMNP 的核心逻辑是：因为
$$ V_tV_t^\top \approx \operatorname{diag}(V_tV_t^\top), $$
所以可以用后者替代前者，效果差不多。这里其实有一个比较模型的漏洞，即便对称矩阵 $A$ 如果在 rms norm 的意义下近似于一个对角矩阵 $D$，使用 $A^{-1/2}$ 和 $D^{-1/2}$ 作为 preconditioner 的效果却可能差之千里。
我们可以构造一个简单的例子来说明这个问题，令
$$ u=\frac{1}{\sqrt m}(1,1,\ldots,1)^\top, $$
并定义
$$ A_{\epsilon,m}=I-(1-\epsilon)uu^\top, \qquad 0<\epsilon<1. $$
这是正定矩阵。它的谱非常简单：
$$ A_{\epsilon,m}u=\epsilon u, \qquad A_{\epsilon,m}x=x\quad \text{for }x\perp u. $$
从 entrywise / RMS 角度看，$A_{\epsilon,m}$ 非常接近 diagonal。因为
$$ \begin{align} (A_{\epsilon,m})_{ii}&=1-\frac{1-\epsilon}{m},\\ (A_{\epsilon,m})_{ij}&=-\frac{1-\epsilon}{m},\quad\forall j\neq i. \end{align} $$
所以 off-diagonal entry 的大小只有 $O(1/m)$。如果用作者那种 average off-diagonal magnitude 的视角来看，每一行的 off-diagonal 平均量级是
$$ \frac{1-\epsilon}{m}, $$
而 diagonal element 接近 $1$。于是 diagonal dominance ratio 大约是
$$ r_i\approx \frac{m}{1-\epsilon}, $$
随 $m$ 增大而发散。按这个指标看，矩阵越来越 diagonal dominant。
但是 inverse square root 呢？沿 $u$ 方向：
$$ A_{\epsilon,m}^{-1/2}u=\epsilon^{-1/2}u. $$
而 diagonal approximation 大约给出
$$ \operatorname{diag}(A_{\epsilon,m})^{-1/2}u\approx u. $$
如果 $\epsilon=10^{-4}$，真实 inverse square root 在 $u$ 方向放大 $100$ 倍，而 diagonal approximation 几乎不放大。
这样的例子还可以构造很多。我们知道，Muon’s update rule 的核心作用就是把谱结构拉平。对于上面这个例子而言，真正的 inverse square root 能够把 $u$ 方向的特征值从 $\epsilon$ 拉平到 $1$，而 diagonal approximation 则无法做到这一点。问题的关键是 RMS 只看平均 entrywise 误差，而 inverse square root 关心的是谱结构，我们尤其关注那些小特征值。
因此，通过说明梯度/Momentum 的 Gram matrix 在 RMS 意义下接近 diagonal，并不能说明行归一化能够替代 Muon 的 preconditioner 作用。在模型训练中，nuances 很重要，update 积微成著造成最后收敛效果天差地别。
笔者的实验
笔者平时对优化器做过很多实验，行归一化是一种简单的动量 Precondition 方法，之前也有在 Pretrain 的 Setting 下尝试过。我的实验使用的是 32k 序列长，32M tokens per batch 的 setting。当然，由于资源有限，我会使用 Moonlight 的思路，不做 HP sweep，按 update RMS 和 Adam 对齐。实验下来效果差强人意。
这里的rownorm.* 是行归一化的实验，都没有跑完，因为看起来效果不是很值得持续跑下去。这里的两组实验的差别在于：
dense4b_stable_rownorm: 是直接沿用 Moonlight 的学习率换算：$0.2\cdot\sqrt{\max(A,B)}\cdot \eta$dense4b_stable_rownorm_fanout_scale: 使用的是$0.2\cdot\sqrt{B}\cdot\eta$
可以看到行归一化的效果其实是跟 Adam 差不多水平的。
当然每个人对于优化器的实验偏好不大一样。我比较倾向于用接近 production setting 来校验优化器效果，但是毕竟资源有限，无法做大规模的 HP sweep 来做最充分的验证。另一些朋友可能比较喜欢 speedrun，但是我觉得这个做法有很多 hacking 的问题，且按下不表了。如果有其他人的实验表明结果不同，欢迎探讨。
工程问题
最后还不得不提到工程实现的问题，抛开算法效果不论，在工程效率上可能有这样两方面的优势：
- 计算 FLOPs 被降低了，Muon 依赖的 NS 迭代涉及到多个 GEMM，可以被替换为行归一化操作。
- 大规模训练有望节省通信开销，我们考虑一个 FSDP 的场景，如果参数在 dim0 被切分，实际上对于行归一化而言，每个节点可以 inplace 做动量归一化并更新自己的参数分片，应该是可以做到像 Adam 一样的 0 通信更新。
但是话说回来，Muon 优化器本身的开销，放到 end2end 训练里，体感上大概就是 3%，取决于并行配置和实现，其实完全可以接受。
按照笔者的经验，行归一化优化器更像是 Adam 的平替（少一个 state）而不是 Muon 的平替。