提升Transformer效率又有新招？基于矩阵分解的线性化Attention方案
©PaperWeekly 原创 · 作者｜苏剑林
单位｜追一科技
研究方向｜NLP、神经网络
标准 Attention 的 复杂度可真是让研究人员头大。前段时间我们在文章 Performer：用随机投影将 Attention 的复杂度线性化中介绍了 Google 的 Performer 模型，它通过随机投影的方式将标准 Attention 转化为线性 Attention。无独有偶，前些天 Arxiv 上放出了 AAAI 2021 的一篇论文，里边又提出了一种从另一个角度把标准 Attention 线性化的方案。
论文标题：
Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention
论文链接：
https://arxiv.org/abs/2102.03902
代码链接：
https://github.com/mlpen/Nystromformer
该方案写的是 Nyström-Based，顾名思义是利用了 Nyström 方法来近似标准 Attention 的。但是坦白说，在看到这篇论文之前，笔者也完全没听说过 Nyström 方法，而纵观整篇论文，里边也全是笔者一眼看上去感觉很茫然的矩阵分解推导，理解起来颇为困难。
不过有趣的是，尽管作者的推导很复杂，但笔者发现最终的结果可以通过一个相对来说更简明的方式来理解，遂将笔者对 Nyströmformer 的理解整理在此，供大家参考。
简单的回顾
如果读者对线性 Attention 还不是很了解，那么建议先通读一下线性 Attention 的探索：Attention 必须有个 Softmax 吗？和 Performer：用随机投影将 Attention 的复杂度线性化。总的来说，线性 Attention 是通过矩阵乘法的结合律来降低 Attention 的复杂度。
1.1 标准Attention
标准的 Scaled-Dot Attention 写成矩阵形式就是（有时候指数部分还会多个缩放因子，这里我们就不显式写出来了）：
这里 （对应 Self Attention）。此外，本文的所有 softmax，都是对矩阵的第二个维度做归一化。
在上式中， 这一步必须要先算出来，然后才能算 softmax，它导致了我们不能使用矩阵乘法的结合律。而 是 个向量的内积，因此时间和空间复杂度都是 。
1.2 线性Attention
而线性 Attention 比较朴素的做法就是：
其中 是值域非负的激活函数。为了方便对比，上式还没有显式地写出归一化因子，只突出了主要计算量的部分。上式左端的复杂度依然是 的，由于矩阵乘法满足结合律，我们可以先算后面两个矩阵的乘法，这样整体复杂度就降为 了。
上式是直接将 Attention 定义为两个矩阵的乘法来利用乘法结合律的，也可以将标准 Attention（近似地）转化为矩阵的乘法来利用结合律，如下一节提到的 Performer；此外，相乘矩阵也不一定是两个，比如本文要介绍的 Nyströmformer 就是将注意力表示为三个矩阵相乘的。
1.3 Performer
对于 Performer 来说，它是通过随机投影来找到矩阵 使得 softmax 中的 ，这样一来标准 Attention 就可以近似为上一节的线性 Attention 来算了，细节请看之前的文章 Performer：用随机投影将 Attention 的复杂度线性化。
如果对 SVM 和核方法等比较熟悉的读者可能会联想到，这个做法其实就是核函数的思想，即低维空间中两个向量的核函数可以映射为高维空间中两个向量的内积。它也可以跟 LSH（Locality Sensitive Hashing）联系起来。
Nyströmformer
在这部分内容中，我们以一个简单的双重 softmax 形式的线性 Attention 为出发点，逐步寻找更加接近标准 Attention 的线性 Attention，从而得到 Nyströmformer。
2.1 双重Softmax
在文章线性 Attention 的探索：Attention 必须有个 Softmax 吗？中我们提到了一种比较有意思的线性 Attention，它使用了双重 softmax 来构建 Attention 矩阵：
可以证明这样构造出来的 Attention 矩阵自动满足归一化要求，不得不说这是一种简单漂亮的线性 Attention 方案。
不过，直接对 做 softmax 似乎有点奇怪，总感觉没有经过相似度（内积）对比就直接 softmax 会有哪里不对劲。为了解决这个问题，Nyströmformer 先分别将 视为 n 个 d 维向量，然后聚成m类来得到 m 个聚类中心构成的矩阵 ，这时候我们可以通过下述公式来定义 Attention：
具体的聚类过程我们稍后再来讨论。现在，softmax 的对象是内积的结果，具有比较鲜明的物理意义，因此可以认为上式比前面的式 (3) 更为合理。如果我们选定一个比较小的 m，那么上式右端的复杂度只是线性地依赖于 n，因此它也是一个线性 Attention。
2.2 向标准靠近
纯粹从改进式 (3) 的角度来看，式 (4) 已经达到目标了，不过 Nyströmformer 并不局限于此，它还希望改进后的结果与标准 Attention 更加接近。
为此，观察到式 (4) 的注意力矩阵 是一个 的矩阵乘以一个 的矩阵，为了微调结果，又不至于增加过多的复杂度，我们可以考虑在中间插入一个 的矩阵 ：
如何选择 呢？一个合理的要求是当 m=n 时应当完全等价于标准 Attention，此时 ，推出：
对于一般的 m， 恰好是一个 矩阵，因此选它作为 至少在矩阵运算上是合理的，而根据 m=n 时的特殊情况我们则“大胆地”推测选它作为 能让新的 Attention 机制更接近标准 Attention，因此 Nyströmformer 最终选择的是：
作为 Attention 矩阵，它是三个小矩阵的乘积，因此通过矩阵乘法的结合律就能转化为线性 Attention。
不过，还有一个理论上的小细节需要补充一下，那就是上式涉及到矩阵的求逆，而 未必是可逆的。当然，从实践上来看，一个实数的方阵不可逆的概率几乎为零（不可逆意味着行列式严格等于 0，从概率上来看不等于 0 自然比等于 0 的概率大得多）。
因此这种情况在具体实验中可以不考虑，但理论上还是得完善的。这个其实也简单，如果是不可逆的矩阵，那就换成“伪逆”就好（记号为 ），它对任意矩阵都存在，并且当矩阵可逆时伪逆跟逆相等。
因此，最终的 Nyströmformer 的 Attention 矩阵形式为：
2.3 迭代求逆阵
从理论上看，式 (8) 已经达到目标了，不过落实到实践上还需要处理好一些细节问题，比如上述伪逆怎么求。伪逆又叫广义逆、Moore-Penrose 逆等，标准的求法是通过 SVD 来求，设矩阵 的 SVD 分解为 ，那么它的伪逆为：
其中对角阵 的伪逆 等于将它对角线所有非零值取倒数所得到的新对角阵。SVD 的求法虽然理论上比较简单易懂，但计算量还是比较大的，而且也不容易求梯度，因此并不是实现伪逆的理想方式。
Nyströmformer采用了迭代求逆的近似方法。具体来说，它采用了论文《Chebyshev-type methods and preconditioning techniques》[1] 提供的迭代算法：
若初始矩阵 满足 ，那么对于下述迭代格式：
成立 。
这里的 可以是任意一种矩阵范数，满足条件的一个比较简单的初始值可以是：
在 Nyströmformer 论文中，作者直接用上述初始值和迭代格式进行迭代，将迭代 6 次的结果来代替 。迭代 6 次看上去很多，但事实上论文所选取的 m 比较小（论文写的是 64），迭代过程中又只涉及到矩阵乘法，因此迭代计算量不会太大，而且只有乘法的话求梯度就很轻松了。这样求伪逆的问题就算是解决了，论文将这个迭代过程简写为 pINV。
2.4 池化当聚类
还需要解决的另一个问题是聚类方法的选择，比较直接的想法自然就是直接套用 K-Means 了。然而，同前面求伪逆所面临的问题一样，在设计模型时不仅要考虑前向计算，还需要考虑反向传播的求梯度，直接套用 K-Means 涉及到 操作，无法求出有意义的梯度，需要将它“软化”才能嵌入到模型中。
这一系列操作下来，其实就相当于胶囊网络的“动态路由”过程，细节我们在再来一顿贺岁宴：从 K-Means 到 Capsule 讨论过。这个方案的主要问题是 K-Means 是一个迭代过程，需要迭代几次才能保证效果，这导致计算量明显加大，不是特别理想。
Nyströmformer 选了一个非常简单的方案：假设序列长度 n 是 m 的整数倍（如果不是，padding 零向量），那么将 的每 n/m 个向量求平均作为 的每个向量。这个操作叫做 Adaptive Average Pooling（原论文称为 Segment-Means，简称 sMEANS），即是一种平均池化方法，通过自适应窗口大小使得平均池化后的特征矩阵具有固定的形状。
Nyströmformer 的实验表明，不需要比较复杂的聚类方法，就这样使用简单的自适应池化就可以取得非常有竞争力的效果了，而且只需要选择 m=64，跟映射前的 d 是一般大小，这比 Performer 要选择比 d 大几倍的 m 要好得多了。
不过，自适应池化的一个明显缺点是会“糅合”每一个区间的信息，导致它不能防止未来信息泄漏而不能做自回归生成（语言模型或者 Seq2Seq 的解码器），这基本是任何带有 Pooling 技术的模型的缺点。
实验与分析
这里我们汇总一下 Nyströmformer 的实验结果，并且分享一下笔者对它的一些看法和思考。
3.1 性能与效果
可能受限于算力，原论文做的实验不算特别丰富，主要是将 small 和 base 版本的 BERT 里边的标准 Attention 替换为 Nyströmformer 进行对比实验，实验结果主要是下面两个图。
其中一个是预训练效果图，其中比较有意思的是 Nyströmformer 在 MLM 任务上的效果比标准 Attention 还要优；另外是在下游任务上的微调效果，显示出跟标准 Attention（即 BERT）比还是有竞争力的。
▲ Nyströmformer在下游任务的微调效果
不过，原论文并没有比较 Nyströmformer 跟同类模型的效果差异，只是提供下面的一张复杂度对比图，因此无法更好地突出 Nyströmformer 的竞争力：
▲ 不同模型的时间和空间复杂度对比图
3.2 个人的思考
总的来说，Nyströmformer 对标准 Attention 进行近似线性化的思路还是比较新颖的，值得学习与参考。不过伪逆部分的处理总感觉有点不大自然，这部分可能是未来的一个改进点，如果可以做到不用近似，那就比较完美了。还有，如何定量地估计 Nyströmformer 与标准 Attention 的误差，也是一个值得思考的理论问题。
从实验上来看，Nyströmformer 跟标准 Attention 相比还是显得有竞争力的，尤其是 MLM 的结果比标准 Attention 还好，显示了 Nyströmformer 的潜力。此外，前面说到包含了 Pooling 导致不能做自回归生成是 Nyströmformer 的一个显著缺点，不知道有没有办法可以弥补，反正笔者目前是没有想到好的方向。
跟 Performer 相比，Nyströmformer 去除了线性化过程中的随机性，因为 Performer 是通过随机投影来达到线性化的，这必然会带来随机性，对于某些有强迫症的读者来说，这个随机性可能是难以接受的存在，而 Nyströmformer 则不存在这种随机性，因此也算是一个亮点。
3.3 Nyström方法
可能有些读者还是想学习一下 Nyström 方法，这里稍微补充一下。要理解 Nyström 方法，需要先简单认识一下矩阵的 CUR 分解。
大家可能都听说过矩阵的 SVD 分解，格式为 ，其中 是正交矩阵而 是对角矩阵。要注意 是正交矩阵意味着它们是稠密的，那么当 很大的时候 SVD 的计算成本和储存成本都很大（哪怕是做了近似）。
现在假设 很大但很稀疏，那么它的 SVD 分解比原始矩阵还不划算得多得多。为此，CUR 分解应运而生，它希望从原矩阵中选择 k 列组成矩阵 、选择 k 行组成矩阵 ，并插入一个 的矩阵 ，使得：
由于 都是原句子的一部分，因此也继承了稀疏性。关于 CUR 分解，读者还可以参考斯坦福的 CS246 课程的《Dimensionality Reduction》[2]一节。跟 SVD 不同的是，CUR 分解在笔者看来更多的是一种分解思想而不是具体的分解算法，它有不同的实现方式，比如 Nyström 方法也算是其中一种，分解形式为
其中 和 选出来的列矩阵和行矩阵，这里为了方便描述，假设了经过排列后选出来的行列均排在矩阵前面。Nyströmformer 其实也没有直接用 Nyström 方法（事实上也直接套用不了，原论文有描述），而是借鉴了 Nyström 方法的分解思想而已。
关于 Nyström 方法，原论文主要引用的是《Improving CUR Matrix Decomposition and the Nyström Approximation via Adaptive Sampling》[3]，但并不推荐新手读这篇论文，而推荐读《Matrix Compression using the Nystro ̈m Method》[4] 和《Using the Nyström Method to Speed Up Kernel Machines》[5]。
要特别说明的是，对于 CUR 分解和 Nyström 方法，笔者也是新学的，可能有理解不当的地方，请读者自行甄别理解，也欢迎熟悉相关理论的读者交流指正。
来一个小结
本文介绍了提升 Transformer 效率的一个新工作 Nyströmformer，它借鉴了 Nyström 方法的思想来构建一个能逼近标准 Attention 的线性 Attention，类似思想的工作还有 Performer，两者相比各有自己的优缺点，都是值得学习的工作。
本文分享了笔者自己对 Nyströmformer 的理解，窃认为这种途径更加易懂一些，如有谬误，肯请读者指正。
参考文献
[1] https://www.researchgate.net/publication/220562466_Chebyshev-type_methods_and_preconditioning_techniques
[2] https://web.stanford.edu/class/cs246/slides/06-dim_red.pdf
[3] https://arxiv.org/abs/1303.4207
[4] https://arxiv.org/abs/1305.0203
[5] https://www.researchgate.net/publication/49459305_Using_the_Nystroem_Method_to_Speed_Up_Kernel_Machines
更多阅读
#投 稿 通 道#
让你的论文被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学习心得或技术干货。我们的目的只有一个，让知识真正流动起来。
???? 来稿标准：
• 稿件确系个人原创作品，来稿需注明作者个人信息（姓名+学校/工作单位+学历/职位+研究方向）
• 如果文章并非首发，请在投稿时提醒并附上所有已发布链接
• PaperWeekly 默认每篇文章都是首发，均会添加“原创”标志
???? 投稿邮箱：
• 投稿邮箱：[email protected]
• 所有文章配图，请单独在附件中发送
• 请留下即时联系方式（微信或手机），以便我们在编辑发布时和作者沟通
????
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
关于PaperWeekly
PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」，小助手将把你带入 PaperWeekly 的交流群里。
