科学空间|Scientific Spaces

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

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

共轴等价

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

点击阅读全文...

回顾前两篇文章《基于流式幂迭代的Muon实现：1. 初识》和《基于流式幂迭代的Muon实现：2. 加速》，我们引入了Muon的流式幂迭代（Streaming Power Iteration）实现方案，初步验证了它的可行性，并进一步讨论了核心运算——QR分解——的加速，使其接近Newton-Schulz迭代实现的效率。

在这篇文章中，我们不再局限于优化单步的QR分解，而是从更整体的视角看待流式幂迭代，并结合具体的计算背景，对其实现细节做进一步的“精雕细琢”，尽可能减少计算瓶颈，使其效率趋近理论极限。

现有结果

流式幂迭代本质上是“边训练边SVD”，它的想法是通过幂迭代来求SVD，并通过缓存上一步的结果，将计算平摊到每一步训练上，使得在优化器中嵌入SVD成为可能。至于Muon，只不过是它的一个基本应用，因为Muon的核心运算$\newcommand{msign}{\mathop{\text{msign}}}\msign$最基本的实现方式就是SVD。具体来说，Muon的更新公式是

点击阅读全文...

在第一篇文章《基于流式幂迭代的Muon实现：1. 初识》中，笔者将流式幂迭代（Streaming Power Iteration）单独抽象出来，作为一种新的Muon实现方式。由于新方案是直接对SVD进行近似计算，所以相比基于Newton-Schulz迭代的标准实现，它具有更丰富的拓展空间，值得继续深入研究。

从计算上看，新方案的主要变化是Newton-Schulz迭代换成了$\newcommand{QR}{\mathop{\text{QR}}}\QR$分解，这带来了一些降速。上篇我们已经讨论了一些基本的加速手段，但尚未比肩标准实现。这篇文章我们继续研究$\QR$的加速，以求尽可能缩小差距。

流式迭代

我们将沿用第一篇文章的所有概念和记号，有相关疑惑的读者请先往前翻看一下。首先，Muon的更新公式是
\begin{equation}\newcommand{msign}{\mathop{\text{msign}}}\begin{aligned}
\boldsymbol{M}_t =&\, \beta\boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt]
\boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t [\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}] \\
\end{aligned}\end{equation}

点击阅读全文...

Muon的核心运算是$\newcommand{msign}{\mathop{\text{msign}}}\msign$，当前标准实现是Newton-Schulz迭代。不得不说，这确实是一个非常高效且GPU友好的算法，Muon能流行起来，起码有一大半是这个算法的功劳。然而，这个算法也给人一种“只此一家，别无分号”的感觉，因为它似乎就局限在算$\msign$了，一旦我们想要魔改Muon（比如$\msign$换成这里的$\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$），那么相应的计算就会变得麻烦起来。

本文提出一种新的实现思路——通过流式幂迭代（Streaming Power Iteration）来近似计算SVD。这并不是完全新的思路，而是已经出现之前的一些优化器工作中，但这里我们将它单独提炼出来，作为一个独立的算法使用。

内容回顾

Muon的细节我们就不展开了，大家自行翻看之前的文章如《Muon优化器赏析：从向量到矩阵的本质跨越》、《Muon续集：为什么我们选择尝试Muon？》、《Muon优化器指南：快速上手与关键细节》即可，这里直接给出它的公式：
\begin{equation}\begin{aligned}
\boldsymbol{M}_t =&\, \beta\boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt]
\boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t [\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}] \\
\end{aligned}\end{equation}

点击阅读全文...

前面我们在《通过msign来计算奇异值裁剪mclip（上）》讨论了奇异值裁剪$\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$的数值计算，核心思路来自 @leloykun 的文章《Numerically Stable Spectral Clipping Via Newton-Schulz Iteration》（现已重新修订和改名），通过寻找基于$\newcommand{msign}{\mathop{\text{msign}}}\msign$的表达式来避免另外寻找Newton-Schulz迭代，在文章中笔者提出了一个计算量更低的嵌套$\msign$方案。

不过前两天，@leloykun 在推特上指出笔者的方案实际计算中存在误差偏大的问题。本文来具体分析一下这个问题，并给出一个更高效、误差更低的新方案。

点击阅读全文...

前面我们用了两篇文章《msign算子的Newton-Schulz迭代（上）》和《msign算子的Newton-Schulz迭代（下）》讨论了矩阵的$\newcommand{msign}{\mathop{\text{msign}}}\newcommand{sign}{\mathop{\text{sign}}}\newcommand{clip}{\mathop{\text{clip}}}\newcommand{mclip}{\mathop{\text{mclip}}}\msign$算子的数值计算，这篇文章我们来关注“奇异值裁剪（Singular Value Clipping）”运算，它最近在 @_arohan_ 的推特上引起了热议，我们此前在《高阶MuP：更简明但更高明的谱条件缩放》也提到过，接下来我们简称为$\mclip$。

基本概念

对于标量$x$，$\clip$运算定义为
\begin{equation}\clip(x) = \max(\min(x, 1), -1) = \left\{\begin{aligned}1, &\quad x\geq 1 \\
x, &\quad x\in(-1, 1)\\
-1, &\quad x\leq -1
\end{aligned}\right.\end{equation}

点击阅读全文...

SVD（Singular Value Decomposition，奇异值分解）是常见的矩阵分解算法，相信很多读者都已经对它有所了解，此前我们在《低秩近似之路（二）：SVD》也专门介绍过它。然而，读者是否想到，SVD竟然还可以求导呢？笔者刚了解到这一结论时也颇感意外，因为直觉上“分解”往往都是不可导的。但事实是，SVD在一般情况下确实可导，这意味着理论上我们可以将SVD嵌入到模型中，并用基于梯度的优化器来端到端训练。

问题来了，既然SVD可导，那么它的导函数长什么样呢？接下来，我们将参考文献《Differentiating the Singular Value Decomposition》，逐步推导SVD的求导公式。

推导基础

假设$\boldsymbol{W}$是满秩的$n\times n$矩阵，且全体奇异值两两不等，这是比较容易讨论的情形，后面我们也会讨论哪些条件可以放宽一点。接着，我们设$\boldsymbol{W}$的SVD为：
\begin{equation}\boldsymbol{W} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}\end{equation}

点击阅读全文...

上一篇文章中我们介绍了“伪逆”，它关系到给定矩阵$\boldsymbol{M}$和$\boldsymbol{A}$（或$\boldsymbol{B}$）时优化目标$\Vert \boldsymbol{A}\boldsymbol{B} - \boldsymbol{M}\Vert_F^2$的最优解。这篇文章我们来关注$\boldsymbol{A},\boldsymbol{B}$都不给出时的最优解，即
\begin{equation}\mathop{\text{argmin}}_{\boldsymbol{A},\boldsymbol{B}}\Vert \boldsymbol{A}\boldsymbol{B} - \boldsymbol{M}\Vert_F^2\label{eq:loss-ab}\end{equation}
其中$\boldsymbol{A}\in\mathbb{R}^{n\times r}, \boldsymbol{B}\in\mathbb{R}^{r\times m}, \boldsymbol{M}\in\mathbb{R}^{n\times m},r < \min(n,m)$。说白了，这就是要寻找矩阵$\boldsymbol{M}$的“最优$r$秩近似（秩不超过$r$的最优近似）”。而要解决这个问题，就需要请出大名鼎鼎的“SVD（奇异值分解）”了。虽然本系列把伪逆作为开篇，但它的“名声”远不如SVD，听过甚至用过SVD但没听说过伪逆的应该大有人在，包括笔者也是先了解SVD后才看到伪逆。

接下来，我们将围绕着矩阵的最优低秩近似来展开介绍SVD。

结论初探

对于任意矩阵$\boldsymbol{M}\in\mathbb{R}^{n\times m}$，都可以找到如下形式的奇异值分解（SVD，Singular Value Decomposition）：
\begin{equation}\boldsymbol{M} = \boldsymbol{U}\boldsymbol{\Sigma} \boldsymbol{V}^{\top}\end{equation}

点击阅读全文...