©PaperWeekly 原创 · 作者｜苏剑林
单位｜追一科技
研究方向｜NLP、神经网络
众所周知，尽管基于 Attention 机制的 Transformer 类模型有着良好的并行性能，但它的空间和时间复杂度都是 近来，也有不少工作致力于降低 Transformer 模型的计算量，比如模型剪枝、量化、蒸馏等精简技术，又或者修改 Attention 结构，使得其复杂度能降低到 前几天笔者读到了论文 Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention [1] Attention （Linear Attention） 这个探索点，继而阅读了一些相关文献，有一些不错的收获，最后将自己对线性化 Attention 的理解汇总在此文中。
Attention
当前最流行的 Attention 机制当属 Scaled-Dot Attention [2] ，形式为： 这里的
本文我们主要关心 Self Attention 场景，所以为了介绍上的方便统一设 1.1 摘掉Softmax
读者也许想不到，制约 Attention 性能的关键因素，其实是定义里边的 Softmax！事实上，简单地推导一下就可以得到这个结论。
这一步我们得到一个 也就是说，去掉 Softmax 的 Attention 的复杂度可以降到最理想的线性级别 摘掉 Softmax 后的线形 Attention 。 问题是，直接去掉 Softmax 还能算是 Attention 吗？它还能有标准的 Attention 的效果吗？为了回答这个问题，我们先将 Scaled-Dot Attention 的定义（1）等价地改写为（本文的向量都是列向量）。 所以，Scaled-Dot Attention 其实就是以 也就是把 这种一般形式的 Attention 在 CV 中也被称为 Non-Local 网络，来自文章 Non-local Neural Networks [4] 。
几个例子
如果直接去掉 Softmax，那么就是 值得指出的是，下面介绍的这几种 Linear Attention，前两种只做了 CV 的实验，第三种是笔者自己构思的，所以都还没有在 NLP 任务上做过什么实验，各位做模型改进的 NLPer 们就有实验方向了。 2.1 核函数形式
一个自然的想法是：如果 其中 [5] 选择的是 非要讲故事的话，式（4）可以联想到“核方法（kernal method）”，尤其是 这方面的思考可以参考论文 Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel [6] ，此处不做过多延伸。 2.2 妙用Softmax
另一篇更早的文章 Efficient Attention: Attention with Linear Complexities [7] 则给出了一个更有意思的选择。它留意到在 其中 、 在这里，笔者给出自己的一种构思。这个构思的出发点不再是式（4），而是源于我们对原始定义（2）的近似。由泰勒展开我们有：
如果 这不同于形式（4），但理论上它应该是最接近原始的 Scaled-Dot Attention 了。
通过修改 Attention 的形式来降低它的计算复杂度，相关的工作有很多，这里简要列举一些。 3.1 稀疏Attention
我们之前介绍过 OpenAI 的 Sparse Attention ，通过“只保留小区域内的数值、强制让大部分注意力为零”的方式，来减少 Attention 的计算量。经过特殊设计之后，Attention 矩阵的非 0 元素只有 但是很明显，这种思路有两个不足之处：
如何选择要保留的注意力区域，这是人工主观决定的，带有很大的不智能性；
它需要从编程上进行特定的设计优化，才能得到一个高效的实现，所以它不容易推广。
3.2 Reformer Reformer 也是有代表性的改进工作，它将 Attention 的复杂度降到了 某种意义上来说，Reformer 也是稀疏 Attention 的一种，只不过它的稀疏 pattern 不是事先指定的，而是通过 LSH（Locality Sensitive Hashing）技术（近似地）快速地找到最大的若干个 Attention 值，然后只去计算那若干个值。 此外，Reformer 通过构造可逆形式的 FFN（Feedforward Network）替换掉原来的 FFN，然后重新设计反向传播过程，从而降低了显存占用量。 所以，相比前述稀疏 Attention，Reformer 解决了它的第一个缺点，但是依然有第二个缺点：实现起来复杂度高。要实现 LSH 形式的 Attention 比标准的 Attention 复杂多了，对可逆网络重写反向传播过程对普通读者来说更是遥不可及。 3.3 Linformer
跟本文所介绍的 Linear Attention 很相似的一个工作是 Facebook 最近放出来的 Linformer，它依然保留原始的 Scaled-Dot Attention 形式，但在进行 Attention 之前，用两个 这样一来， 但是，笔者认为“对于超长序列 m 可以保持不变”这个结论是值得质疑的，原论文中对于长序列作者只做了 MLM 任务，而很明显 MLM 并不那么需要长程依赖，所以这个实验没什么说服力。因此，Linformer 是不是真的 Linear，还有待商榷。 Linformer 的另一个缺点是 相比之下，本文介绍的几种 Linear Attention 都能做到这一点。以式（3）和式（4）为例，如果要 Mask 掉未来信息，那么只需要把求和 实现上式有两种方式：第一方式是设 这说明这种 Attention 可以作为一个 RNN 模型用递归的方式实现，它的空间复杂度最低，但是要串性计算，适合预测解码时使用；第二种是直接将 3.5 下采样技术
从结果上来看，Linformer 的 近来也有类似的工作发出来，比如IBM的PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination [8] 和 Google 的 Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing [9] 。
除了 Pooling 之外，其实还有其他的下采样技术，比如可以通过 stride > 1 的一维卷积来实现，基于这个思路，或许我们可以把 FFN 里边的 Position-Wise 全连接换成 stride > 1 的一维卷积？总之这方面应该也能玩出很多花样来，不过跟 Linformer 一样，这样糅合之后做自回归生成就很难了。
本文介绍了一些从结构上对 Attention 进行修改从而降低其计算复杂度的工作，其中最主要的 idea 是去掉标准 Attention 中的 Softmax，就可以使得 Attention 的复杂度退化为理想的 相比于其他类似的改进结构的工作，这种修改能在把复杂度降到 [1] https://arxiv.org/abs/2006.16236 [2] https://arxiv.org/abs/1706.03762 [3] https://kexue.fm/archives/7325 [4] https://kexue.fm/archives/1711.07971 [5] https://arxiv.org/abs/2006.16236 [6] https://arxiv.org/abs/1908.11775 [7] https://arxiv.org/abs/1812.01243 [8] https://arxiv.org/abs/2001.08950 [9] https://arxiv.org/abs/2006.03236
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读 ，也可以是学习心得 或技术干货 。我们的目的只有一个，让知识真正流动起来。
📝 来稿标准：
• 稿件确系个人原创作品 ，来稿需注明作者个人信息（姓名+学校/工作单位+学历/职位+研究方向）
• 如果文章并非首发，请在投稿时提醒并附上所有已发布链接
• PaperWeekly 默认每篇文章都是首发，均会添加“原创”标志
📬 投稿邮箱：
• 投稿邮箱： hr@paperweekly.site
• 所有文章配图，请单独在附件中发送
• 请留下即时联系方式（微信或手机），以便我们在编辑发布时和作者沟通
🔍
现在，在「知乎」 也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」 订阅我们的专栏吧
关于PaperWeekly
PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」 ，小助手将把你带入 PaperWeekly 的交流群里。