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，是因为它的