为什么官方版Muon比MuP版多出一个max(1, ⋅)？

在文章《Muon优化器指南：快速上手与关键细节》中，我们罗列了Muon的几个版本，它们的区别是学习率的矩阵形状相关的缩放因子不同，其中“官方版（KellerJordan版）”只比“MuP版”多出了一个$\max(1,\cdot)$的截断操作。本文就来专门讨论一下，这个截断操作是怎么来的呢？

几个版本 #

Muon的更新规则，可以统一地写成
\begin{equation}\begin{aligned}\newcommand{msign}{\mathop{\text{msign}}}
\boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt]
\boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t (\alpha \msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1})
\end{aligned}\end{equation}
几个不同的版本区别在于$\alpha$，分别是：
$$\alpha = \left\{
\begin{aligned} &1 & \color{skyblue}{(\text{朴素版})} \\[5pt]
& \sqrt{\max(1, d_{out}/d_{in})} & \color{skyblue}{(\text{KellerJordan版})} \\[5pt]
& \sqrt{d_{out}/d_{in}} & \color{skyblue}{(\text{MuP版})} \\[5pt]
& 0.2\times\sqrt{\max(d_{out},d_{in})} & \color{skyblue}{(\text{Moonlight版})}
\end{aligned}\right.$$

其中矩阵$\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$，代表线性层$\boldsymbol{y}=\boldsymbol{x}\boldsymbol{W}$的训练参数，其中输入$\boldsymbol{x}\in\mathbb{R}^{d_{in}}$是行向量。

本文主要关心“KellerJordan版”和“MuP版”，前者在后者的基础上多加了个$\max(1,)$。按照我们在《高阶MuP：更简明但更高明的谱条件缩放》、《MuP之上：2. 线性层与最速下降》等文章的分析，在MuP相关的谱条件约束下，最速下降应该就是MuP版Muon才对，多出来的$\max(1,\cdot)$该如何解释呢？

特征增量 #

简单起见，下面的讨论都省掉下标$t$。不失一般性，我们假设动量$\boldsymbol{M}$是满秩的，那么$\boldsymbol{\Phi} = \msign(\boldsymbol{M})$的奇异值全为1，那么当$d_{in} \leq d_{out}$时，$\boldsymbol{\Phi} \boldsymbol{\Phi}^{\top} = \boldsymbol{I}_{d_{in}}$，当$d_{in} > d_{out}$时，$\boldsymbol{\Phi}^{\top} \boldsymbol{\Phi} = \boldsymbol{I}_{d_{out}}$。

记$\Delta \boldsymbol{W} = \eta\alpha \boldsymbol{\Phi}$，我们要做的事情，就是寻找$\alpha$与$d_{in},d_{out}$的关系。从《为什么我们偏爱各向同性？基于最速下降的理解》中我们知道，参数实际上只是模型的副产物，特征层面的变化可能会更加本质。将$\Delta \boldsymbol{W}$转换到特征层面是$\Delta \boldsymbol{y} = \boldsymbol{x} \Delta\boldsymbol{W} = \eta\alpha \boldsymbol{x}\boldsymbol{\Phi}$，那么$\Vert\Delta \boldsymbol{y}\Vert_{RMS} = \alpha \Vert\boldsymbol{x}\boldsymbol{\Phi}\Vert_{RMS}$。

接下来要分情况讨论。首先是当$d_{in} \leq d_{out}$时，$\boldsymbol{\Phi}$可以写成$\boldsymbol{U}[\boldsymbol{I}_{d_{in}}, \boldsymbol{0}_{d_{in}\times (d_{out}-d_{in})}]\boldsymbol{V}^{\top}$的形式，其中$\boldsymbol{U}\in\mathbb{R}^{d_{in}\times d_{in}}, \boldsymbol{V}\in\mathbb{R}^{d_{out}\times d_{out}}$都是正交矩阵，那么
\begin{equation}\begin{aligned}
\Vert\Delta \boldsymbol{y}\Vert_{RMS} =&\, \eta\alpha\big\Vert\boldsymbol{x}\boldsymbol{U}[\boldsymbol{I}_{d_{in}}, \boldsymbol{0}_{d_{in}\times (d_{out}-d_{in})}]\boldsymbol{V}^{\top}\big\Vert_{RMS} \\[4pt]
=&\, \eta\alpha\big\Vert\boldsymbol{x}\boldsymbol{U}[\boldsymbol{I}_{d_{in}}, \boldsymbol{0}_{d_{in}\times (d_{out}-d_{in})}]\big\Vert_{RMS} \\[4pt]
=&\, \eta\alpha\big\Vert[\boldsymbol{x}\boldsymbol{U}, \boldsymbol{0}_{d_{out}-d_{in}}]\big\Vert_{RMS} \\[4pt]
=&\, \eta\alpha\sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{x}\boldsymbol{U}\Vert_{RMS} \\[4pt]
=&\, \eta\alpha\sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{x}\Vert_{RMS} \\
\end{aligned}\end{equation}
注意全程都是等号，所以我们只需要设$\alpha = \sqrt{d_{out}/d_{in}}$，就可以让“每一个”$\Delta \boldsymbol{y}$的RMS都是$\eta\Vert\boldsymbol{x}\Vert_{RMS}$，即所有Token的相对更新幅度一致。

各向同性 #

很遗憾，对于第二种情况$d_{in} > d_{out}$，上述“完全一致”的目标并无法实现。具体来说，此时$\boldsymbol{\Phi}$的SVD要写成$\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\boldsymbol{V}^{\top}$，于是
\begin{equation}\begin{aligned}
\Vert\Delta \boldsymbol{y}\Vert_{RMS} =&\, \eta\alpha\left\Vert\boldsymbol{x}\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\boldsymbol{V}^{\top}\right\Vert_{RMS} \\[5pt]
=&\, \eta\alpha\left\Vert\boldsymbol{x}\boldsymbol{U}\begin{bmatrix}\boldsymbol{I}_{d_{out}} \\ \boldsymbol{0}_{(d_{in}-d_{out})\times d_{out}}\end{bmatrix}\right\Vert_{RMS} \\[5pt]
=&\, \eta\alpha\big\Vert(\boldsymbol{x}\boldsymbol{U})_{[:d_{out}]}\big\Vert_{RMS}
\end{aligned}\end{equation}
$\boldsymbol{x}\boldsymbol{U}$是一个$d_{in}$维向量，$d_{in} > d_{out}$，所以$(\boldsymbol{x}\boldsymbol{U})_{[:d_{out}]}$只是截取了$\boldsymbol{x}\boldsymbol{U}$的前$d_{out}$个维度去算RMS，这时候它的RMS是不确定的，最大可以去到$\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$（Worst Case），最小可以是0。

我们知道正交矩阵不改变RMS，所以$\Vert\boldsymbol{x}\boldsymbol{U}\Vert_{RMS}=\Vert\boldsymbol{x}\Vert_{RMS}$，当$\boldsymbol{x}$的分布足够各向同性时，我们可以认为这代表着$\boldsymbol{x}\boldsymbol{U}$每个分量的平均尺度都是$\Vert\boldsymbol{x}\Vert_{RMS}$，于是取前$d_{out}$个分量去RMS，平均来说也约等于$\Vert\boldsymbol{x}\Vert_{RMS}$，即$\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \alpha\Vert\boldsymbol{x}\Vert_{RMS}$，因此只需要取$\alpha = 1$，就可以实现跟上一节类似的效果。

各向异性 #

将前两节的结果综合起来，我们就得到
\begin{equation}\alpha = \sqrt{\max\left(1, \frac{d_{out}}{d_{in}}\right)}\end{equation}
这正是KellerJordan版Muon出现的$\max(1,\cdot)$。

然而，上一节的结论依赖于输入$\boldsymbol{x}$足够各向同性的假设，这在训练早期或许近似成立，但随着训练的推进，特征的分布会逐渐各向异性，集中在让$\Vert\Delta \boldsymbol{y}\Vert_{RMS}$最大化的“Worst Case”上面，此时平均近似$\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\alpha\Vert\boldsymbol{x}\Vert_{RMS}$就不够准确了，反而是最大值$\eta\alpha\sqrt{d_{in}/d_{out}}\Vert\boldsymbol{x}\Vert_{RMS}$会更加准确。

这种情况下，让$\Vert\Delta \boldsymbol{y}\Vert_{RMS}\approx \eta\Vert\boldsymbol{x}\Vert_{RMS}$的$\alpha$是$\sqrt{d_{out}/d_{in}}$，跟$d_{in} \leq d_{out}$时的结论一致，重新得到MuP版结果。也就是说，对于训练的中后期，MuP版Muon更加科学。对于这种不一致性，我们有两种策略：一是始终用MuP版Muon，这会稍微降低一点初期的收敛速度，但毕竟中后期才是训练的“重中之重”；二是将缩放因子改为
\begin{equation}\alpha = \sqrt{\max\left(\tau_t, \frac{d_{out}}{d_{in}}\right)}\end{equation}
其中$\tau_t$从1到0单调衰减，这样实现了从KellerJordan版到MuP版的渐变，代价是多了一个Schedule要调。

文章小结 #

本文主要从“特征增量”均匀性的角度解释了KellerJordan版中$\max(1,\cdot)$的来源。

转载到请包括本文地址：https://kexue.fm/archives/11772

更详细的转载事宜请参考：《科学空间FAQ》