[论文解读] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Angelos Katharopoulos, Apoorv Vyas+2 more|arXiv (Cornell University)|Jun 29, 2020

Neural Networks and Applications被引用 311

一句话总结

这篇论文为变换器提出线性注意力，以实现 O(N) 时间/内存，使自回归推理可快数千倍，同时在因果设置中与 RNN 等价。

ABSTRACT

Transformers achieve remarkable performance in several tasks but due to their quadratic complexity, with respect to the input's length, they are prohibitively slow for very long sequences. To address this limitation, we express the self-attention as a linear dot-product of kernel feature maps and make use of the associativity property of matrix products to reduce the complexity from $\mathcal{O}\left(N^2 ight)$ to $\mathcal{O}\left(N ight)$, where $N$ is the sequence length. We show that this formulation permits an iterative implementation that dramatically accelerates autoregressive transformers and reveals their relationship to recurrent neural networks. Our linear transformers achieve similar performance to vanilla transformers and they are up to 4000x faster on autoregressive prediction of very long sequences.

研究动机与目标

• 解决 Transformers 中自注意力对长序列的二次时间/内存成本。
• 提出一种基于核的线性注意力形式，利用结合性将复杂度降至 O(N)。
• 实现线性复杂度和常量内存的因果掩码。
• 通过将 Transformers 重新表述为 RNN，展示自回归推理的加速。
• 在图像、语音和合成任务上，实证比较线性变换器与 Softmax 与 Reformer。

提出的方法

• 用基于特征映射 phi 的核注意力替代 Softmax 注意力，使每层的计算复杂度为 O(N)。
• 利用结合性将注意力改写为 phi(Q) (phi(K)^T V) 并为跨查询复用一次性计算求和。
• 用线性复杂度和常量内存强制实现因果掩码，逐步更新累计和 S_i 和 Z_i。
• 推导前向和反向的线性时间梯度，以在训练中保持线性内存使用。
• 推导出具有因果掩码的变换器可以被看作具有两个记忆(s 和 z)的 RNN，以及递归更新方程。
• 提供 PyTorch 实现和 CUDA 加速梯度；在图像、音频和合成任务上与 Softmax 和 Reformer 进行对比。

实验结果

研究问题

• RQ1 在不牺牲性能的前提下，能否将变换器中的自注意力重构为线性时间和内存？
• RQ2 线性注意力在自回归序列生成的因果掩码下的表现如何？
• RQ3 在因果掩码下，变换器能否被解释为 RNN，这是否使推理更快？
• RQ4 与完整的变换器和最先进的替代方法相比，线性变换器在视觉和语音任务上是否保持竞争力？

主要发现

方法	比特/维度	图像/秒
Softmax	0.621	0.45
LSH-1	0.745	0.68
LSH-4	0.676	0.27
Linear (ours)	0.644	142.8
（表格续表用于 CIFAR-10）
Softmax	3.47	0.004
LSH-1	3.39	0.015
LSH-4	3.51	0.005
Linear (ours)	3.40	17.85
（表格续表用于 WSJ）
Bi-LSTM	-	1047
Softmax	-	2711
LSH-4	-	2250
Linear (ours)	-	824

• 线性注意力将时间和内存复杂度从 O(N^2) 降至每层 O(N)。
• 采用合适的特征映射（基于 elu 的），线性变换器在测试任务上达到与完整变换器相同的性能。
• 自回归推理变得快数千倍，因为内部状态 (S_i, Z_i) 可以增量更新。
• 实验显示 MNIST 和 CIFAR-10 图像生成在近 Softmax 的困惑度下，但吞吐量大幅提升（数百至数千倍更快）。
• 在 WSJ ASR 中，线性变换器在 PER 和训练速度方面超过 LSTM 和 Reformer 基线，而 Softmax 仍然最准确但更慢。
• 论文揭示一个形式上的联系：具有因果掩码的变换器可以被视为具有两个记忆组件的 RNN。