流形上的最速下降：6. Muon + 双旋转

我们知道，用Adam、Muon等优化器更新矩阵参数时，奇异值和左右奇异向量都会随之变化，它们通常都是耦合在一起。也正是因为这种耦合性，我们无法简单地调控矩阵参数的奇异值，因此在奇异值出现异常增长时，我们无法简单有效地阻止它，这可能导致训练失败。

本文在《Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation》（后面简称Pion）的启发下，提出了一种单独更新矩阵左右奇异向量的Muon变体——“旋转Muon（MuonR）”，它能维持矩阵的奇异值分布不变，从而保证训练稳定性。

前文回顾 #

由于左右奇异向量组成的矩阵一定是正交矩阵，所以我们先来简单回顾一下正交约束下的Muon。设参数$\boldsymbol{W}\in\mathbb{R}^{n\times n}$，满足$\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$，设更新量为$\Delta\boldsymbol{W}=-\eta \boldsymbol{\Phi}$，我们希望参数更新之后依然满足正交性，那么对应的谱范数最速下降问题就是
\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\quad(\boldsymbol{W} - \eta \boldsymbol{\Phi})^{\top}(\boldsymbol{W} - \eta \boldsymbol{\Phi})=\boldsymbol{I}\end{equation}
可以解得$\boldsymbol{\Phi} = \boldsymbol{W}\boldsymbol{O}$，其中$\newcommand{msign}{\mathop{\text{msign}}}\boldsymbol{O}=\msign([\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}})$，这里$[\boldsymbol{X}]_{\text{skew}} = (\boldsymbol{X} - \boldsymbol{X}^{\top})/2$是反对称化算符。考虑上缩回操作，完整的更新规则是：
\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \boldsymbol{W}(\boldsymbol{I} - \eta\boldsymbol{O})\left(\boldsymbol{I} - \boldsymbol{O}^{\top}\boldsymbol{O} + \frac{\boldsymbol{O}^{\top}\boldsymbol{O}}{\sqrt{1+\eta^2}}\right)\label{eq:orth-steepest}\end{equation}
特别地，如果$[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$是满秩的，那么将简化成
\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \frac{\boldsymbol{W}(\boldsymbol{I} - \eta\boldsymbol{O})}{\sqrt{1+\eta^2}}\label{eq:orth-steepest-full}\end{equation}
推导过程在《流形上的最速下降：2. Muon + 正交》可以找到，这里就不详细展开了。不管是式$\eqref{eq:orth-steepest}$还是$\eqref{eq:orth-steepest-full}$，它们都是完全解析的，仅在Muon的基础上增加了几步矩阵乘法，复杂度没有明显增加，所以这个结果是完全实用的。

现在考虑矩阵$\boldsymbol{W}\in\mathbb{R}^{n\times m}(n \geq m)$，如果同时满足$\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$，那么我们称$\boldsymbol{W}$位于Stiefel流形上，它是正交矩阵概念的一般化，上述结果理论上可以推广到Stiefel流形上，但非方阵下需要求解一个非线性方程组，比较难实用化，个中细节可以参考《流形上的最速下降：3. Muon + Stiefel》。

瞬时重参 #

接着，我们将注意力转移到任意参数矩阵$\boldsymbol{W}\in\mathbb{R}^{n\times m}$上，目标是使参数更新过程中奇异值保持不变，从而杜绝奇异值异常增长的可能。

为了实现这个目标，我们采用“瞬时重参”的思想：在更新开始前，先将$\boldsymbol{W}$矩阵重新参数化为$\tilde{\boldsymbol{W}} = \boldsymbol{L}\boldsymbol{W}\boldsymbol{R}$，其中$\boldsymbol{L}\in\mathbb{R}^{n\times n},\boldsymbol{R}\in\mathbb{R}^{m\times m}$，并且都初始化为单位阵。这样一来，初始化时满足$\tilde{\boldsymbol{W}}=\boldsymbol{W}$，并且通过记$\boldsymbol{G} = \nabla_{\boldsymbol{W}}\mathcal{L}$，我们可以写出
\begin{equation}\nabla_{\boldsymbol{L}}\mathcal{L} = \boldsymbol{G}\boldsymbol{W}^{\top},\qquad \nabla_{\boldsymbol{R}}\mathcal{L} = \boldsymbol{W}^{\top}\boldsymbol{G}\end{equation}
随后，我们约定冻结$\boldsymbol{W}$，只更新$\boldsymbol{L}$和$\boldsymbol{R}$，同时在更新过程中保持$\boldsymbol{L},\boldsymbol{R}$的正交性，这样更新后的$\tilde{\boldsymbol{W}}$依然具有跟$\boldsymbol{W}$一样的奇异值。现在从$\boldsymbol{L},\boldsymbol{R}$的视角看，问题再次变成了正交流形上的最速下降，并且这时候$\boldsymbol{L},\boldsymbol{R}$都是方阵，对应的最速下降完全可解析解！根据式$\eqref{eq:orth-steepest-full}$可以直接写出更新规则
\begin{gather}\boldsymbol{L}\quad\leftarrow\quad (\boldsymbol{I} - \eta\boldsymbol{O}_L)\left(\boldsymbol{I} - \boldsymbol{O}_L^{\top}\boldsymbol{O}_L + \frac{\boldsymbol{O}_L^{\top}\boldsymbol{O}_L}{\sqrt{1+\eta^2}}\right)\\
\boldsymbol{R}\quad\leftarrow\quad (\boldsymbol{I} - \eta\boldsymbol{O}_R)\left(\boldsymbol{I} - \boldsymbol{O}_R^{\top}\boldsymbol{O}_R + \frac{\boldsymbol{O}_R^{\top}\boldsymbol{O}_R}{\sqrt{1+\eta^2}}\right)\end{gather}
其中$\boldsymbol{O}_L = \msign([\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}),\boldsymbol{O}_R = \msign([\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}})$。将新的$\boldsymbol{L},\boldsymbol{R}$与$\boldsymbol{W}$乘起来，那么就得到完整的更新规则
\begin{equation}\boldsymbol{W} \quad \leftarrow\quad \boldsymbol{L}\boldsymbol{W}\boldsymbol{R}\end{equation}
这便是基于“瞬时重参”思想推导出来“旋转Muon（Muon under Rotation，MuonR）”。实际场景下通常还有动量，我们将它理解为平滑过后的梯度，所以只需要将梯度$\boldsymbol{G}$换成动量$\boldsymbol{M}$。

一些细节 #

由于$\boldsymbol{L},\boldsymbol{R}$的更新都需要算一次$\msign$，所以即便在最理想（$n=m$）的情况下，MuonR计算量也是Muon的两倍。不过对于足够大的模型来说，这个计算量翻倍对端到端的训练时间影响是比较微弱的，通常还可以接受。如果想要减少这部分开销，可以考虑$\boldsymbol{L},\boldsymbol{R}$交替更新的做法，平摊一下计算量。

事实上，MuonR最大的问题在于，它自始至终保持矩阵的全体奇异值不变，这意味着我们必须在初始化时就把参数的全体奇异值确定下来。这并不容易，因为不同位置的矩阵可能需要不同的尺度，强行设置为同一组值大概率是次优的。

一个可行的思路是，在适当的随机初始化基础上，给每个矩阵前或后添加一个Element-wise相乘的向量，用来补偿尺度的自由度。对于紧接在RMSNorm后面的矩阵，RMSNorm自带的gamma参数已经承担了这一角色，所以这部分矩阵则可以省略这一操作。

至于矩阵的初始奇异值如何选取，我们可以考虑常规的随机初始化，也可以考虑按Zipf's law的形式构建。进一步地，我们还可以尝试将奇异值熵调节成《矩阵参数的奇异值熵越高越好吗？》计算出来的最优熵，以期达到更优的效果。

当然，如果我们确实能事先确定矩阵的奇异值——比如本就期望某处参数始终保持正交性——那就不用考虑这些了，直接套用MuonR就行。

中途切换 #

还有一种可选的做法是“中途切换”，仅将MuonR作为“维稳”的手段。

具体来说，一开始我们还是用常规的Muon，并监测矩阵的谱范数/F范数，一旦矩阵的范数超出我们期望的范围，那就切换到MuonR，由于两种Muon都只依赖于同一个梯度/动量，它们只是计算不同，所以这种切换是允许的。MuonR不改变矩阵的奇异值，那么谱范数/F范数都不再增长，正好可以作为“维稳”手段使用。

但要注意的是，我们要让切换前后的更新幅度要尽可能对齐，避免引入“突变”。为此，我们考虑MuonR的一阶近似
\begin{equation}\boldsymbol{L}\boldsymbol{W}\boldsymbol{R} \approx (\boldsymbol{I} - \eta\boldsymbol{O}_L) \boldsymbol{W} (\boldsymbol{I} - \eta\boldsymbol{O}_R) \approx \boldsymbol{W} - \eta(\boldsymbol{O}_L \boldsymbol{W} + \boldsymbol{W} \boldsymbol{O}_R)\end{equation}
由于$\boldsymbol{O}_L,\boldsymbol{O}_R$的奇异值都不超过1（注意我们无法保证$[\boldsymbol{G}\boldsymbol{W}^{\top}]_{\text{skew}}$和$[\boldsymbol{W}^{\top}\boldsymbol{G}]_{\text{skew}}$都满秩，所以不能直接使用$\boldsymbol{O}_L,\boldsymbol{O}_R$的正交性），，所以有$\Vert\boldsymbol{O}_L \boldsymbol{W}\Vert_F \leq \Vert \boldsymbol{W}\Vert_F$和$\Vert \boldsymbol{W}\boldsymbol{O}_R\Vert_F\leq \Vert\boldsymbol{W}\Vert_F$，于是
\begin{equation}\Vert\boldsymbol{O}_L \boldsymbol{W} + \boldsymbol{W} \boldsymbol{O}_R\Vert_F \leq \Vert\boldsymbol{O}_L \boldsymbol{W}\Vert_F + \Vert\boldsymbol{W} \boldsymbol{O}_R\Vert_F \leq 2\Vert\boldsymbol{W}\Vert_F\end{equation}
常规Muon是$\boldsymbol{W} - \eta \msign(\boldsymbol{G})$，$\msign(\boldsymbol{G})$的F范数普遍为$\sqrt{\min(n,m)}$，因此为了对齐更新量F范数，从Muon切换到MuonR，大致上要将学习率乘以$\frac{\sqrt{\min(n,m)}}{2\Vert\boldsymbol{W}\Vert_F}$。

实践中，上式第一个不等号可能不够紧，$\boldsymbol{O}_L \boldsymbol{W}$和$\boldsymbol{W} \boldsymbol{O}_R$更多是近乎正交的关系，根据勾股定理，结果应该约等于$\sqrt{2}\Vert\boldsymbol{W}\Vert_F$，那么这个倍数应当多乘以$\sqrt{2}$。不过考虑到$\sqrt{2}$与$1$差别不是特别大，并且为了保证极端情况下的可用性，建议还是保留上述形式。

异同分析 #

本文在开头就明确表示，MuonR受启发自Pion，我们先来看看它们的联系与区别。

首先，将更新规则限定为左乘和右乘正交矩阵这种双旋转形式的思路主要来源于Pion；在确定这一更新形式后，通过“瞬时重参”获取对应的梯度这一步是比较自然的；接着，Pion和MuonR就开始“分道扬镳”：

1、Pion通过矩阵指数$\exp(\text{反对称矩阵})$来实现正交性，实际计算中则通过展开到二阶来近似；

2、Pion走的是Adam路线，并且将$\boldsymbol{L},\boldsymbol{R}$的梯度分别滑动平均，这使得它的缓存变量达到了4组；

3、MuonR走的是Muon路线，跟Muon一样只缓存动量，这使得我们可以随时跟Muon切换；

4、MuonR基于正交流形下最速下降的解析解，只需有限额外步骤即可精确地实现正交性。

总的来说，Pion的正交性设计比较经验化，且四组缓存变量让人有点“望而却步”；MuonR则是Muon、正交流形最速下降等一系列工作下相对自然的产物，笔者认为它整体上要更符合第一性原理一些。

文章小结 #

本文提出了MuonR，一种将更新形式约束为左右旋转矩阵的Muon变体，它能够保持矩阵的奇异值分布不变，是一种简洁的维持训练稳定性的训练方案。

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

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