Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

文章标题：Transformer即RNN：采用线性注意力的快速自回归Transformer
作者/机构：Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François Fleuret

---

A1 主要贡献

核心问题：
标准的Transformer模型虽然在多种任务中表现出色，但其核心组件自注意力（self-attention）的计算和内存复杂性与输入序列长度 N 呈二次方关系，即 $O(N^2)$。这使得处理非常长的序列时，其计算成本过高、速度过慢，从而限制了模型的上下文长度，影响了时间连贯性和捕捉长期依赖的能力。尽管现有的一些高效Transformer方法（如稀疏分解、局部敏感哈希）在训练上降低了复杂性，但它们并未加速自回归推理过程。

研究目标：
本文旨在提出一种线性Transformer模型，该模型能够显著减少内存占用，并将计算复杂度降低至与上下文长度呈线性关系，即 $O(N)$。同时，该模型需要能大幅提升自回归推理的速度。

创新点：
1. 线性化自注意力机制：作者将自注意力表达为核函数特征图（kernel feature maps）的线性点积。通过利用矩阵乘法的结合律，成功将计算复杂度从 $O(N^2)$ 降低到 $O(N)$。
2. 高效的因果掩码：本文提出了一种适用于线性化注意力的因果掩码（causal masking）实现，其同样具有线性的复杂度和恒定的内存占用。
3. 揭示Transformer与RNN的关系：该线性化框架揭示了自回归Transformer与循环神经网络（RNN）之间的内在联系。基于此，作者将Transformer层重写为RNN形式，使其能够在自回归推理任务中实现数千倍的速度提升。

---

A3 背景知识/关键Observation/设计原则

2.1. 高效的Transformer

• 传统效率提升方法：现有工作通过权重剪枝【21, Michel et al. Are sixteen heads really better than one? 2019. NeurIPS】、权重分解【17, Lan et al. Albert: A lite bert for self-supervised learning of language representations. 2020. ICLR】、权重分层或知识蒸馏来提高Transformer的内存效率。例如，Clark等人【4, Clark et al. ELECTRA: Pre-training text encoders as discriminators rather than generators. 2020. ICLR】提出了新的预训练目标以提高样本效率；Lample等人【16, Lample et al. Large memory layers with product keys. 2019. NeurIPS】使用乘积键注意力来增加层容量。然而，这些方法虽然能带来训练或推理的加速，但并未改变自注意力与序列长度的二次方复杂度，这阻碍了其扩展到长序列。本文的方法则从根本上降低了渐进复杂度和内存占用。
• 上下文扩展方法：另一类研究致力于扩展Transformer的“上下文”范围。例如，Transformer-XL【7, Dai et al. Transformer-XL: Attentive language models beyond a fixed-length context. 2019. ACL】通过学习超越固定长度上下文的依赖关系，实现了语言建模的SOTA，但代价是显著增加了计算成本。Sukhbaatar等人【35, Sukhbaatar et al. Adaptive attention span in transformers. 2019. ACL】通过学习每个注意力头的最佳注意力跨度来扩展上下文长度。这两种方法都没有改变模型的渐进复杂度。相比之下，本文通过改进渐进复杂度，使得模型能够使用更大的上下文。
• 与本文更相关的工作：与本文模型最相关的是Child等人【3, Child et al. Generating long sequences with sparse transformers. 2019. arXiv】和Kitaev等人【14, Kitaev et al. Reformer: The efficient transformer. 2020. arXiv】的工作。前者引入了注意力矩阵的稀疏分解，将复杂度从二次方降至 $O(N\sqrt{N})$。后者提出了Reformer，使用局部敏感哈希（LSH）将复杂度进一步降至 $O(N \log N)$。但Reformer为了使用LSH，要求键（keys）与查询（queries）相同，这使得它无法用于键和查询必须不同的解码任务。相比之下，线性Transformer对查询和键没有限制，与序列长度呈线性扩展，并且能在自回归任务中实现三个数量级的推理加速。

2.2. 理解自注意力

• 基于核的视角：已有研究从理论上理解自注意力。Tsai等人【37, Tsai et al. Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel. 2019. EMNLP-IJCNLP】提出了基于核的注意力公式，将注意力视为在输入上应用核平滑器，其中核分数是输入间的相似度。这种公式有助于更好地理解注意力组件和整合位置嵌入。本文则利用核公式来加速自注意力的计算并降低其计算复杂度，并发现如果对查询和键应用具有正相似性分数的核函数，线性注意力可以正常收敛。
• 自注意力与卷积/RNN的关系：最近，Cordonnier等人【6, Cordonnier et al. On the relationship between self-attention and convolutional layers. 2020. ICLR】证明了具有足够多头的多头自注意力可以表示任何卷积层。本文则展示了，一个使用自回归目标训练的自注意力层可以被看作一个循环神经网络（RNN），这一发现可用于显著加速自回归Transformer模型的推理时间。

2.3. 线性化的Softmax

• Softmax瓶颈：多年来，Softmax一直是训练具有大量类别的分类模型的瓶颈。近期工作，如Blanc和Rendle【2, Blanc & Rendle. Adaptive sampled softmax with kernel based sampling. 2017. arXiv】以及Rawat等人【29, Rawat et al. Sampled softmax with random fourier features. 2019. NeurIPS】，通过特征图的线性点积来近似Softmax，以通过采样加速训练。受这些工作的启发，本文将Transformer中的Softmax注意力线性化。与本文同期的工作，Shen等人【31, Shen et al. Efficient attention: Attention with linear complexities. 2020. arXiv】也探索了线性化注意力在图像目标检测任务中的应用。相比之下，本文不仅线性化了注意力计算，还开发了一个在训练和推理上都具有线性复杂度和恒定内存的自回归Transformer模型，并从核的角度揭示了每个Transformer都可以被视为一个RNN。

---

A2 方法细节

本节形式化地提出了线性Transformer。通过将传统的softmax注意力改为基于特征图的点积注意力，实现了更好的时间和内存复杂度，并获得了一个能像RNN一样以线性时间进行序列生成的因果模型。

3.1. Transformer

• Transformer层级结构：一个Transformer定义为一个函数 $T: R^{N \times F} \rightarrow R^{N \times F}$，由L个Transformer层 $T_1(\cdot), \dots, T_L(\cdot)$ 组合而成。
  

  

  
  其中函数 $f_l(\cdot)$ 独立地转换每个特征，通常由一个小型两层前馈网络实现。$A_l(\cdot)$ 是自注意力函数，是Transformer中唯一跨序列操作的部分。

• 标准自注意力：自注意力函数 $A_l(\cdot)$ 为每个位置计算所有其他位置特征表示的加权平均值，权重与表示之间的相似度得分成正比。具体地，输入序列x通过三个矩阵 $W_Q \in R^{F \times D}$, $W_K \in R^{F \times D}$, $W_V \in R^{F \times M}$ 投影得到相应的表示Q, K, V。所有位置的输出 $A_l(x) = V'$ 计算如下：
  

  

  
  注意，上式中的softmax函数是按行应用于 $QK^T$ 的。Q、K、V通常被称为“查询”、“键”和“值”。

• 通用注意力形式：我们可以为任何相似性函数编写一个通用的注意力方程，如下所示，其中用i索引矩阵表示取其第i行作为一个向量：
  

  

• Softmax注意力的等价性：如果我们将相似性函数替换为 $sim(q, k) = exp(\frac{q^Tk}{\sqrt{d_k}})$，那么方程3就等价于方程2。

3.2. 线性化注意力

• 基于核的注意力：方程2中的注意力定义是通用的，可用于定义其他注意力实现，如多项式注意力或RBF核注意力【37, Tsai et al. Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel. 2019. EMNLP-IJCNLP】。为了让方程3定义一个注意力函数，我们只需要对 $sim(\cdot)$ 施加非负约束，这包括了所有的核函数 $k(x, y): R^{2 \times F} \rightarrow R^+$。

• 核函数重写注意力：给定这样一个带有特征表示 $\phi(x)$ 的核函数，我们可以将方程2重写为：
  

  

• 利用矩阵乘法结合律简化：通过利用矩阵乘法的结合律，我们可以进一步将其简化为：
  

  

• 分子向量化形式：上述方程的分子以向量化形式书写时更易理解：
  

  

  
  注意，特征映射 $\phi(\cdot)$ 是按行应用于矩阵Q和K的。

• 复杂度分析：从方程2可以看出，softmax注意力的计算成本与序列长度N的平方成正比，即 $O(N^2)$。内存需求也是如此，因为必须存储完整的注意力矩阵来计算关于查询、键和值的梯度。相比之下，我们提出的线性Transformer（方程5）的时间和内存复杂度为 $O(N)$，因为我们可以计算一次 $\sum_{j=1}^{N} \phi(K_j)V_j^T$ 和 $\sum_{j=1}^{N} \phi(K_j)$，并对每个查询重用它们。

3.2.1. 特征图与计算成本

• 成本对比：对于softmax注意力，总的乘法和加法成本与 $O(N^2 \max(D, M))$ 成比例，其中D是查询和键的维度，M是值的维度。相反，对于线性注意力，我们首先计算维度为C的特征图，随后计算新的值需要 $O(NCM)$ 的加法和乘法。
• 核函数选择：上述分析未考虑核函数和特征函数的选择。指数核对应的特征函数是无限维的，使得精确的softmax注意力线性化不可行。而多项式核则有精确的有限维特征图，并且已被证明与指数或RBF核效果相当【37, Tsai et al. Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel. 2019. EMNLP-IJCNLP】。一个2次多项式线性化Transformer的计算成本为 $O(ND^2M)$，这在 $N > D^2$ 时具有优势，这在实践中是成立的，因为我们希望能够处理数万个元素的序列。
• 本文使用的特征图：在处理较短序列的实验中，我们采用了一个能产生正相似性函数的特征图，定义如下：
  

  

  
  其中 $elu(\cdot)$ 是指数线性单元激活函数【5, Clevert et al. Fast and accurate deep network learning by exponential linear units (ELUs). 2015. arXiv】。我们更偏好 $elu(\cdot)$ 而非 $relu(\cdot)$，以避免在x为负时梯度为0。这个特征图产生的注意力函数需要 $O(NDM)$ 次乘法和加法。实验部分表明，该特征图的性能与完整Transformer相当，同时显著降低了计算和内存需求。

3.3. 因果掩码

• 自回归模型的因果掩码：Transformer架构可以通过掩蔽注意力计算来高效训练自回归模型，使得第i个位置只能受位置j（当且仅当 $j \le i$）的影响。形式上，这种因果掩码将方程3改变为：
  

  

• 线性化因果掩码：遵循3.2节的思路，我们将掩码注意力线性化如下：
  

  

• 引入累加状态变量：通过引入 $S_i$ 和 $Z_i$ 如下：
  

  

  
  

  

• 简化后的线性因果注意力：我们可以将方程9简化为：
  

  

  
  注意，$S_i$ 和 $Z_i$ 可以在常数时间内从 $S_{i-1}$ 和 $Z_{i-1}$ 计算得出，从而使得带有因果掩码的线性Transformer的计算复杂度与序列长度呈线性关系。

3.3.1. 梯度计算

• 内存优化问题：在任何深度学习框架中，对方程12的朴素实现都需要存储所有中间值 $S_i$ 来计算梯度，这会使内存消耗增加 $\max(D, M)$ 倍，从而阻碍了因果线性注意力在更长序列或更深模型中的应用。
• 累加和形式的梯度：为了解决这个问题，我们将方程9中分子的梯度推导为累加和的形式。这使得我们能够在线性时间和常数内存下完成因果线性注意力的前向和后向传播。

• 梯度公式：给定分子 $\bar{V}_i$ 和标量损失函数L相对于分子的梯度 $\nabla_{\bar{V}} L$，我们可以推导出 $\nabla_{\phi(Q_i)}L$, $\nabla_{\phi(K_i)}L$ 和 $\nabla_{V_i}L$ 如下：
  

  

  
  

  

  
  

  

• 复杂度和实现：方程9和13-15中的累加和项可以在线性时间内计算，并且相对于序列长度只需要常数内存。对于给定的C维特征图，这产生了一个计算复杂度为 $O(NCM)$、内存为 $O(N\max(C, M))$ 的算法。算法1中给出了分子前向和后向传播的伪代码实现。

3.3.2. 训练与推理

• 训练与推理的差异：在训练自回归Transformer模型时，完整的真实序列是可用的，这使得层级并行成为可能，因此Transformer比RNN训练效率更高。然而，在推理时，时间步i的输出是时间步i+1的输入，这使得自回归模型无法并行化，并且Transformer每个时间步的成本随着当前序列长度的平方增长。
• 线性Transformer的优势：我们提出的线性Transformer模型结合了两者的优点。在训练时，计算可以并行化并充分利用GPU。在推理时，我们模型的每次预测所需的时间和内存是恒定的。我们可以简单地将 $\phi(K_j)V_j^T$ 矩阵存储为内部状态，并在每个时间步像RNN一样更新它。这使得推理速度比其他Transformer模型快数千倍。

3.4. Transformer即RNN

• Transformer与RNN的联系：在文献中，Transformer模型被认为是与RNN根本不同的方法。然而，从3.3节的因果掩码公式和上一节的讨论中可以明显看出，任何具有因果掩码的Transformer层都可以写成一个模型：给定一个输入，修改一个内部状态，然后预测一个输出——即一个循环神经网络（RNN）。这与Universal Transformer【8, Dehghani et al. Universal transformers. 2018. arXiv】不同，我们考虑的是关于时间的递归，而不是深度。
• RNN形式化：我们将方程1的Transformer层形式化为一个RNN。得到的RNN有两个隐藏状态，即注意力内存s和归一化内存z。下标表示递归中的时间步。
  

  

  
  

  

  
  

  

  
  

  

  
  

  

• 通用性和意义：在上述方程中，$x_i$ 表示第i个输入，$y_i$ 表示特定Transformer层的第i个输出。我们的公式对特征函数没有任何限制，理论上可以用于表示任何Transformer模型，甚至包括使用softmax注意力的模型。这个公式是朝着更好地理解Transformer与流行的循环网络（如LSTM【12, Hochreiter & Schmidhuber. Long short-term memory. 1997. Neural computation】）之间关系的第一步。

---

A4 实验环境

• 对比模型：

  • softmax: 带softmax注意力的标准Transformer【38, Vaswani et al. Attention is all you need. 2017. NIPS】。
  • Reformer (lsh-X): SOTA的加速Transformer架构【14, Kitaev et al. Reformer: The efficient transformer. 2020. arXiv】，使用PyTorch重现版本，不使用可逆层。X表示哈希轮数。
  • linear: 本文提出的线性Transformer。

• 软件配置：

  • linear模型使用公式7的特征图。
  • 代码基于PyTorch【25, Paszke et al. Pytorch: An imperative style, high-performance deep learning library. 2019. NeurIPS】实现，恒定内存梯度计算部分由约200行CUDA代码实现。
  • 优化器：RAdam【19, Liu et al. On the variance of the adaptive learning rate and beyond. 2019. arXiv】或Adam【13, Kingma & Ba. Adam: A method for stochastic optimization. 2014. arXiv】。

• 硬件配置：

  • 合成任务基准测试：NVIDIA GTX 1080 Ti (11GB显存)。
  • CIFAR-10实验：NVIDIA P40 (24GB显存)。

• 实验任务与模型参数：

  • 合成任务（序列复制）：
    • 数据集：最大长度128，包含10种不同符号。
    • 模型：4层，8个注意力头。
    • 训练参数：批大小64，RAdam优化器，学习率从$10^{-3}$降至$10^{-4}$。
  • 图像生成（MNIST）：
    • 数据集：MNIST，序列长度784。
    • 模型：8层，8个注意力头，嵌入维度256（每头32维），前馈网络维度是嵌入维度的4倍。输出层使用10个逻辑斯蒂混合模型【30, Salimans et al. Pixelcnn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. 2017. arXiv】。
    • 训练参数：批大小10，RAdam优化器，学习率$10^{-4}$，训练250个epoch。
  • 图像生成（CIFAR-10）：
    • 数据集：CIFAR-10，序列长度3072。
    • 模型：16层，配置同MNIST。
    • 训练参数：softmax批大小1，linear和reformer批大小4。所有模型训练7天。
  • 自动语音识别（ASR）：
    • 数据集：WSJ (80小时)【26, Paul & Baker. The design for the wall street journal-based csr corpus. 1992. workshop on Speech and Natural Language】，使用40维mel-scale filterbanks特征。序列平均长度800，最长2400。
    • 对比模型：双向LSTM（3层，隐藏层大小320）。
    • Transformer模型：9层，6个注意力头，嵌入维度配置同图像实验。
    • 训练参数：LSTM使用Adam优化器，Transformer使用RAdam优化器，学习率均根据验证集错误率进行调整。

---

A4 实验结果

4.1. 合成任务

• 收敛性分析 (图2)：在线性Transformer的收敛性分析中，它在序列复制任务上的表现稳定，最终达到的损失与softmax模型相同，并且优于由哈希引入噪声的lsh模型。
  

  

• 内存与计算需求 (图1)：通过对不同长度的序列进行前向/后向传播测试，结果表明：softmax模型的内存和时间消耗随序列长度呈二次方增长；而linear和Reformer模型则呈线性增长。在所有配置下，linear模型都比基线模型速度更快、内存占用更少。
  

  

4.2. 图像生成

• MNIST (表1, 图3)：linear模型在性能（bits/dim）上与softmax模型几乎持平（0.83 vs 0.82），但在图像生成吞吐量上快了300多倍（142张/秒 vs 0.45张/秒）。这是因为其恒定的内存需求使其能在单GPU上同时生成10000张图像。从图3的样本来看，生成图像清晰，图像补全效果令人信服。
  

  

  
  

  

• CIFAR-10 (表2, 图4)：由于序列更长，linear模型的优势更加明显。在固定的7天训练时间内，linear模型完成了softmax模型3倍的训练轮数（150 vs 49），并取得了更好的困惑度（3.14 vs 3.26）。在生成吞吐量上，linear模型比softmax快4000多倍（17.85张/秒 vs 0.004张/秒）。图4展示了模型生成的图像具有空间一致性，且能令人信服地完成图像补全。
  

  

  
  

  

4.3. 自动语音识别

• WSJ数据集上的性能 (表3)：在非自回归的ASR任务中，使用CTC损失进行评估。linear模型在音素错误率（PER）和训练速度上均大幅优于LSTM和Reformer基线。softmax模型虽然取得了最低的PER（7.3%），但其训练速度明显慢于其他模型，linear模型每个epoch的训练速度是其3倍以上（290秒 vs 900秒）。
  

  

---

A5 结论

本文提出了线性Transformer，通过利用矩阵乘法的结合律，成功地将自注意力的计算和内存成本降低到与序列长度呈线性关系。研究表明，该模型可以与因果掩码结合使用，并保持其线性的渐进复杂度。最终，本文将Transformer模型表达为一种循环神经网络（RNN），使其能够在自回归任务上实现数千倍的推理加速。

未来工作展望：
1. 深入研究Transformer与RNN的关系：这一特性为未来研究RNN和Transformer中信息的存储与检索机制开辟了多种方向。
2. 探索新的特征图：另一个值得探索的研究方向是为线性注意力选择不同的特征图。例如，使用随机傅里叶特征来近似RBF核，可能允许我们直接使用以softmax注意力预训练的模型。

---

A6 附录

A. 梯度推导

• 目标：本节详细推导了因果掩码线性Transformer的梯度，并证明它们可以在线性时间和常数内存中计算。具体地，是推导一个标量损失函数相对于以下方程分子的梯度：
  

  

  
  分母和整个分数的梯度可以由自动求导系统高效处理。不失一般性，我们可以假设Q和K已经包含了由$\phi(\cdot)$映射过的向量。

• 分子定义：给定分子
  

  

  
  和 $\nabla_{\bar{V}} L$，目标是计算 $\nabla_Q L, \nabla_K L$ 和 $\nabla_V L$。

• 梯度推导过程：首先，将上述方程以非向量形式表示单个元素：
  

  

  
  随后，通过对任意 $Q_{lt}$ 求偏导来推导Q的梯度：
  

  

  
  将上式写成梯度矩阵的乘积形式，即证明了主论文中的方程13：
  

  

  
  对于K和V，由于 $K_j$ 会影响所有 $i \ge j$ 的 $\bar{V}_i$，因此需要对i进行求和。对 $K_{lt}$ 的偏导如下：
  

  

  
  其向量化形式证明了主论文中的方程14：
  

  

  
  对V的梯度推导过程类似，证明了方程15。值得注意的是，Q和K的梯度累加和矩阵大小相同，但一个是从1到N的正向累加，另一个是从N到1的反向累加，类似于RNN中的BPTT。

B. 训练过程

• 训练曲线对比（图5）：
  • MNIST (图5a)：所有模型训练250个epoch。由于序列较短，训练时间差异不大。linear模型的收敛情况与softmax相当，显著优于两种reformer变体。
  • CIFAR-10 (图5b)：所有模型训练固定的7天时间。lsh-1和linear模型完成的epoch数远多于softmax和lsh-4，并取得了更好的性能。随着序列长度进一步增加，这种差距预计会更大。
  • 语音识别 (图5c)：softmax在收敛性上显著优于Reformer和linear。但linear每个epoch的速度快3倍，这意味着它完成的epoch数大约是softmax的4倍。尽管softmax在该任务上更优，但linear在收敛和最终性能上均显著优于Reformer。
    

    

C. 图像生成吞吐量讨论

• C.1. 有状态的softmax注意力：为了进行更公平的比较，本节创建了一个名为stateful-softmax的基线，它将softmax自回归transformer实现为一个循环模型，即缓存所有的键和值。这个循环模型的状态大小与序列长度成正比，这与我们提出的状态大小固定的模型有本质区别。如表4所示，stateful-softmax比普通transformer快得多，但其复杂度仍然是二次的。对于CIFAR-10，我们的linear模型仍然比它快50多倍。
  

  

• C.2. 统一批次大小：本节评估了生成单个图像的延迟（latency），即将批次大小设为1。结果如表5所示，所有方法在GPU上都未充分利用资源，吞吐量远低于表4。linear transformer在所有方法中速度最快，在CIFAR-10上生成一张图片比softmax快近6.6倍。一个有趣的现象是，linear模型在CPU上的运行速度在所有情况下都比GPU快，这是因为其RNN式的注意力计算成本极低，主要的计算瓶颈变成了不可避免的序列外层循环。
  

  

D. 图像生成的定性结果

• 样本质量：本节提供了图像生成实验的定性结果。由于所有模型的困惑度大致相同，生成的样本在质量上没有显著差异。一个有趣的观察是，Reformer模型在无条件样本中的变化明显较少。此外，所有模型在图像补全任务上的表现都比无条件生成要好得多，表明这是一个相对容易的任务。