科学空间|Scientific Spaces

通过前几篇文章的推导和计算，我们可以发现，第一篇《MuP之上：1. 好模型的三个特征》所提的三个稳定性指标通常可以分为“参数稳定性”和“增量稳定性”两部分，而在《MuP之上：2. 线性层与最速下降》和《MuP之上：3. 特殊情况特殊处理》中，我们演示了将增量稳定性与最速下降结合来获得新的更新规则（优化器）的过程.

然而，对于参数稳定性，我们之前只是停留在初始化上。这篇文章的任务，正是探讨如何在整个训练过程中维持参数的稳定性，将理论的实践补充完整。

问题背景

以《MuP之上：2. 线性层与最速下降》为例，三个稳定性指标分别是：

点击阅读全文...

本系列文章的主题是“流式幂迭代”，顾名思义，它由“流式”和“幂迭代”两部分构成，其中“幂迭代”是求矩阵SVD的一种经典的多步迭代方案，而“流式”则是指将原本需要多步迭代的算法平摊到每一步训练上，使得计算成本变得可以接受，其核心思想在于：与其一次性完成复杂计算，不如在训练过程中持续逼近目标。

作为该系列的延伸，本文将介绍另外一些“流式”思想的应用，进一步展示如何通过流式转化将相对昂贵的操作巧妙地融入训练流程。

正交投影

有些场景下，我们会希望约束某些参数矩阵的正交性。正交矩阵具有良好的数值稳定性，可以避免一些数值爆炸或消失问题，同时在某些设计中能带来更好的理论保证。当然，哪些地方适合约束参数为正交矩阵，我们需要具体场景具体分析，这里不做展开。

点击阅读全文...

经过《基于流式幂迭代的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}

点击阅读全文...

在上一篇文章《MuP之上：1. 好模型的三个特征》中，我们提出了前向稳定性、依赖稳定性、更新稳定性这三个核心指标，并给出了相应的数学定义。同时，我们提出以它们是否满足$\Theta(1)$来刻画一个模型的好坏，这将作为我们后续分析和计算的理论基石。接下来，我们将会把它们跟最速下降思想结合，给每个参数定制“稳中求快”的更新规则。

\begin{align}
&\text{前向稳定性:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c1} \\[5pt]
&\text{依赖稳定性:}\quad\max_{\boldsymbol{x}_1,\boldsymbol{x}_2} \Vert \boldsymbol{f}(\boldsymbol{x}_1;\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x}_2;\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c2} \\[5pt]
&\text{更新稳定性:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega} + \Delta\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c3}
\end{align}

我们以线性层作为第一个例子，其结果对部分读者来说应该不陌生，它就是去年逐渐兴起的Muon优化器。当然，我们的目的并不是重新发现Muon，而是展示从第一性原理出发来设计模型和优化器的过程，为我们后续处理其他参数提供统一的方法论。

点击阅读全文...

从机器学习时代的数据白化预处理，到深度学习时代的BatchNorm、InstanceNorm、LayerNorm、RMSNorm等花样百出的Normalization方法，本质上都体现了我们对“各向同性（Isotropy）”的偏爱。为什么我们会倾向于各向同性的特征呢？它有什么实际上的好处呢？这个问题能找到很多答案，比如对齐尺度、减少冗余、去相关性等等，但多是流于表面的感觉。

近日，笔者在读论文《The Affine Divergence: Aligning Activation Updates Beyond Normalisation》时，悟到了该问题在优化视角下的一个新理解，个人认为它相对来说还是比较贴近本质的，所以写出来跟大家分享和讨论一下。

最速下降

我们从最简单的线性层出发
\begin{equation}\boldsymbol{Y} = \boldsymbol{X}\boldsymbol{W}\end{equation}

点击阅读全文...