基于流式幂迭代的Muon实现：4. 原理

经过《基于流式幂迭代的Muon实现：1. 初识》、《基于流式幂迭代的Muon实现：2. 加速》和《基于流式幂迭代的Muon实现：3. 雕琢》三篇文章，想必大家已经对流式幂迭代（Streaming Power Iteration）的思想、实现、加速等细节有所了解，总的来说，这称得上是一种颇有竞争力的Muon实现方式，并且得益于它直接近似计算SVD，所以它还具备更好的拓展性。

受限于篇幅，当时我们对相关运算的数学原理描述得相对简略，因此在这篇文章中，我们补充部分关于幂迭代和QR分解的数学推导，以建立更完整的理论图景。不过，这里的推导依然是侧重解释性而非严格性，主要是为了帮大家（包括笔者）理清思路，还请专业读者海涵。

共轴等价[#](https://spaces.ac.cn/archives/11710#%E5%85%B1%E8%BD%B4%E7%AD%89%E4%BB%B7)

在开始推导之前，我们需要先引入“共轴等价”的概念。对于矩阵$\mathbf{\mathit{A}} , \mathbf{\mathit{B}} \in \mathbb{R}^{n \times m}$，如果存在一个符号矩阵$\mathbf{\mathit{S}}$满足$\mathbf{\mathit{A}} = \mathbf{\mathit{B}} \mathbf{\mathit{S}}$，那么称$\mathbf{\mathit{A}}$与$\mathbf{\mathit{B}}$“共轴等价（Coaxial Equivalent）”，它们互为对方的“共轴矩阵”。这里的“符号矩阵（Signature matrix）”是指为对角线为$\pm 1$的对角矩阵，即$\text{diag} ⁡ \left(\right. \pm 1 , \pm 1 , \hdots , \pm 1 \left.\right)$。

需要说明的是，“共轴（Coaxial）”这个词是笔者自行提出用来描述这种等价关系的，因为在坐标系视角下，满足条件的矩阵$\mathbf{\mathit{A}} , \mathbf{\mathit{B}}$实际上描述了同一个坐标系，只不过某些轴的正方向选择有所不同，有些文献似乎称它们为“符号等价（Sign Equivalent）”，但笔者感觉不如“共轴”来得直观。

之所以引入共轴的概念，是因为很多矩阵分解都是在共轴等价的意义下才唯一的（所以笔者也疑惑，这么常用的一种等价关系居然没有标准命名）。比如QR分解，设矩阵$\mathbf{\mathit{A}} \in \mathbb{R}^{n \times m} \left(\right. n \geq m \left.\right)$且满秩，$\mathbf{\mathit{Q}}_{1} \mathbf{\mathit{R}}_{1}$和$\mathbf{\mathit{Q}}_{2} \mathbf{\mathit{R}}_{2}$是它的两种QR分解，那么$\mathbf{\mathit{Q}}_{1}$与$\mathbf{\mathit{Q}}_{2}$共轴、$\mathbf{\mathit{R}}_{1}$与$\mathbf{\mathit{R}}_{2}$共轴。

这种唯一性不难理解，因为在Gram-Schmidt正交化的过程中，我们可以随意反转正交化后的向量而不改变正交性。一般的教程通常是限定$\mathbf{\mathit{R}}$的对角线元素为正来表达QR分解的唯一性，这也是一种途径，但有时候会带来不必要的麻烦。此外，SVD的唯一性也是在共轴等价下成立的。

幂之迭代[#](https://spaces.ac.cn/archives/11710#%E5%B9%82%E4%B9%8B%E8%BF%AD%E4%BB%A3)

在《基于流式幂迭代的Muon实现：1. 初识》中，我们是通过逐特征向量求解的方式来引入幂迭代的，然后直接切换到并行版本，认为它具有相同的收敛结果。这里则直接从并行版本出发，证明它的收敛性。

还是设矩阵$\mathbf{\mathit{A}} \in \mathbb{R}^{n \times m} \left(\right. n \geq m \left.\right)$，并且设它满秩（秩为$m$），考虑幂迭代

$$ (\text{1}) \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) , \mathbf{\mathit{V}}_{0} = \mathbf{\mathit{I}} $$

设$\mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1}$的QR分解是$\mathbf{\mathit{Q}}_{t} \mathbf{\mathit{R}}_{t}$，那么$\text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) = \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \mathbf{\mathit{R}}_{t}^{- 1}$，那么逐步迭代得

$$ (\text{2}) \mathbf{\mathit{V}}_{t} = \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right)^{t} \mathbf{\mathit{R}}_{1}^{- 1} \mathbf{\mathit{R}}_{2}^{- 1} \hdots \mathbf{\mathit{R}}_{t}^{- 1} $$

接着设$\mathbf{\mathit{A}}$的SVD为$\mathbf{\mathit{U}} \mathbf{\Sigma} \mathbf{\mathit{V}}^{\top}$，那么$\mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} = \mathbf{\mathit{V}} \mathbf{\Sigma}^{2} \mathbf{\mathit{V}}^{\top}$，代入到上式得

$$ (\text{3}) \mathbf{\mathit{V}}_{t} = \mathbf{\mathit{V}} \mathbf{\Sigma}^{2 t} \mathbf{\mathit{V}}^{\top} \mathbf{\mathit{R}}_{1}^{- 1} \mathbf{\mathit{R}}_{2}^{- 1} \hdots \mathbf{\mathit{R}}_{t}^{- 1} = \mathbf{\mathit{V}} \underset{\mathbf{\mathit{B}}}{\underbrace{\left(\right. \mathbf{\Sigma}^{2 t} \mathbf{\mathit{V}}^{\top} \mathbf{\Sigma}^{- 2 t} \left.\right)}} \underset{\mathbf{\mathit{C}}}{\underbrace{\left(\right. \mathbf{\Sigma}^{2 t} \mathbf{\mathit{R}}_{1}^{- 1} \mathbf{\mathit{R}}_{2}^{- 1} \hdots \mathbf{\mathit{R}}_{t}^{- 1} \left.\right)}} $$

结果分$\mathbf{\mathit{V}} , \mathbf{\mathit{B}} , \mathbf{\mathit{C}}$三部分，我们要证明的是在共轴意义下$\mathbf{\mathit{V}}_{t} \rightarrow \mathbf{\mathit{V}}$。已知$\mathbf{\mathit{V}}_{t} , \mathbf{\mathit{V}}$都是正交矩阵，然后由上三角矩阵的封闭性可知$\mathbf{\mathit{C}}$是一个上三角矩阵，如果我们能够证明$\mathbf{\mathit{B}}$会趋于一个上三角阵，那么由QR分解的唯一性便可得到$\mathbf{\mathit{V}}_{t} \rightarrow \mathbf{\mathit{V}}$。关于$\mathbf{\mathit{B}} = \mathbf{\Sigma}^{2 t} \mathbf{\mathit{V}}^{\top} \mathbf{\Sigma}^{- 2 t}$，用分量形式写出来是

$$ (\text{4}) B_{i , j} = V_{j , i} \left(\right. \sigma_{i} / \sigma_{j} \left.\right)^{2 t} $$

假设$\sigma_{1} > \sigma_{2} > \hdots > \sigma_{m}$，那么当$i > j$时$\left(\right. \sigma_{i} / \sigma_{j} \left.\right)^{2 t} \rightarrow 0$，即$\mathbf{\mathit{B}}$的下三角部分趋于0，换言之$\mathbf{\mathit{B}}$趋于上三角阵，条件得证。由此还可知，幂迭代的收敛速度取决于相邻奇异值之比，$\sigma_{i} / \sigma_{i + 1}$越大，收敛越快，一旦存在相等的奇异值，幂迭代就会失效，但实践中可以假设两个奇异值严格相等的概率为0，从而回避这个问题。

本质思考[#](https://spaces.ac.cn/archives/11710#%E6%9C%AC%E8%B4%A8%E6%80%9D%E8%80%83)

让我们仔细思考一下整个证明过程，理解一下它本质上在做什么。

首先，$\mathbf{\mathit{C}} = \mathbf{\Sigma}^{2 t} \mathbf{\mathit{R}}_{1}^{- 1} \mathbf{\mathit{R}}_{2}^{- 1} \hdots \mathbf{\mathit{R}}_{t}^{- 1}$这一块看起来复杂，但它的唯一作用仅仅是“一个上三角矩阵”，即它等于什么根本不重要，只要保持上三角矩阵的形式即可。真正发挥核心作用的是$\mathbf{\mathit{B}} = \mathbf{\Sigma}^{2 t} \mathbf{\mathit{V}}^{\top} \mathbf{\Sigma}^{- 2 t}$会趋于一个上三角矩阵，这使得前面的$\mathbf{\mathit{V}}$能够被QR分解提炼出来。

真正达到这个效果的，其实是$\left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right)^{t}$这一步！换句话说，理论上有$\mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right)^{t} \left.\right)$，即前面那么多次$\text{QR}$理论上都是多余的，只要在最后做一次$\text{QR}$或者说正交化就行了。当然，这种等价性只有理论价值，直接计算$\left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right)^{t}$会数值爆炸/坍缩，导致结果不可用。所以，定期执行$\text{QR}$（而不是最后才执行一次），本质作用就是保持数值稳定性。

我们也可以从“预处理器（Preconditioner）”的角度去理解多次QR的操作。预处理器指对输入的一些预处理操作，这些操作理论上不改变输出，但通常对数值计算有益。对于$\text{QR}$来说，右乘任意满秩的上三角阵$\mathbf{\mathit{R}}$并不会改变结果，即$\text{QR} ⁡ \left(\right. \mathbf{\mathit{A}} \left.\right) = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}} \mathbf{\mathit{R}} \left.\right)$，所以任意右乘的满秩上三角阵都可以称为$\text{QR}$的预处理器。

我们知道，$\text{QR}$本身也可以写成右乘上三角阵的形式，所以$\text{QR}$本身就是$\text{QR}$的一种预处理器，因此我们可以往$\left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right)^{t}$里边插入一些$\text{QR}$，变成

$$ (\text{5}) \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right)^{t} \rightarrow \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \hdots \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \left.\right) \left.\right) \left.\right) $$

最后两端都$\text{QR}$，结果不会改变，但由于每次$\text{QR}$出来一个正交矩阵，右乘一个正交矩阵几乎没有数值爆炸的风险，所以$\text{QR}$是它自身良好的预处理器。

相关变体[#](https://spaces.ac.cn/archives/11710#%E7%9B%B8%E5%85%B3%E5%8F%98%E4%BD%93)

从预处理器的角度，我们还可以理解很多幂迭代变体。比如，我们在开篇《基于流式幂迭代的Muon实现：1. 初识》就用过的$\text{ColNorm}$版：

$$ (\text{6}) \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) \rightarrow \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \text{ColNorm} ⁡ \left(\right. \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) \left.\right) $$

左右两个$\mathbf{\mathit{V}}_{t}$是相等的，因为对一个矩阵做$\text{ColNorm}$等价于右乘一个对角阵：

$$ (\text{7}) \underset{\text{ColNorm} ⁡ \left(\right. \left[\right. \mathbf{\mathit{a}}_{1} , \mathbf{\mathit{a}}_{2} , \hdots , \mathbf{\mathit{a}}_{m} \left]\right. \left.\right)}{\underbrace{\left[\right. \frac{\mathbf{\mathit{a}}_{1}}{\parallel \mathbf{\mathit{a}}_{1} \parallel} , \frac{\mathbf{\mathit{a}}_{2}}{\parallel \mathbf{\mathit{a}}_{2} \parallel} , \hdots , \frac{\mathbf{\mathit{a}}_{m}}{\parallel \mathbf{\mathit{a}}_{m} \parallel} \left]\right.}} = \left[\right. \mathbf{\mathit{a}}_{1} , \mathbf{\mathit{a}}_{2} , \hdots , \mathbf{\mathit{a}}_{m} \left]\right. \left[\right. \begin{matrix}\parallel \mathbf{\mathit{a}}_{1} \parallel^{- 1} & 0 & \hdots & 0 \\ 0 & \parallel \mathbf{\mathit{a}}_{2} \parallel^{- 1} & \hdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \hdots & \parallel \mathbf{\mathit{a}}_{m} \parallel^{- 1}\end{matrix} \left]\right. $$

对角阵是上三角阵的一个特例，所以$\text{ColNorm}$也是$\text{QR}$的一个预处理器，它不改变结果。基于同样的理由，我们也可以把$\text{ColNorm}$往外挪一步，这同样不改变$\text{QR}$的结果

$$ (\text{8}) \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) \rightarrow \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \text{ColNorm} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) \left.\right) $$

此外，《基于流式幂迭代的Muon实现：2. 加速》和《基于流式幂迭代的Muon实现：3. 雕琢》的主角是双重$\text{QR}$的幂迭代：

$$ (\text{9}) \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) \rightarrow \mathbf{\mathit{V}}_{t} = \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}}^{\top} \text{QR} ⁡ \left(\right. \mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1} \left.\right) \left.\right) $$

基于这两节的讨论，容易证明给$\mathbf{\mathit{A}} \mathbf{\mathit{V}}_{t - 1}$多加一步$\text{QR}$理论上也并不改变$\mathbf{\mathit{V}}_{t}$的结果。

有限误差[#](https://spaces.ac.cn/archives/11710#%E6%9C%89%E9%99%90%E8%AF%AF%E5%B7%AE)

我们多次用到了“理论上”之类的字眼，是因为很多结果都是在无限精度的、精确的$\text{QR}$基础上才能得到的，实际计算中的$\text{QR}$是有限精度的，而且为了效率我们通常还会用不那么准确的“SCQR（Shifted Cholesky QR）”， 所以分析误差的累积效应就非常有必要。

对于SCQR，大家这几篇读下来应该不再陌生，它将矩阵$\mathbf{\mathit{A}}$的QR分解分为两步：1、对$\mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} + \lambda \mathbf{\mathit{I}}$做Cholesky分解得到$\mathbf{\mathit{R}}$，2、通过$\mathbf{\mathit{Q}} = \mathbf{\mathit{A}} \mathbf{\mathit{R}}^{- 1}$得到正交矩阵$\mathbf{\mathit{Q}}$。如果$\lambda = 0$，那么SCQR跟精确QR等价，然而$\lambda = 0$在实际计算时几乎总会是因为条件数过大而失败，所以设置适当的$\lambda > 0$，这就带来了误差。

幸运的是，这种误差并不会累积！不管是$\mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}}$还是$\mathbf{\mathit{A}}^{\top} \mathbf{\mathit{A}} + \lambda \mathbf{\mathit{I}}$，它们做Cholesky分解的结果都是一个上三角矩阵，然后$\mathbf{\mathit{A}}$会右乘这个上三角阵的逆，我们知道右乘上三角阵不改变$\text{QR}$的结果。因此，如果不考虑矩阵乘法的误差，那么可以认为SCQR对$\text{QR}$本身是“无损”的。

换言之，只要最后一步换用精确的$\text{QR}$，那么不管前面做多少次不精准的SCQR，结果都还是精确的，反过来说，就是总的误差等价于单步SCQR（最后一步）的误差，而不会累积起来。依赖于这一原理的，还有我们在《基于流式幂迭代的Muon实现：2. 加速》提到的“SCQR2”加速技巧。

相反，如果我们所提的一些变换不能写成“右乘上三角阵”的形式，那么就会有损地破坏了原始矩阵的信息，长期迭代下来就会累积误差，直至完全失效。上篇《基于流式幂迭代的Muon实现：3. 雕琢》我们提到的 @Ji_Ha_Kim 同学的化简方案，就是一个经典反例。

乔列分解[#](https://spaces.ac.cn/archives/11710#%E4%B9%94%E5%88%97%E5%88%86%E8%A7%A3)

最后，我们再来简单复习一下“Cholesky分解”，它的目的是将给定正定对称矩阵$\mathbf{\mathit{B}}$表示成$\mathbf{\mathit{L}} \mathbf{\mathit{L}}^{\top}$的形式，其中$\mathbf{\mathit{L}}$是一个下三角阵。然后记$\mathbf{\mathit{R}} = \mathbf{\mathit{L}}^{\top}$，我们就可以得到上三角阵的形式$\mathbf{\mathit{R}}^{\top} \mathbf{\mathit{R}}$。

Cholesky分解很快，是因为它有可以直接递归计算的解析解。具体来说，从如下等式

$$ (\text{10}) \left[\right. \begin{matrix}a_{1 , 1} & a_{1 , 2} & \hdots & a_{1 , m} \\ a_{2 , 1} & a_{2 , 2} & \hdots & a_{2 , m} \\ \vdots & \vdots & \ddots & \vdots \\ a_{m , 1} & a_{m , 2} & \hdots & a_{m , m}\end{matrix} \left]\right. = \left[\right. \begin{matrix}l_{1 , 1} & & & \\ l_{2 , 1} & l_{2 , 2} & & \\ \vdots & \vdots & \ddots & \\ l_{m , 1} & l_{m , 2} & \hdots & l_{m , m}\end{matrix} \left]\right. \left[\right. \begin{matrix}l_{1 , 1} & l_{2 , 1} & \hdots & l_{m , 1} \\ & l_{2 , 2} & \hdots & l_{m , 2} \\ & & \ddots & \vdots \\ & & & l_{m , m}\end{matrix} \left]\right. $$

逐项对应相等，可解得递归公式

$$ (\text{11}) l_{j , j} = \pm \sqrt{a_{j , j} - \sum_{k = 1}^{j - 1} l_{j , k}^{2}} , l_{i , j} = \frac{1}{l_{j , j}} \left(\right. a_{i , j} - \sum_{k = 1}^{j - 1} l_{i , k} l_{j , k} \left.\right) 对于\textrm{ } i > j $$

Cholesky QR是Cholesky分解的经典应用之一，它跟施密特正交化类似，但允许我们跳过正交化的流程，直接得到三角矩阵$\mathbf{\mathit{R}}$。它们也面临着相同的数值困难：施密特正交化是逐向量正交化的，一旦遇到线性相关的向量或者零向量，那么正交化便无法进行下去，这在奇异值上体现为有效秩低或者条件数大，而在这种情况下Cholesky分解同样容易失败。

Cholesky分解的另一经典引用是正定对称矩阵求逆。比如解方程$\mathbf{\mathit{B}} \mathbf{\mathit{X}} = \mathbf{\mathit{C}}$，其中$\mathbf{\mathit{B}}$是正定对称的，那么我们可以先将它分解为$\mathbf{\mathit{L}} \mathbf{\mathit{L}}^{\top}$，然后变成$\mathbf{\mathit{L}} \left(\right. \mathbf{\mathit{L}}^{\top} \mathbf{\mathit{X}} \left.\right) = \mathbf{\mathit{C}}$，这只需要解两次三角系数方程，每一步都比较高效。其中，@Ji_Ha_Kim 所提的分式迭代求$\text{msign}$便是基于这个思路来求的逆矩阵。

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

这篇文章主要补充讨论了一下流式幂迭代的部分数学推导，包括幂迭代的收敛性、QR分解的预处理器、Cholesky分解的简单介绍等，希望能帮助大家从原理上更好地理解这个方法。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11710](https://spaces.ac.cn/archives/11710 "基于流式幂迭代的Muon实现：4. 原理")_

_更详细的转载事宜请参考：_[《科学空间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》")