提升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 来说，它是通过随机投影来找到矩阵