从谱范数梯度到新式权重衰减的思考

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

在文章《Muon优化器赏析：从向量到矩阵的本质跨越》中，我们介绍了一个名为 “Muon” 的新优化器，其中一个理解视角是作为谱范数正则下的最速梯度下降，这似乎揭示了矩阵参数的更本质的优化方向。

众所周知，对于矩阵参数我们经常也会加权重衰减（Weight Decay），它可以理解为 范数平方的梯度，那么从 Muon 的视角看，通过谱范数平方的梯度来构建新的权重衰减，会不会能起到更好的效果呢？

那么问题来了，谱范数的梯度或者说导数长啥样呢？用它来设计的新权重衰减又是什么样的？接下来我们围绕这些问题展开。

基础回顾

谱范数（Spectral Norm），又称 “2 范数”，是最常用的矩阵范数之一，相比更简单的 范数（Frobenius Norm），它往往能揭示一些与矩阵乘法相关的更本质的信号，这是因为它定义上就跟矩阵乘法相关：对于矩阵参数 ，它的谱范数定义为

这里 是列向量，右端的 是向量的模长（欧氏范数）。换个角度看，谱范数就是使得下面不等式对 恒成立的最小常数 ：

不难证明，当 取 范数 时，上式也是恒成立的，所以可以写出 （因为 只是让上式恒成立的其中一个 ，而 则是最小的那个 ）。这个结论也表明，如果我们想要控制输出的幅度，以谱范数作为正则项要比F范数更为精准。

早在 6 年前的《深度学习中的Lipschitz约束：泛化与生成模型》中，我们就讨论过谱范数，当时的应用场景有两个：一是 WGAN 对判别器明确提出了 Lipschitz 约束，而实现方式之一就是基于谱范数的归一化；二是有一些工作表明，谱范数作为正则项，相比 范数正则有更好的性能。

梯度推导

现在让我们进入正题，尝试推导谱范数的梯度 。我们知道，谱范数在数值上等于它的最大奇异值，对此我们在《低秩近似之路（二）：SVD》的“矩阵范数”一节有过证明。这意味着，如果 可以 SVD 为 ，那么

其中 是 的奇异值。对两边求微分，我们得到

留意到

同理 ，所以

注意，这个证明过程有一个关键条件是 ，因为如果 的话， 既可以表示成 又可以表示成 ，用同样方法求出的梯度分别是 和 ，结果不唯一意味着梯度不存在。当然，从实践角度看，两个数完全相等的概率是很小的，因此可以忽略这一点。

注：这里的证明过程参考了 Stack Exchange 上的回答 [1]，但该回答里面没有证明 和 ，这部分由笔者补充完整。

权重衰减

根据这个结果以及链式法则，我们有：

对比 范数下的结果：

这样对比着看就很清晰了： 范数平方作为正则项所得出的权重衰减，同时惩罚全体奇异值；而谱范数平方对应的权重衰减，只惩罚最大奇异值。如果我们目的是压缩输出的大小，那么压缩最大奇异值是“刚刚好”的做法，压缩全体奇异值虽然可能达到相近的目的，但同时也可能压缩参数的表达能力。

根据 “Eckart-Young-Mirsky 定理”，式 （7） 最右侧的结果还有一个含义，就是 矩阵的“最优1秩近似”。也就是说，谱范数的权重衰减将每一步减去它自身的操作，改为每一步减去它的最优 1 秩近似，弱化了惩罚力度，当然某种程度上也让惩罚更加“直击本质”。

数值计算

对于实践来说，最关键的问题来了：怎么计算 呢？SVD 当然是最简单直接的方案，但计算复杂度无疑也是最高的，我们必须找到更高效的计算途径。

不失一般性，设 。首先注意到：

由此可见计算 只需要知道 ，然后根据我们在《低秩近似之路（二）：SVD》中的讨论， 实际上是矩阵 的最大特征值对应的特征向量。这样一来，我们便将问题从一般矩阵 的 SVD 转化成了实对称矩阵 的特征值分解，这其实已经降低复杂度了，因为特征值分解通常要比 SVD 明显快。

如果还觉得慢，那么我们就需要请出很多特征值分解算法背后的原理——“幂迭代（Power Iteration）[2]”：

当 时，迭代

以 的速度收敛至 。

幂迭代每步只需要算两次“矩阵-向量”乘法，复杂度是 ， 步迭代的总复杂度是 ，非常理想，缺点是 接近时收敛会比较慢。但幂迭代的实际表现往往比理论想象更好用，早期很多工作甚至只迭代一次就得到不错的效果，因为 接近表明两者及其特征向量一定程度上可替换，而幂迭代即便没完全收敛，得到的也是两者特征向量的一个平均，这也完全够用了。

迭代证明

这一节我们来完成幂迭代的证明。不难看出，幂迭代可以等价地写成

为了证明这个极限，我们从 出发，代入计算可得

由于 是 的一组标准正交基，所以 可以写成 ，于是我们有

以及

由于随机初始化的缘故， 的概率是非常小的，所以我们可以认为 ，那么

当 时，所有的 都小于 1，因此当 时对应项都变成了 0，最后的极限是 。

相关工作

最早提出谱范数正则的论文，应该是 2017 年的《Spectral Norm Regularization for Improving the Generalizability of Deep Learning》[3]，里边对比了权重衰减、对抗训练、谱范数正则等方法，发现谱范数正则在泛化性能方面表现最好。

论文当时的做法，并不是像本文一样求出 ，而是直接通过幂迭代来估计 ，然后将 加权到损失函数中，让优化器自己去求梯度，这样做效率上稍差一些，并且也不好以权重衰减的形式跟优化器解耦开来。本文的做法相对来说更加灵活一些，允许我们像 AdamW 一样，将权重衰减独立于主损失函数的优化之外。

当然，从今天 LLM 的视角来看，当初的这些实验最大问题就是规模都太小了，很难有足够的说服力，不过鉴于谱范数的 Muon 优化器“珠玉在前”，笔者认为还是值得重新思考和尝试一下谱范数权重衰减。当然，不管是 范数还是谱范数的权重衰减，这些面向“泛化”的技术往往也有一些运气成份在里边，大家平常心期待就好。

个人在语言模型的初步实验结果显示，Loss 层面可能会有微弱的提升（希望不是幻觉，当然再不济也没有出现变差的现象）。实验过程就是用幂迭代求出 的近似值（初始化为全一向量，迭代 10 次），然后将原来的权重衰减 改为 ， 的取值不做改变。

文章小结

本文推导了谱范数的梯度，由此导出了一种新的权重衰减，并分享了笔者对它的思考。

（文：PaperWeekly）