FLASH：可能是近来最有意思的高效Transformer设计

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

高效 Transformer，泛指所有概率 Transformer 效率的工作，笔者算是关注得比较早了，最早的博客可以追溯到 2019 年的《为节约而生：从标准 Attention 到稀疏 Attention》，当时做这块的工作很少。后来，这类工作逐渐多了，笔者也跟进了一些，比如线性 Attention、Performer、Nyströmformer，甚至自己也做了一些探索，比如之前的“Transformer 升级之路”。再后来，相关工作越来越多，但大多都很无趣，所以笔者就没怎么关注了。

▲ 本文模型脉络图

大抵是“久旱逢甘霖”的感觉，最近终于出现了一个比较有意思的高效 Transformer工作——来自 Google 的《Transformer Quality in Linear Time》[1] ，经过细读之后，笔者认为论文里边真算得上是“惊喜满满”了～

何喜之有

什么样的结果值得我们用“惊喜”来形容？有没有言过其实？我们不妨先来看看论文做到了什么：

1、提出了一种新的 Transformer 变体，它依然具有二次的复杂度，但是相比标准的 Transformer，它有着更快的速度、更低的显存占用以及更好的效果；

2、提出一种新的线性化 Transformer 方案，它不但提升了原有线性 Attention 的效果，还保持了做 Decoder 的可能性，并且做 Decoder 时还能保持高效的训练并行性。

说实话，笔者觉得做到以上任意一点都是非常难得的，而这篇论文一下子做到了两点，所以我愿意用“惊喜满满”来形容它。更重要的是，论文的改进总的来说还是比较自然和优雅的，不像很多类似工作一样显得很生硬。此外，笔者自己也做了简单的复现实验，结果显示论文的可复现性应该是蛮好的，所以真的有种“Transformer危矣”的感觉了。

门控注意

闲话少说，进入主题。我们知道标准的 Transformer 其实是 Attention 层和 FFN 层交替构建的，而这篇论文的核心是提出了一个融合了两者的新设计 GAU（Gated Attention Unit，门控注意力单元），它是新模型更快、更省、更好的关键，此外它使得整个模型只有一种层，也显得更为优雅。

威力初显

怎么做到 Attention 和 FFN 的融合呢？首先，标准的 FFN 是两层 MLP 模型：

这里 而 是激活函数。后来，《GLU Variants Improve Transformer》[2] 发现使用了 GLU（Gated Linear Unit，门控线性单元）的 FFN 效果更好，并为后来的 mT5 所用，其形式为：

这里 而 是逐位对应相乘（Hadamard 积）。GLU 更有效并不是一件让人意外的事情，早在 2017 年 Facebook 的《Convolutional Sequence to Sequence Learning》[3] 中 GLU 就起到了关键作用，此外笔者之前研究的 DGCNN 也肯定了 GLU 的有效性。

一般情况下的 GLU 是 不加激活函数而 加 Sigmoid，但这篇论文 都加了激活函数 Swish [4]（也叫 SiLU [5]，Sigmoid Linear Unit），这可以在附录中的源码找到，此处跟主流 GLU 用法略有不同，特别指出一下。

强强联合

既然 GLU 式的 FFN 更有效，那么我们就以它为基础进行修改。注意到 FFN 不能取代 Attention，是因为它的各个 token 之间没有进行交互，也就是矩阵 的每一行都是独立运算的。为了补充这点不足，一个自然的想法就是把 token 之间的联系补充到 上去，而为了体现出跟 Attetion 的结合，那么一个比较自然的设计就是

其中 是 Attention 矩阵，它负责融合 token 之间的信息。这样出来的 就包含了 token 之间的交互，原则上它可以取代 Attention。至于 怎么算，我们等会再说。

在式（3）中，如果 等于单位阵 ，那么它就是 GLU 式的 FFN；而如果 是全 1 矩阵，那么它就是普通的注意力机制。所以说，（3）是 Attention 和 FFN 的一个简单而自然的融合，我们期望它能同时替换掉 Attention 和 FFN，甚至有更好的表现。

弱化注意

刚才说了，GLU 本身就很强，不然 Facebook 也无法凭借 CNN+GLU 做到了当时 Seq2Seq 的 SOTA，而既然 GLU 那么强，那么一个猜测是它会弱化对 Attention 的依赖。也就是说，虽然在式（3）中 是不可或缺的，但或许我们可以简化它的形式。事实上确实如此，原论文使用了如下的简化版 Attention 矩阵：

这里 ， 即注意力的 head_size，文中取了 ，而 是简单的仿射变换（像 Layer Norm 中的乘 加 ）， 则是 后再平方。

跟标准的 Self Attention 类似，这里的注意力矩阵还是 的内积而来，复杂度还是 的，不同的是这里简化了 的来源变换，并且激活函数换用了 。大家可能对这个激活函数比较陌生，事实上这是作者在他之前的论文《Primer: Searching for Efficient Transformers for Language Modeling》[6] 用 NAS 的方式搜出来的。最后的 是简单的归一化因子，用以消除长度的影响。这个设计的成功也表明，注意力机制中的 softmax 不是必须的，可以换成常规的激活函数加简单的归一化。

▲ GAU示意图及其伪代码

以一当十

接下来请各位看官不要眨眼了，真正的“重磅”即将来袭了。可能 GLU 真的太强了，它对 Attention 的依赖真的非常非常弱，以至于作者们发现：只用一个头就够了！

▲ GAU与多头注意力的一些消融分析

我们知道标准的 Transformer 用的是多头注意力机制，在运算过程中需要产生 大小的矩阵， 是 batch_size 而 是头数，试想一下，当 、 甚至更大时， 已经够“惨”的了，还要活生生地乘个 ，不管对时间还是空间复杂度无疑都是“雪上加霜”。而如今，只要一个头的 GAU，就可以达到相同甚至更好的效果，不仅提高了计算速度，还降低了显存占用量，几乎算得上是“免费的午餐”了。

当 GAU 只有一个头时， 的参数量就很少了，主要参数量在 上，所以 GAU 的参数量大约为 ；而在标准的 Transformer 中，Attention 的参数量为 ，FFN 的参数量为 （标准 FFN 中一般是 ），所以总参数量为 。因此，从参数量看，当 时，两层 GAU 大致上就等于原来的 Attention+FFN。

所以，在 GAU 的实验中，作者都固定 ，那么“ 层 Attention+n 层 FFN”的标准 Transformer 模型，对应的就是“ 层GAU”的新模型，我们记为 FLASH-Quad，其中 Quad 是“Quadratic”的简写，表明复杂度依然是二次的，至于  FLASH 的含义，后面再谈。

高效线性

其实 FLASH-Quad 已经是标准 Transformer 的一个非常优秀的替代品了，但作者们还不满意其二次复杂度，继而提出了具有线性复杂度的 FLASH（Fast Linear Attention with a Single Head）。为此，作者提出了一种“分块混合注意力（Mixed Chunk Attention）”的方案，它不单可以用于前述 GAU 中，也可以用于标准的 Attention 中，是一种较为通用的线性化技巧。

现有方法

主流的高效 Transformer 工作对 Attention 的改进思路大体上可以两大类，分别是“稀疏化”和“线性化”。

本文开头提到的《为节约而生：从标准 Attention 到稀疏 Attention》，就是“稀疏化”的工作之一，后面诸如 Reformer [7] 等也算是此列，还有一些跟 Pooling 结合的如 Linformer [8] 也可以理解为广义的“稀疏化”。这类工作的特点是引入一定的归纳先验，强制大部分注意力为 0，从而理论上可以少减少计算量。但这种方案的缺点是往往需要专门的编程优化才能实现加速，或者是难以用来做 Decoder（Pooling 类工作），此外效果好坏比较依赖于其引入的归纳先验，显得不够自然。

至于“线性化”，我们在《线性 Attention的探索：Attention 必须有个 Softmax 吗？》有过介绍，研究的人相对多一些，后面的 Performer、Nyströmformer 以及最近的 cosFormer [9]、Flowformer [10] 都可以归入此类。简单来看，这类工作是将标准 Attention的 改为 从而实现了线性复杂度。

这类方法的好处是易于实现，但有两个主要问题，一是低秩性会导致效果明显变差（参考《Transformer 升级之路：从 Performer 到线性 Attention》）；另外是用来做 Decoder（Causal）时会牺牲训练并行性，因为它需要转化为 RNN 来计算，又或者不牺牲并行性，但需要 的空间复杂度，相比于标准 Attention 的 ，起码要 才有优势，而哪怕 ，都要 了，多数情况下不现实。

分块混合

FLASH 采取了“局部-全局”分块混合的方式，结合了“稀疏化”和“线性化”的优点。首先，对于长度为 的输入序列，我们将它不重叠地划分为 个长度为 的块（不失一般性，假设 能被 整除，论文取 ），设 为第 块，其中 的定义同前。跟式（4）一样，我们将 通过 4 个简单的仿射变换分别得到 。

其中 我们用来算块内的自注意力：

这代表的是每个块的 token 内部自行交互，本质上也算是“稀疏化”的一种，其复杂度大致是 ，正比于 。实现时相当于头数为 、序列长度为 的多头注意力，可以充分地并行，而如果想要做 Decoder，那么 mask 掉注意力矩阵的上三角部分即可。

剩下的 则用来做全局的 Attention，我们直接用前述线性 Attention 的方式来做：

注意，这个操作跟直接用完整矩阵 做线性 Attention 是完全等价的，写成这样只是更好地体现跟分块的联系。如果是做 Decoder，那么要防止泄漏未来信息，所以要改为 cumsum 形式：

这种情况下，为了保持并行性，我们只需要 的空间复杂度，而如果不分块直接用线性 Attention，那么是 （要是原始的用法还要加上多头，那就是 ），在当前参数设置下有 ，所以是更省显存了。

最后，将两种 Attention 结果结合起来，整合到 GAU 中，得到线性版本的 GAU

基于线性版本 GAU 搭建的 Transformer 模型，便是作者笔下的 FLASH 模型了。

一些讨论

笔者认为，之所以这样分块做“局部-全局”的混合注意力，除了是想降低计算成本外，还因为这样做能得到更贴合实际情况的注意力分布。按照我们对 NLP 的经验理解，自然语言中的关联主要还是集中在局部的，而全局的、极度长距离的关联虽然存在，但不会是主导地位，所以这种混合式的注意力设计更有利于模型凸出局部关联但不舍弃长程关联。原论文还做了消融实验，显示相对来说局部注意力比全局注意力更重要，而混合式的效果最好。

▲ 全局注意力和局部注意力的消融实验

此外，可能会有些读者担心这种非重叠的分块会不会不利于边界词的预测？原论文提到了这一点，它说引入更复杂的重叠式局部注意力确实有利于提升效果，但也引入了额外的计算成本，在增加同样计算成本的情况下，引入重叠式局部注意力带来的增益还不如直接多加几层目前的非重叠式 GAU。所以说，目前的非重叠足够好地平衡了速度和效果。

最后，这种“分块混合”的线性化方案本质上是通用的，它不仅可以用于 GAU 中，也可以用于标准的 Transformer 中，即保留标准的 Attention+FFN 组合，然后 Attention 用分块混合的方式进行线性化，原论文称之为“MC-TFM”，并也进行了相应的比较，结果显示 GAU 在线性化方面也显得更有优势。

实验分析

关于 GAU 和 FLASH 的实验结果，笔者认为最值得留意的有两个。

第一个是新设计的门控注意力单元 GAU 与标准的多头注意力之间 MHSA 的比较，其实也就是 FLASH-Quad 和标准 Transformer 的比较了，如下图：

▲ GAU与多头注意力的对比

注意横轴是速度，纵轴是效果，这种图越靠近右上角的点意味着越理想（速度和效果都最优），所以上图显示不管哪种规格的模型，GAU 都比相应的多头注意力模型更有优势。

第二个则是 FLASH 模型的实验表格：

▲ FLASH与标准Transformer的对比

该表格更直接地显示出：

1、尽管 FLASH-Quad 和 Transformer 都是二次复杂度，但 FLASH-Quad 效果更好、速度更快；

2、在序列足较长时，线性复杂度的 FLASH 比 FLASH-Quad 更快，并且效果相仿。

说实话，即便是 FLASH-Quad 这个依然是二次复杂度的模型的速度提升幅度，很多号称是线性复杂度的工作都未必能做到，GAU 的强大可见一斑。对了，论文还特别指出笔者之前提的旋转位置编码 RoPE 能明显提高 Transformer 和 FLASH 的效果，所以论文实验的 Transformer+、Transformer++、FLASH-Quad 和 FLASH 都是带有 RoPE 编码的，在此沾沾自喜一下。

另外，上述表格并没有给出显存占用的对比。事实上，笔者测试发现，在 base 量级和序列长度为 1024 时，FLASH-Quad可用的最大 batch_size 将近是 Transformer 的两倍，这意味着 FLASH-Quad 明显降低了显存消耗。同时，笔者简单尝试了 small 版本 FLASH-Quad 的中文预训练，发现效果甚至比 RoFormer（RoPE+Transformer）要好些，所以论文所报告的结果确实不虚。不过最近的卡有限，就没法进行更深入的测试了，以后有新结果再跟大家分享。

延伸思考

至此，对 GAU、FLASH 的介绍也基本结束了。到发博客时，作者还没有在 Gihub 上开放完整源代码，但是附录已经贴出了几乎可以直接抄来用的关键源码（tensorflow 版），所以代码的实现应但是没有困难的，有兴趣有算力的同学，可以自行参考实验。另外论文有什么读不懂的地方，也可以直接参考源代码。

下面进行“挑骨头”环节，说一下我觉得这篇论文还做的不够完美的地方。

首先，笔者认为 FLASH-Quad 和 FLASH 解耦得不够好。如本文开头的观点，FLASH-Quad 和 FLASH 都算得上是“重磅”级别的结果，甚至对笔者来说  FLASH-Quad 更有价值，因为自注意力的二次复杂度本身也带来了足够多的自由度，可以玩很多像 UniLM 这样的花样，所以 FLASH-Quad 本身应该是一个很独立、很值得肯定的模型，但在原论文中，它更像是 FLASH 的一个过渡产品，这我认为是过于“冷落”了 FLASH-Quad。幸好，作者单独分离出了 GAU 的概念，也算是缓解了这个不足。

然后，GAU 既可以代替 Attention，也可以代替 FFN，从设计上来看，它旨在代替的是 Self-Attention，作者似乎不关心它对 Cross Attention 的可代替性，论文也没有相应的实验。那么，GAU 是否有可能代替 Cross Attention 呢？从式（3）的形式看，理论上是有可能的，但不知道 GAU 代替 Cross Attention 时能否依然只保留一个头，因为只需一个头可谓是 GAU 替代 Self Attention 的最大亮点了，它是更快更省的关键。此外，论文只做了 LM 和 MLM 的语言模型实验，并没有做“预训练+微调”的实验，不确定 GAU 的迁移性能如何。或许等我有卡了，我也去补充一波实验。

最后，有一个笔者不大理解的地方，就是 GAU/FLASH-Quad/FLASH 同时用上了加性绝对、加性相对以及 RoPE 三种位置编码，理论上三者只用其一就行了，笔者自己做的 GAU 实验也只用 RoPE 但效果依然挺好，所以这里同时用三种有什么讲究吗？此外，从论文附录所给的源码看，作者并没有仔细处理好 padding 的问题，以及做Decoder是归一化因子递归也没有写好（前 项求和应该除以 而不是 ），这些都是不大不小的可改善的细节。当然，不排除作者的原始代码是正确的，附录只是出于可读性目的做了简化，因为附录里边的代码还是以“伪代码”自称。

本文小结

本文介绍了 Google 新出的一个高效 Transformer 工作，里边将 Attention 和 FFN 融合为一个新的 GAU 层，从而得到了 Transformer 变体 FLASH-Quad，作者还进一步提出了一种“分块混合”线性化方案，得到了具有线性复杂度的FLASH。目前的实验结果显示，不管 FLASH-Quad 还是 FLASH，跟标准 Transformer 相比都是更快、更省、更好。也许不久之后，All You Need 的就不再是 Attention 而是 GAU 了。

参考文献

[1] https://arxiv.org/abs/2202.10447

[2] https://arxiv.org/abs/2002.05202

[3] https://arxiv.org/abs/1705.03122

[4] https://arxiv.org/abs/1710.05941

[5] https://arxiv.org/abs/1606.08415

[6] https://arxiv.org/abs/2109.08668

[7] https://arxiv.org/abs/2001.04451

[8] https://arxiv.org/abs/2006.04768

[9] https://arxiv.org/abs/2202.08791

[10] https://arxiv.org/abs/2202.06258

特别鸣谢

感谢 TCCI 天桥脑科学研究院对于 PaperWeekly 的支持。TCCI 关注大脑探知、大脑功能和大脑健康。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·