nil’s blogs

Prologue 最近有一些工作认为使用行归一化对梯度矩阵做 Precondition，可以替代 Muon 的正交化达到接近甚至超越的结果。例如前段时间 Yiping 做的Row/Column Normalization and Hyperparameter Transfer。虽然论文给了一些有趣的理论结果，但是当时我对他们的实验 setting 是有一些质疑的。最近又有新的工作RMNP，声称使用行归一化达到了超越 Muon 的结果。 事实果真如此吗？按照笔者的经验，行归一化的效果其实跟 Adam 更接近，没有超越 Muon。本篇文章以 RMNP 为例来讨论下。 RMNP 假设梯度的 Momentum 矩阵是 $V_t$，Muon 使用： $$ \Delta W\propto (V_tV_t^\top)^{-1/2}V_t. $$ 而 RMNP 的作者认为，因为 $V_tV_t^\top$ 大概率是一个对角主导的矩阵（对角线元素远大于非对角线元素），因此可以用如下的近似： $$ \Delta W\propto \operatorname{diag}(V_tV_t^\top)^{-1/2}V_t. $$ 这等价于对 Momentum 矩阵每一行单独归一化。 这个假设不难理解，实际上就是非常经典的维度灾难的问题，我对作者的这个假设先按下不表。看看作者的实验效果如何： RMNP 实验效果 这里有个很有意思的结论，作者从“行归一化与正交化在高维下渐进等价”这个角度出发，用行归一化取代 Muon，取得了超越 Muon 的结果。这就好像「在模仿吴亦凡的比赛中吴亦凡获得了第二名」，听起来令人匪夷所思。 渐进等价的逻辑漏洞 RMNP 的核心逻辑是：因为 $$ V_tV_t^\top \approx \operatorname{diag}(V_tV_t^\top), $$ 所以可以用后者替代前者，效果差不多。这里其实有一个比较明显的漏洞，即便对称矩阵 $A$ 如果在 rms norm 的意义下近似于一个对角矩阵 $D$，使用 $A^{-1/2}$ 和 $D^{-1/2}$ 作为 preconditioner 的效果却可能差之千里。 ...

最近在细读Tensor Programs的时候，发现一个有意思的细节，Yang 2022.在将mup应用到Transformer的时候，对Attention Logits使用$1/d_k$进行放缩，而不是传统上的$1/\sqrt{d_k}$。 我不禁回想起很多年前在知乎的一个问题： transformer中的attention为什么scaled? - Nil-9的回答 当时笔者沿着Vaswani 2017.的一个脚注，展开写了这样一篇回答，其数学原理并不复杂，浓缩版本是这样的： 假设$\boldsymbol{q}$,$\boldsymbol{k}\in\mathbb{R}^{d_k}$的每个分量均是是独立采样的随机变量，均值和方差分别是0、1，$\forall i$满足分量$q_i,k_i$互相独立，则点积$\langle\boldsymbol{q},\boldsymbol{k}\rangle$的量级是$\Theta(\sqrt{d_k})$。为了防止点积过大导致softmax出现梯度消失问题，因此需要再点积基础上乘上$1/\sqrt{d_k}$，将点积稳定在$\Theta(1)$量级。 这个问题后来也成为了某些公司的所谓「机器学习八股文」之一。但这个放缩是不是严谨的呢？我们来仔细思考一下。 这里比较成问题的是「互相独立」的假设。我们先建立如下的基本的结论。 假设$X,Y$是两个随机变量，已知： $$ \mathbb{E}[X]= \mathbb{E}[Y] = 0\\ \text{Var}[X]=\text{Var}[Y] = 1 $$ 那么乘积$XY$满足： $$ \mathbb{E}[XY]= \rho $$ 如果进一步地，$X,Y$均为正态分布随机变量，则有 $$ \text{Var}[XY]= 1 + \rho^2 $$ 其中$\rho$是$X$和$Y$的相关系数。 proof ...

TL;DR 笔者通过自定义$XX^\top$算子，加速Muon优化器中的核心操作——Newton-Schulz迭代，算子单测在8192维度上相比原生实现计算时间降低约一半，端到端的Newton-Schulz迭代运行时间降低至原来的0.71倍。相关代码已经发布在： https://github.com/nil0x9/flash-muon 读者可以通过如下方式来试用优化版本的Muon实现或者核心算子。 git clone --recurse-submodules https://github.com/nil0x9/flash-muon.git pip install -e ./ 具体做法 Muon优化器的核心机制是通过Newton-Schulz迭代法来代替SVD分解，实现GPU友好的$\boldsymbol{U}\boldsymbol{V}^T$计算（之前笔者写过这篇blog介绍Muon背后的数学原理）。 Newton-Schulz迭代法的核心公式如下所示，通过指定合适的多项式系数$a,b,c$，将SVD分解的 $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^T$中的$\boldsymbol{\Sigma}$在迭代中消除 $$ \boldsymbol{X}\leftarrow a\boldsymbol{X} + b(\boldsymbol{X}\boldsymbol{X}^\top)\boldsymbol{X} + c(\boldsymbol{X}\boldsymbol{X}^\top)(\boldsymbol{X}\boldsymbol{X}^\top)\boldsymbol{X} $$ 其中$\boldsymbol{X}$是梯度矩阵或梯度的动量矩阵。在Jordan的实现 中这个迭代过程由如下的PyTorch代码实现： for _ in range(steps): A = X @ X.T B = b * A + c * A @ A X = a * X + B @ X 注意到这里提前计算了$\boldsymbol{A}=\boldsymbol{X}\boldsymbol{X}^\top$来减少重复计算。由于$\boldsymbol{A}$是一个对称矩阵，则循环内第二行的$\boldsymbol{AA} = \boldsymbol{AA}^\top$。这里我们可以看到一个被使用了两次的运算lambda x: matmul(x, x.T)，如果这个算子可以有比matmul更高效的实现，就可以实现更快的Newton-Schulz迭代。 幸运的是$\boldsymbol{X}\boldsymbol{X}^\top$是一个对称矩阵，理论上我们可以只计算上三角部分，下三角部分可以直接复用下三角部分的结果。 这一点在Laker Newhouse的这篇文章中有提及，但是可惜的是，他们并没有实现一个可用的kernel版本。 笔者沿着这个思路实现了一个可用的CUDA kernel。它的设计思路其实非常直观（也可能比较clumsy，欢迎批评指正）——在一般的GEMM算子中，每个线程会从全局内存load计算一定分块的矩阵算元到自己的寄存器内，在计算完对应分块的结果矩阵后，将这部分的结果store到输出算元对应的全局内存空间。GEMM的相关分析教程实在是太多了，笔者就不赘述了。——而对于$\boldsymbol{X}\boldsymbol{X}^\top$这种对称运算，笔者让下三角对应的线程块在kernel一开始直接退出，上三角部分的线程在计算结束后做完一般的store操作后，根据线程和线程块的id，找到对应的下三角部分的全局内存空间，将当前线程对应的结果分块转置后store到下三角的分块。 用图像表达如下： matmul_transpose的kernel设计 ...

Prologue 如何衡量神经网络参数空间中的距离（i.e., 选取合适的范数）？这个问题我们之前已经在这篇文章中有所涉及，本文是更完整的拓展。 首先应该明确为什么这个问题很重要？因为在梯度下降这样的优化算法框架下，总是依赖目标参数的某种距离的衡量。例如，当目标参数可以对应到欧几里得空间的坐标时，就可以得到局部目标函数下降最快的方向是梯度负方向。 尽管在深度学习优化中，直接使用欧几里得范数（衡量参数空间的距离）非常吸引人，并且实际上也有效。但是，这种做法会丢失结构信息，因为它实质上是将参数看做一个扁平的向量。那么得到的优化方向就可能是相对低效的。 对于常见的神经网络而言，对参数矩阵使用谱范数似乎是个更好的选择[1][2]。使用谱范数的相关约束，可以引出在现在的大规模模型训练中非常重要的两个工作： $\mu P$ (Maximal Update Parameterization)：使用特定的初始化与学习率的参数化，可以做到在小模型上调节部分超参数，迁移到同构的大模型。 Muon优化器的更新规则，即对梯度矩阵做SVD分解，消除奇异值矩阵后将$-\boldsymbol{U}\boldsymbol{V}^\top$作为对应参数的更新方向。 从而涵盖了给定神经网络结构下的参数初始化、参数更新的步长和方向三个方面（模型成功训练的三个基本要素）。 feature learning的谱条件 为什么需要使用谱范数来衡量参数的大小和参数更新步长呢？在Greg Yang的系列工作中，考虑的主要是feature learning的问题。Yang在他的TP4中发现，在标准参数化或者NTK参数化的设定下，无穷宽的神经网络没有办法学习特征（一次更新后特征与初始化状态无异），这就与传统的linear transfer learning的智慧形成了悖论，因为如果没有feature learning，那么预训练就对后续的transfer没有任何增益。 于是为了达成非平凡的feature learning，就需要对神经网络的特征$\boldsymbol{h}_\ell(\boldsymbol{x})\in\mathbb{R}^{n_\ell}$的范数与一步更新的特征范数变化量$\boldsymbol{h}_\ell(\boldsymbol{x})$做如下的约束[1]： $$ \|\boldsymbol{h}_\ell\|_2=\Theta(\sqrt{n_\ell})\text{ and } \|\Delta\boldsymbol{h}_\ell\|_2=\Theta(\sqrt{n_\ell}), \text{ for }\ell=1,\ldots,L-1\tag{1} $$ 这个条件意味着，对于中间的任何特征向量而言，每个元素平均下来的范数是常数阶的，一次更新后的变化量也是常数阶的（不会随着宽度$n_\ell$趋近无穷而爆炸或者弥散）。 对于常见的由一系列矩阵参数构成的神经网络而言，实现这样的约束，需要对参数矩阵做如下的谱条件约束： $$ \|\boldsymbol{W}_\ell\|_2=\Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right)\text{ and } \|\Delta\boldsymbol{W}_\ell\|_2=\Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right), \text{ for }\ell=1,\ldots,L-1\tag{2} $$ 这里的$\|\cdot\|_2$是矩阵的谱范数，即从向量$\ell_2$空间映射到向量$\ell_2$空间的算子诱导范数 $$ \|\boldsymbol{A}\|_2=\underset{\boldsymbol{x}\in\mathbb{R}^n}{\max} \frac{\|\boldsymbol{A}\boldsymbol{x}\|_2}{\|\boldsymbol{x}\|_2}\text{ for } \boldsymbol{A}\in\mathbb{R}^{m\times n} $$ 这个谱条件能够保证公式$(1)$成立，证明思路主要利用到谱范数的性质$\|\boldsymbol{Av}\|_2\le\|\boldsymbol{A}\|_2\|\boldsymbol{v}\|_2$，证明出上述谱条件诱导的$\|\boldsymbol{h}_\ell\|_2,\|\Delta\boldsymbol{h}_\ell\|_2$的上界均为$\Theta(\sqrt{n_\ell})$，接着证明这个上界依概率总是取得。具体可以看这篇论文的第三节（相比Greg的TP系列，还算是很容易理解的）。 公式$(2)$中的第二个条件是很直观的，对于SGD优化器而言，直接按照每个参数矩阵的维度设定对应的学习率即可 $$ \eta_\ell = \Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right) $$ 对于第一个条件$\boldsymbol{W}_\ell = \Theta\left(\sqrt{\frac{n_\ell}{n_{\ell-1}}}\right)$，需要修改初始化的方式——我们可以从一个标准正态分布中i.i.d.采样一个权重矩阵$\boldsymbol{W}'\in\mathbb{R}^{n_\ell\times n_{\ell-1}}$，接着做缩放$\boldsymbol{W} = \sigma\boldsymbol{W}'$（这里略去下标$\ell$）得到最终的初始化参数矩阵。 这个$\sigma$怎么确定呢？由于矩阵的stable rank按定义可以将矩阵的谱范数和Frobenius范数联系起来： $$ \text{stable-rank}(\boldsymbol{W}) = \frac{\|\boldsymbol{W}\|_F^2}{\|\boldsymbol{W}\|_2^2} $$ ...

上篇文章介绍了Muon等新兴深度学习优化器背后的原理，即约束参数矩阵的诱导范数下得到新的更新方向。 在Muon对参数更新方向$-\boldsymbol{U}\boldsymbol{V}^\top$的计算中用到了Newton-Schulz迭代方法，本质上是在寻找这样一个多项式函数 $$ f(x)=ax+bx^3+cx^5+\ldots $$ 使其满足对任意$x\in(0, 1]$，对$x$应用多次$f(\cdot)$，都能收敛到1附近。这里我们尝试设计一个能work的参数组合。 我的一个简单的想法是，设计一个多项式函数，使$x=1$是它的一个吸引不动点： 定义1（不动点）当$x_0$被函数$f(\cdot)$映射到自身，即$f(x_0)=x_0$时，称$x_0$是函数$f(\cdot)$的一个不动点。 定义2（吸引不动点）$f$的吸引不动点是$f$的不动点$x_0$使得，对在足够接近$x_0$的定义域中的任何$x$值而言，迭代函数序列$x,f(x),f(f(x)),f(f(f(x))),\ldots$收敛于$x_0$。 要令$x=1$是$f(x)$的一个吸引不动点，要满足如下的必要条件： $f(1)=1$ $|f'(1)|<1$ 使用这两个条件是无法确定具体的参数值$a,b,\ldots$的，但是对于三阶（参数包括$a,b$两个）或者五阶（参数包括$a,b,c$三个）的Newton-Schulz迭代，可以大大缩小搜索的空间。下面展开看下。 三阶迭代 先讨论三阶迭代的形式 $$ f(x)=ax+bx^3 $$ 代入上面的两个必要条件： $$ \begin{split} f(1)=a+b = 1\\ -1 < f'(1)=a+3b < 1\\ \end{split} $$ 根据第一个条件，可以把$b$用$1-a$重参数化，然后就有可行的条件 $$ 1 < a < 2 $$ 我们记五次迭代后的函数$\phi(x)=f(f(f(f(f(x)))))$，可视化看一下不同$a$取值下对应的情况（理想情况下，对于$(0,1]$区间内的$x$，曲线要尽可能接近$y=1$） 三阶迭代下，a取不同取值时对应的φ(x) 注意到在$a$接近1的时候，$\phi(x)$收敛到1附近的邻域是比较窄的，随着$a\to 2$，收敛到1附近的「邻域」范围逐渐拓宽，但在$a=2$附近，曲线开始出现一定的抖动。对于优化器而言，这样的局部近似的方差是可以容忍的，因此我们可以选取一个比较接近2的值作为$a$的参数，例如$a=1.99,b=-0.99$。 作为对比，在Bernstein & Newhouse 2024.中，作者给出的参数是$a=3/2,b=-1/2$。可以在下图中对照两种设定下的$\phi(x)$. 两种φ(x)对比 可以看到Bernstein给出的参数虽然更平滑地收敛于1，但是对于在0附近的初始$x$，普遍无法收敛到1。也就是说对于较小的奇异值对应的$\boldsymbol{u}_i, \boldsymbol{v}_i$，倾向于在更新中被忽略。 在$x=0$附近$\phi(x)$能否快速接近1，主要取决于参数$a$的大小，这是因为$\phi'(0)=a^5$。所以应该在尽可能保证$\phi(x)\approx 1,\forall x\in(0,1]$的同时，让$a$尽可能大。 五阶迭代 现在来考虑五阶迭代的形式 $$ f(x)=ax+bx^3+cx^5 $$ 代入上面的两个必要条件： $$ \begin{split} f(1)=a+b+c = 1\\ -1 < f'(1)=a+3b+5c < 1\\ \end{split} $$ ...