慢即是快：2. Adam的尽头还是Muon？

在上一篇文章《慢即是快：1. Shampoo的尽头是Muon？》中，我们从凸优化的平均收敛结论出发，建立了“正交最慢下降”框架，并用于改进任意Kronecker积形式的预条件优化器。结果发现，在此框架的优化之下，这些优化器都有趋于Muon的迹象，于是笔者发出了“Shampoo的尽头是Muon”的悲观感慨。

除了Kronecker积外，预条件矩阵的另一种更简单的结构化近似是对角近似，常见例子是RMSProp、Adam等。注意对角近似虽然计算上比Kronecker积简单，但它俩并没有包含关系，所以一个新的问题是：对这类优化器施加正交最慢下降，会得到新的结果吗？

目标分析[#](https://spaces.ac.cn/archives/11964#%E7%9B%AE%E6%A0%87%E5%88%86%E6%9E%90)

沿用上一篇记号，梯度$\mathbf{\mathit{G}} \in \mathbb{R}^{n \times m} , \mathbf{\mathit{g}} = \text{vec} ⁡ \left(\right. \mathbf{\mathit{G}} \left.\right) \in \mathbb{R}^{n m}$，引入矩阵$\mathbf{\mathit{N}} \in \mathbb{R}_{+}^{n \times m}$，所谓对角近似，是指向量化表述下的预条件矩阵是一个对角阵，即$\mathbf{\mathit{H}}^{- 1} = \text{diag} ⁡ \left(\right. \text{vec} ⁡ \left(\right. \mathbf{\mathit{N}} \left.\right) \left.\right)$，那么

$$ (\text{1}) \mathbf{\mathit{H}}^{- 1} \mathbf{\mathit{g}} = \text{vec} ⁡ \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{G}} \left.\right) $$

或者说更新规则为$\mathbf{\mathit{W}} - \eta \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{G}}$，这里的$\bigodot$是Hadamard积。我们要做的是给$\mathbf{\mathit{N}}$引入Kronecker积形式的正交变换，使它在正交最慢下降原理之下，变成一个更好的优化器。 具体来说，我们引入正交矩阵$\mathbf{\mathit{X}} \in \mathbb{R}^{n \times n} , \mathbf{\mathit{Y}} \in \mathbb{R}^{m \times m}$，设$\mathbf{\mathit{O}} = \mathbf{\mathit{Y}} \bigotimes \mathbf{\mathit{X}}$，将对角的预条件矩阵换成$\mathbf{\mathit{O}}^{\top} \text{diag} ⁡ \left(\right. \text{vec} ⁡ \left(\right. \mathbf{\mathit{N}} \left.\right) \left.\right) \mathbf{\mathit{O}}$，即更新量从$\mathbf{\mathit{N}} \bigodot \mathbf{\mathit{G}}$变成了$\mathbf{\mathit{X}}^{\top} \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \mathbf{\mathit{Y}}$，这个变换需要一点Kronecker积基础，但这里就不展开了。

直观来看，就是给$\mathbf{\mathit{G}}$找一组新的正交基，然后在上面做对角预条件，再变换回来。根据正交最慢下降原理，我们要最小化它与$\mathbf{\mathit{G}}$的内积，即

$$ (\text{2}) \underset{\mathbf{\mathit{X}} , \mathbf{\mathit{Y}}}{min} \text{tr} ⁡ \left(\right. \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \mathbf{\mathit{Y}} \left.\right) \text{s}.\text{t}. \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{X}} = \mathbf{\mathit{I}}_{n} , \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Y}} = \mathbf{\mathit{I}}_{m} $$

统一形式[#](https://spaces.ac.cn/archives/11964#%E7%BB%9F%E4%B8%80%E5%BD%A2%E5%BC%8F)

注意到

$$ (\text{3}) \text{tr} ⁡ \left(\right. \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \mathbf{\mathit{Y}} \left.\right) = \text{tr} ⁡ \left(\right. \mathbf{\mathit{Y}} \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \left.\right) = \parallel \mathbf{\mathit{N}}^{\bigodot / 2} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \parallel_{F}^{2} $$

这里的$\mathbf{\mathit{N}}^{\bigodot / 2}$是指在Hadamard积的意义下求平方根，即逐元素开平方。此外，上篇我们求解了

$$ (\text{4}) \underset{\mathbf{\mathit{X}} , \mathbf{\mathit{Y}}}{min} \text{tr} ⁡ \left(\right. \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{P}} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}} \mathbf{\mathit{Y}} \left.\right) \text{s}.\text{t}. \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{X}} = \mathbf{\mathit{I}}_{n} , \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Y}} = \mathbf{\mathit{I}}_{m} $$

目标函数可以改写成

$$ (\text{5}) \text{tr} ⁡ \left(\right. \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{P}} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}} \mathbf{\mathit{Y}} \left.\right) = \text{tr} ⁡ \left(\right. \mathbf{\mathit{Y}} \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{P}} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}} \left.\right) = \parallel \mathbf{\mathit{P}}^{1 / 2} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}}^{1 / 2} \parallel_{F}^{2} $$

可以发现两个问题的目标函数都具有$f \left(\right. \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right)$的形式，所以接下来我们一般地考虑

$$ (\text{6}) \underset{\mathbf{\mathit{X}} , \mathbf{\mathit{Y}}}{min} f \left(\right. \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \text{s}.\text{t}. \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{X}} = \mathbf{\mathit{I}}_{n} , \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Y}} = \mathbf{\mathit{I}}_{m} $$

驻点条件[#](https://spaces.ac.cn/archives/11964#%E9%A9%BB%E7%82%B9%E6%9D%A1%E4%BB%B6)

先求上述问题的驻点条件。为此，我们引入拉格朗日乘子$\mathbf{\mathit{A}} \in \mathbb{R}^{n \times n} , \mathbf{\mathit{B}} \in \mathbb{R}^{m \times m}$，定义辅助函数

$$ (\text{7}) F = f \left(\right. \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) - \text{tr} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \left(\right. \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{X}} - \mathbf{\mathit{I}}_{n} \left.\right) \left.\right) - \text{tr} ⁡ \left(\right. \mathbf{\mathit{B}}^{\top} \left(\right. \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Y}} - \mathbf{\mathit{I}}_{m} \left.\right) \left.\right) $$

记$\mathbf{\mathit{Z}} = \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} , \mathbf{\mathit{K}} = \nabla_{\mathbf{\mathit{Z}}} f \left(\right. \mathbf{\mathit{Z}} \left.\right)$，分别求$F$对$\mathbf{\mathit{X}} , \mathbf{\mathit{Y}}$的梯度并让它们等于零得

$$ (\text{8}) \nabla_{\mathbf{\mathit{X}}} F = \mathbf{\mathit{K}} \mathbf{\mathit{Y}} \mathbf{\mathit{G}}^{\top} - \mathbf{\mathit{X}} \left(\right. \mathbf{\mathit{A}} + \mathbf{\mathit{A}}^{\top} \left.\right) = 0 \\ (\text{9}) \nabla_{\mathbf{\mathit{Y}}} F = \mathbf{\mathit{K}}^{\top} \mathbf{\mathit{X}} \mathbf{\mathit{G}} - \mathbf{\mathit{Y}} \left(\right. \mathbf{\mathit{B}} + \mathbf{\mathit{B}}^{\top} \left.\right) = 0 $$

第一式两端乘$\mathbf{\mathit{X}}^{\top}$，第二式两端乘$\mathbf{\mathit{Y}}^{\top}$，结合$\mathbf{\mathit{X}} , \mathbf{\mathit{Y}}$的正交性，可以得到驻点条件

$$ (\text{10}) \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{K}} \mathbf{\mathit{Y}} \mathbf{\mathit{G}}^{\top} 对称 , \mathbf{\mathit{G}}^{\top} \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{K}} \mathbf{\mathit{Y}} 对称 $$

在上一篇文章我们已经证明过，这意味着$\mathbf{\mathit{X}}^{\top} \mathbf{\mathit{K}} \mathbf{\mathit{Y}}$具有$\mathbf{\mathit{U}} \mathbf{\mathit{D}} \mathbf{\mathit{V}}^{\top}$的形式，其中$\mathbf{\mathit{U}} \mathbf{\Sigma} \mathbf{\mathit{V}}^{\top}$是$\mathbf{\mathit{G}}$的SVD，$\mathbf{\mathit{D}}$是一个对角阵。这是上述优化问题最优解的共同特征，与$f$的具体选择无关，算得上一个极强的结论了。

洗去基底[#](https://spaces.ac.cn/archives/11964#%E6%B4%97%E5%8E%BB%E5%9F%BA%E5%BA%95)

先来用这里的结论，过一下上篇文章的目标$\parallel \mathbf{\mathit{P}}^{1 / 2} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}}^{1 / 2} \parallel_{F}^{2}$，此时$\mathbf{\mathit{K}} = 2 \mathbf{\mathit{P}} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}}$，我们关心的更新量是$\mathbf{\Phi} = \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{P}} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}} \mathbf{\mathit{Y}} = \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{K}} \mathbf{\mathit{Y}} / 2$，根据上一节的结果，$\mathbf{\Phi}$具有$\mathbf{\mathit{U}} \mathbf{\mathit{D}} \mathbf{\mathit{V}}^{\top}$的形式，即$\mathbf{\mathit{P}} , \mathbf{\mathit{Q}}$的特征矩阵完全被消去。

让笔者非常意外的是，同样的结论也能推广到本文的主角——也就是目标$\parallel \mathbf{\mathit{N}}^{\bigodot / 2} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \parallel_{F}^{2}$——此时$\mathbf{\mathit{K}} = 2 \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top}$，我们关心的更新量是$\mathbf{\Phi} = \mathbf{\mathit{X}}^{\top} \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \mathbf{\mathit{Y}} = \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{K}} \mathbf{\mathit{Y}} / 2$，那么依然根据驻点条件的结果，$\mathbf{\Phi}$具有$\mathbf{\mathit{U}} \mathbf{\mathit{D}} \mathbf{\mathit{V}}^{\top}$的形式。

换句话说，不管是Kronecker积还是对角近似，预条件矩阵的“基”都会被正交最慢下降洗掉：更新量的左、右奇异矩阵一律退化为梯度（或动量）的左、右奇异矩阵，预条件所能改变的至多只有奇异值——这与Muon的构造如出一辙。这就是标题“尽头是Muon”的第一重含义。

可否均匀[#](https://spaces.ac.cn/archives/11964#%E5%8F%AF%E5%90%A6%E5%9D%87%E5%8C%80)

如果能证明$\mathbf{\mathit{D}}$的对角元趋于均衡，那么就更加跟Muon趋同了，这是“尽头是Muon”的第二重含义。

上一篇文章，我们已经对Kronecker积情形做了详细论证，但那个过程不好推广，因此这里改用一种启发式方法来感知$\mathbf{\mathit{D}}$的特性。不失一般性设$n \geq m$，我们有

$$ (\text{11}) \text{tr} ⁡ \left(\right. \mathbf{\mathit{G}}^{\top} \mathbf{\Phi} \left.\right) = \text{tr} ⁡ \left(\right. \mathbf{\mathit{V}} \mathbf{\Sigma}^{\top} \mathbf{\mathit{D}} \mathbf{\mathit{V}}^{\top} \left.\right) = \text{tr} ⁡ \left(\right. \mathbf{\Sigma}^{\top} \mathbf{\mathit{D}} \left.\right) = \sum_{i = 1}^{m} \sigma_{i} d_{i} \geq m \left(\left(\right. \prod_{i = 1}^{m} \sigma_{i} d_{i} \left.\right)\right)^{1 / m} = m \left(\left(\right. \prod_{i = 1}^{m} \sigma_{i} \prod_{i = 1}^{m} d_{i} \left.\right)\right)^{1 / m} $$

其中$\mathbf{\Sigma} = \text{diag} ⁡ \left(\right. \sigma_{1} , \sigma_{2} , \hdots \left.\right) , \mathbf{\mathit{D}} = \text{diag} ⁡ \left(\right. d_{1} , d_{2} , \hdots \left.\right)$，等号成立的条件是$d_{i} \propto \sigma_{i}^{- 1}$。这个结果表明，只要$\mathbf{\Phi}$的奇异值之积为常数，那么$\text{tr} ⁡ \left(\right. \mathbf{\mathit{G}}^{\top} \mathbf{\Phi} \left.\right)$的理想最小值在$\mathbf{\mathit{D}} \propto \mathbf{\Sigma}^{- 1}$（对角元取逆，形状不变）取到。 当然，由于我们只能改动$\mathbf{\Phi}$中的正交矩阵，所以这个理想的最小值很难实现，只能猜测它是最理想（$\mathbf{\mathit{D}} \propto \mathbf{\Sigma}^{- 1}$）和最糟糕（$\mathbf{\mathit{D}} \propto \mathbf{\Sigma}$）之间的一个中间结果，而它俩的几何平均正是$\mathbf{\mathit{D}} \propto \mathbf{\mathit{I}}$，即奇异值均匀化。

谱之乘积[#](https://spaces.ac.cn/archives/11964#%E8%B0%B1%E4%B9%8B%E4%B9%98%E7%A7%AF)

那$\mathbf{\Phi}$的奇异值之积是否为常数呢？我们还是按照先Kronecker积后对角近似的顺序分别讨论。首先对于“Kronecker积+方阵”情形，确实是常数：

$$ (\text{12}) \prod_{i = 1}^{m} d_{i} = \left|\right. det \left(\right. \mathbf{\Phi} \left.\right) \left|\right. = \left|\right. det \left(\right. \mathbf{\mathit{X}}^{\top} \mathbf{\mathit{P}} \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \mathbf{\mathit{Q}} \mathbf{\mathit{Y}} \left.\right) \left|\right. = \left|\right. det \left(\right. \mathbf{\mathit{P}} \left.\right) det \left(\right. \mathbf{\mathit{G}} \left.\right) det \left(\right. \mathbf{\mathit{Q}} \left.\right) \left|\right. $$

即$\mathbf{\Phi}$在正交矩阵$\mathbf{\mathit{X}} , \mathbf{\mathit{Y}}$的变换下，行列式绝对值保持不变，从而奇异值之积不变，上一节的结论可用。不过，对于非方阵，上述恒等式不再成立，对此，我们有两种方案：一是按照上一篇文章的方法仔细讨论，二是不加证明地认为上式对非方阵也近似成立。 更棘手的是对角近似情形：Hadamard积虽然简单，但它并不是本质的矩阵运算，它的引入会导致矩阵的各种分析都变得异常困难，所以不管是否方阵，我们无法对$\mathbf{\Phi} = \mathbf{\mathit{X}}^{\top} \left(\right. \mathbf{\mathit{N}} \bigodot \mathbf{\mathit{X}} \mathbf{\mathit{G}} \mathbf{\mathit{Y}}^{\top} \left.\right) \mathbf{\mathit{Y}}$的奇异值/行列式做出有效的估计。

不过，由于$\mathbf{\Phi}$可变的只有正交矩阵，很多时候正交变换不改变行列式，所以我们猜测“$\mathbf{\Phi}$的行列式也不会有明显变化”，从而得到$\mathbf{\mathit{D}}$也有均匀趋势。注意，这个猜测很“无力”：它肯定不是恒等成立，至于近似程度如何，也无法判断，它更多是我们关于行列式与正交矩阵的“信仰”。

文章小结[#](https://spaces.ac.cn/archives/11964#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文将“正交最慢下降”推广到了对角预条件（RMSProp、Adam家族）。通过提炼出两个目标共同的抽象形式，我们统一得到“不管是Kronecker积还是对角近似，都会被正交最慢下降消去原有基底，至多改变梯度或动量的奇异值谱”的结论。进一步地，我们简单讨论了奇异值是否会趋于均衡。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11964](https://spaces.ac.cn/archives/11964 "慢即是快：2. Adam的尽头还是Muon？")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")