Transformer 是 RNN：具有线性注意力的快速自回归 Transformer

Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

安杰洛斯·卡塔罗普洛斯 Angelos Katharopoulos · Idiap Research Institute / EPFL · 2020-06-29 · arXiv:2006.16236 ↗ · 被引 2903

摘要 · Abstract

Transformer 在多项任务中表现出色，但由于其相对于输入长度的二次复杂度，处理极长序列时速度极慢。为解决这一局限，我们将自注意力表示为核特征映射的线性点积，并利用矩阵乘积的结合性，将复杂度从 O(N²)降至 O(N)，其中 N 为序列长度。我们证明，这种公式允许迭代实现，从而大幅加速自回归 Transformer，并揭示其与递归神经网络的联系。我们的线性 Transformer 性能与原始 Transformer 相当，且在极长序列的自回归预测中速度提升高达 4000 倍。

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\right)$ to $\mathcal{O}\left(N\right)$, 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.

核心贡献 · Key contributions

提出通过核特征图实现线性注意力，将复杂度从 O(N^2) 降至 O(N)。
Proposes linear attention via kernel feature maps, reducing complexity from O(N^2) to O(N).
推导出自回归 Transformer 的因果掩码，实现线性时间和恒定内存。
Derives causal masking with linear time and constant memory for autoregressive transformers.
揭示 Transformer 可视为循环神经网络，实现每步恒定时间推理。
Reveals transformers as recurrent neural networks, enabling constant-time inference per step.
在长序列自回归推理上实现高达 4000 倍加速，性能与原始 Transformer 相当。
Achieves up to 4000x faster autoregressive inference on long sequences with comparable performance.
通过累积和提供因果线性注意力的恒定内存梯度计算。
Provides constant-memory gradient computation for causal linear attention via cumulative sums.
在图像生成和语音识别上验证，展示线性缩放和竞争性精度。
Validates on image generation and speech recognition, showing linear scaling and competitive accuracy.

局限 · Limitations

特征图选择（如 elu）可能在某些任务中不如 softmax 注意力。
Feature map choice (e.g., elu) may not match softmax attention in all tasks.
线性注意力需要有限维特征图，限制了精确 softmax 近似。
Linear attention requires finite-dimensional feature maps, limiting exact softmax approximation.
在极短序列上，线性复杂度的优势可能因开销而不明显。
Performance on very short sequences may not benefit from linear complexity due to overhead.
该方法假设非负相似性函数，可能限制注意力的表达能力。
The method assumes non-negative similarity functions, which may restrict attention expressiveness.
实证评估仅限于图像和语音领域，对其他任务的泛化性未验证。
Empirical evaluation is limited to image and speech domains; generalizability to other tasks is unverified.

论文章节 · Sections（共 17）

摘要 Abstract
1 引言 1 Introduction
2 相关工作 2 Related Work
2.1 高效 Transformer 2.1 Efficient Transformers
2.2 理解自注意力 2.2 Understanding Self-Attention
2.3 线性化 softmax 2.3 Linearized softmax
3 线性 Transformer 3 Linear Transformers
3.1 Transformer 3.1 Transformers
3.2 线性化注意力 3.2 Linearized Attention
3.3 因果掩码 3.3 Causal Masking
3.4 Transformer 即 RNN 3.4 Transformers are RNNs
4 实验 4 Experiments
4.1 合成任务 4.1 Synthetic Tasks
4.2 图像生成 4.2 Image Generation
4.3 自动语音识别 4.3 Automatic Speech Recognition
5 结论 5 Conclusions
致谢 Acknowledgements

阅读逐段中英对照全文 →