长短期 Transformer ：用于语言和视觉的高效 Transformer

2.1 高效Transformer

近年来，已经引入了许多方法来处理完全注意力机制的二次成本。 一般来说，它们可以分为以下几类： i) 具有预定义模式的稀疏注意力机制（例如，滑动窗口），包括用于建模图像的Sparse Transformer [12] 、Image Transformer [13] 、Axial Transformer [25] ，以及用于建模语言的Longformer [14] 、分块自注意力 [26] 、ETC [15] 、Big Bird [16] 。 ii) 低秩投影注意力，包括Linformer [17] 、Nyströmformer [18] 、Synthesizer [19] 。 例如，Linformer使用线性层将原始高分辨率键( K )和值( V )（长度为 n ）投影到大小为 r 的低分辨率( r ≪ n )，并允许所有查询符元( Q )关注这些压缩表示。 iii) 基于记忆的机制，例如压缩Transformer [10] 和集合Transformer [27] ，它们使用额外的存储器来缓存全局长程信息，用于计算远距离符元之间的注意力。 iv) 注意力矩阵的基于核的近似方法，包括Performer [28] 、线性Transformer [29] 和随机特征注意力 [30] 。 vi) 基于相似性和聚类的方法，包括Reformer [31] 、路由Transformer [32] 和Sinkhorn Transformer [33] 。

我们的方法无缝地整合了低秩投影和局部窗口注意力，以利用它们的优势来建模长程和短期相关性。 特别地，我们的长程注意力使用动态低秩投影来编码输入序列，并且优于Linformer [17] 使用的先前低秩投影方法。 同样，其他一些方法也尝试结合不同方法的优势。 例如，Longformer [14] 和ETC [15] 用任务驱动的全局符元增强局部窗口注意力。 这种全局符元可能不适用于某些任务（例如，自回归建模）。 BigBird [16] 进一步结合了局部窗口和全局符元注意力以及随机稀疏注意力。 它不适用于自回归任务，因为引入了全局符元和随机稀疏模式。 为了压缩边缘设备上的模型占用空间，Lite Transformer [34] 结合了卷积和自注意力，但对于长序列，它仍然具有二次复杂度。

2.2 视觉Transformer

视觉Transformer (ViT) [4] 将图像分割成小的图像块，并将这些图像块视为输入词符。 它使用标准的Transformer进行图像分类，并且在有足够训练数据的情况下，已被证明优于卷积神经网络（例如，ResNet [35] ）。 DeiT [36] 应用了师生策略来缓解ViT的数据效率问题，并且仅使用标准ImageNet数据集 [37] 就显示出强大的可比性能。 最近的一些工作，包括金字塔视觉Transformer (PVT) [5] 、Swin-Transformer [38] 、T2T-ViT [39] 、视觉Longformer (ViL) [11] 和卷积视觉Transformer (CvT) [6] ，并没有在单个低分辨率的图像块（例如， 16 × 16 图像块）上应用Transformer，而是堆叠了一个ViT金字塔来形成一个多尺度架构，并在更高的分辨率（例如， 56 × 56 = 3136 图像块用于具有 224 × 224 像素的图像）上对长序列的图像块进行建模。 这些方法大多数的自注意力计算复杂度与输入图像大小成二次方关系。

为了降低复杂度，Swin-Transformer [38] 通过仅在每个局部窗口内限制自注意力的计算，实现了线性复杂度。 HaloNet [40] 对分块图像应用局部注意力，其计算复杂度仅与块的大小成二次方关系。 Perceiver [41] 使用数据和潜在数组之间的交叉注意力来替换数据上的自注意力，从而消除二次复杂度的瓶颈。 另一项同时进行的工作，视觉Longformer (ViL) [11] ，通过将Longformer [14] 应用于视觉，实现了线性复杂度。 ViL 使用特定任务的全局符元增强局部窗口注意力，但全局符元不适用于解码任务（例如，图像合成 [23, 24] ）。 相反，我们的方法通过结合局部窗口注意力和全局动态投影注意力将二次成本降低到线性成本，这可以应用于编码和解码任务。

3.1 预备知识和符号

多头注意力是 Transformer [1] 的核心组成部分，它通过关注不同表示子空间中的整个输入序列来计算每个符元的上下文表示。 其定义为

其中 Q , K , V ∈ ℝ n × d 是查询、键和值嵌入， W O ∈ ℝ d × d 是输出的投影矩阵，第 i 个头 H i ∈ ℝ n × d k 是缩放点积注意力， d k = d / h 是每个头的嵌入维度，

其中 W i Q , W i K , W i V ∈ ℝ d × d k 是学习到的投影矩阵， A i ∈ ℝ n × n 表示每个注意力头的完整注意力矩阵。 计算和存储 A i 的复杂度为 O ( n 2 ) ，当 n 很大时，这可能是难以承受的。 为简便起见，我们下面的讨论基于一维输入序列的情况。 给定预定的顺序，将其扩展到二维图像数据是很简单的。

3.2 通过分段滑动窗口实现短期注意力

我们使用简单而有效的滑动窗口注意力来捕获细粒度的局部相关性，其中每个查询都关注固定大小邻域内的附近符元。 类似的技术也已在 [14, 16, 11] 中采用。 具体来说，为了提高效率，我们将输入序列划分为长度为 w 的不相交段。 一个段内的所有符元都关注其所属段内的所有符元，以及其所属段左右两侧 w / 2 个连续的符元（必要时进行零填充），从而导致对总共 2 w 个键值对的注意力跨度。 请参见附录中的图 5 。 对于第 i 个头部中位置 t 处的每个查询 Q t ，我们将它窗口内的 2 w 键值对表示为 K ~ t , V ~ t ∈ ℝ 2 w × d 。 使用PyTorch实现时，这种分段滑动窗口注意力比每个符元滑动窗口注意力更快，其中每个符元都关注自身及其左右 w 个符元，并且其内存消耗随序列长度线性缩放；更多细节请参见 [14] 和我们的图 3 。

通过为滑动窗口注意力的不同头部引入不同的扩张率 [14] ，可以增强滑动窗口注意力以部分捕获远程相关性。 但是，不同头部的扩张率配置需要进一步调整，并且高效实现具有不同扩张率的多头注意力并非易事。 一个更有效的替代方案是使用随机稀疏注意力 [16] 增强滑动窗口注意力，但这并不能保证像全注意力那样在每一层都捕获远程相关性。 在下一节中，我们将提出我们的远程注意力来解决这个问题。

3.3 通过动态投影实现远程注意力

先前的研究表明，自注意力矩阵可以用低秩矩阵的乘积很好地近似 [17] 。 通过用低秩矩阵的乘积 [42, 19, 18, 43, 28] 替换全注意力，每个查询能够关注所有符元。 Linformer [17] 是此类别中最具代表性的模型之一。 它学习一个固定的投影矩阵来减少键和值的长度，但是固定投影对于保持语义的位置变化不够灵活。

从这些观察结果出发，我们将第 i 个头部的动态低秩投影参数化为 P i = f ( K ) ∈ ℝ n × r ，其中 r ≪ n 是低秩大小，而 P i 取决于输入序列的所有键 K ∈ ℝ n × d 。 它将 ( n × d k ) 维键嵌入 K W i K 和值嵌入 V W i V 投影到更短的 ( r × d k ) 维键 K ¯ i 和值 V ¯ i 嵌入。 与Linformer [17] 不同，低秩投影矩阵是动态的，它取决于输入序列，旨在更灵活且更能适应例如插入、删除、释义以及其他改变序列长度的操作。 请参见表 2 中的示例。 注意，查询嵌入 Q W i Q ∈ ℝ n × d k 保持相同的长度，我们让每个查询都关注 K ¯ i 和 V ¯ i 。 通过这种方式，完整的 ( n × n ) 注意力矩阵可以分解为两个矩阵的乘积，这两个矩阵具有 r 列或行。 具体来说，我们将动态投影矩阵 P i ∈ ℝ n × r 和低秩注意力的键值嵌入 K ¯ i , V ¯ i ∈ ℝ r × d k 定义为

其中 W i P ∈ ℝ d × r 是可学习的参数， ¹ softmax对所有 n 符元的第一维上的投影权重进行归一化，这在我们的实验中稳定了训练。 注意，在所有我们考虑的实验中 K = V ，所以如果 P i 依赖于 V ，它将保持不变。公式 3 的计算复杂度为 O ( r n ) 。

为了了解完整的注意力是如何被低秩矩阵的乘积所取代的，我们将每个长程注意力的头部 H i ∈ ℝ n × d k 计算为：

因此，完整的注意力现在被两个低秩矩阵 A ¯ i ∈ ℝ n × r 和 P i ⊺ ∈ ℝ r × n 的隐式乘积所取代，计算复杂度降低到 O ( r n ) 。 注意，查询在所有符元上的有效注意力权重之和仍然为1。 我们的全局注意力允许每个查询关注同一自注意力层内的所有符元嵌入。 相反，稀疏注意力机制 [14, 16] 需要堆叠多层来构建这种相关性。

应用于自回归模型： 在自回归模型中，每个符元只能关注之前的符元，因此长程注意力对于不同的符元应该具有不同的范围。 实现我们全局注意力的一个直接方法是循环更新每个查询的 K ¯ i , V ¯ i ，但这需要由于softmax的非线性而为每个符元重新计算公式( 3 )中的投影，这导致 O ( r n 2 ) 的计算复杂度。 为了保持线性复杂度，对于自回归模型，我们首先将输入序列划分成长度为 l 的等长段，并应用我们的动态投影从每个段中提取 K ¯ i , V ¯ i 。 每个符元只能关注 K ¯ i , V ¯ i 不包含其未来符元的片段。 形式上，设 Q t 为位置 t 处的查询， K ( l − 1 ) s : l s , V ( l − 1 ) s : l s 为来自第 s 个片段的键值对，以及 s t = ⌊ t / l ⌋ 。 对于自回归模型，我们通过关注 K i , t , V i , t 来计算 Q t 的远程注意力，定义为

通过这种方式，动态低秩投影仅对每个片段并行应用一次，从而保持线性复杂度和较高的训练速度。 相比之下，由于需要递归，随机特征注意力 [30] 的训练速度较慢。

3.4 聚合远程和短期注意力

为了聚合局部和远程注意力，我们没有为不同的头采用不同的注意力机制 [12, 14, 34] ，而是让第 i 个头的每个查询都关注来自局部窗口和全局低秩投影的键和值的并集，因此它可以学习选择来自两者中的重要信息。 在我们对自回归语言模型的初步试验中，我们发现这种聚合策略比分离头部效果更好。 具体来说，对于第 i 个头，我们将全局低秩投影的键和值表示为 K ¯ i , V ¯ i ∈ ℝ r × d k ，并将局部键和值表示为 K ~ t , V ~ t ∈ ℝ 2 w × d ，它们位于查询 Q t 位置 t 的局部窗口内。 然后，位置 t 处的第 i 个注意力 H i , t 为

其中 [ ⋅ ; ⋅ ] 表示沿第一维连接矩阵。 此外，我们发现 K ~ t W i K 和 K ¯ i 的初始范数之间存在尺度不匹配，这使得在语言和视觉任务的初始化阶段，注意力偏向局部窗口。 我们引入一种归一化策略（DualLN）来对齐范数并提高聚合的有效性。

双层归一化（DualLN）： 对于具有层归一化（LN）的Transformer（参见 [44] 的图示）， K i , V i 嵌入是LN层的输出，因此它们在初始化时均值为零，方差为一。 均值为零的向量的 ℓ 2 范数与其方差成比例。 我们注意到，加权平均值将降低此类均值为零向量的方差，从而降低其范数。 结果，公式（ 3 ）中加权平均值 K ¯ i , V ¯ i 的低秩注意力嵌入向量的范数将小于来自滑动窗口注意力的常规键和值嵌入（参见图 2 左图示）。 这种尺度不匹配会导致两个副作用。 首先，局部秩分量的内积 Q t W i Q K ¯ i ⊺ 的幅度往往小于局部窗口的幅度，因此长程注意力的注意力分数系统性地较小。 其次，即使低秩和局部窗口分配相同的注意力分数，低秩注意力的键值对 K ¯ i , V ¯ i 对 H i 方向的影响也会自然减小，因为 V ¯ i 的范数较小。 这两种效应都会导致低秩分量上的梯度较小，并阻碍模型学习有效利用长程相关性。

为避免此类问题，我们在局部窗口和全局低秩注意力的键和值投影之后添加了两组层归一化，以便它们的尺度在初始化时对齐，但网络仍然可以在训练后学习重新加权范数。 具体来说，聚合注意力现在计算为

其中 LN ( ⋅ ) L , LN ( ⋅ ) G 分别表示局部和全局注意力的层归一化。 在实践中，为了保持局部注意力和动态投影之间的一致性，我们使用 LN L ( K ) , LN L ( V ) 而不是 K , V 来计算公式 3 中的 K ¯ i , V ¯ i 。 如图 2 右图所示，使用双层归一化（DualLN）训练的Transformer-LS模型的验证损失始终低于未使用双层归一化（DualLN）的模型。

4.1 长程竞技场和 IMDb 上的双向建模

为了评估长短 Transformer 作为长文本的双向编码器，我们在三个 NLP 任务上训练我们的模型，即最近提出的长程竞技场 (LRA) 基准 [20] 中的 ListOps 、 Text 和 Retrieval ，遵循 Peng 等人 [30] 和 Tay 等人 [46] 的设置。 为了进行公平比较，我们使用 PyTorch 实现以及与 [18] 中相同的数据预处理/分割、训练超参数和模型大小，除了 Retrieval ，我们意外地使用了更多预热步骤，并改进了所有模型的结果。 更多细节请参见附录 B 。 这三个任务的结果在表 1 中给出。 LRA其他两个基于图像的任务的结果，以及在JAX中实现的模型结果，见附录 C 和 C.2 。

此外，我们遵循Longformer [14] 的预训练流程，基于RoBERTa-base和RoBERTa-large [47] 对我们的模型进行预训练，并在IMDb情感分类数据集上对其进行微调。 结果见表 3 。

结果。 从表 3 可以看出，我们的基础模型优于Longformer-base，我们的大型模型比RoBERTa-large有所改进，这证明了学习对长序列建模的益处。 与LRA上模型的比较结果见表 1 。 具有每个任务最佳配置的Transformer-LS（最佳）结果见附录 B 中的表 7 。 w , r 我们还报告了在所有任务上使用固定超参数 w = 8 , r = 32 的结果。 总的来说，我们的Transformer-LS（最佳）明显优于其他高效Transformer，并且在所有三个任务上，使用 w , r = 8 , 32 的模型表现良好，同时计算量仅为其他高效Transformer的约50%到70%。 聚合局部和远程注意力的好处在 ListOps 上最为显著，这需要模型理解涉及长期和短期关系的树状结构。 在 Retrieval 中，其中测试了文档级编码能力，我们发现我们的全局注意力比窗口注意力更有效。 仅使用动态投影的测试精度比Linformer在 Text 上高约10%（即66.28对56.12），后者在序列长度上的方差最大（即标准差893）。 这表明与Linformer学习到的但固定的投影相比，动态投影在学习具有序列长度高方差的数据的表示方面具有更高的灵活性。 同样，Linformer、Nyströmformer和我们的模型在 ListOps 上优于全注意力，这表明它们可能具有更好的归纳偏差，并且高效Transformer的有效性可能超越效率本身。

动态投影的鲁棒性。 在表 2 中，我们将Linformer和提出的动态投影(DP)针对LRA的文本和检索任务上的插入和删除的鲁棒性进行了比较。 我们使用原始的、干净的训练集训练模型，只扰动它们的测试集。 对于插入操作，我们在每个测试样本的10个随机位置插入10个随机标点符号。 对于删除操作，我们删除测试样本中的所有标点符号。 这两种变换在大多数情况下都保留标签。 通过设计，动态投影对位置变化更鲁棒。

4.2 自回归语言建模

我们将我们的方法与其他高效的Transformer模型在字符级语言建模任务上进行了比较，其中每个输入符元都是一个字符。

设置。 我们在enwik8和text8数据集上训练和评估我们的模型，每个数据集包含1亿个字符，并按照 [48] 中的方法将其划分为9000万、500万和500万用于训练、验证和测试。 我们较小的12层模型和较大的30层模型都是预归一化Transformer，其宽度和深度与Longformer [20] 相同，不同之处在于我们在每一层的投影片段中添加了相对位置编码。 我们采用了Transformer-XL [9] 的缓存机制，将缓存大小设置为与输入序列长度相同。 我们遵循与Longformer类似的训练计划，并分三个阶段训练我们的模型，每个阶段的序列长度递增。 三个阶段的输入序列长度分别为2048、4096和8192。 相比之下，Longformer使用具有48GB内存的GPU分五个阶段训练其模型（我们的最大内存为32GB），其中最后一个阶段的序列长度为23040。 Longformer的窗口大小随着深度的增加而增加，其第五阶段的平均窗口大小为4352，而我们最后一个阶段平均有效的被关注符元数量为1280。 每个实验在大约8个V100 GPU上运行大约8天才能完成。 详细的超参数见附录 D 。在测试中，与Longformer相同，我们将数据集分割成长度为32K、步长为512的重叠序列，并评估在给定前面32K个字符的情况下预测接下来512个符元的BPC。

结果 表 4 显示了在text8和enwik8上的比较。 我们的方法取得了最先进的结果。 在text8上，我们使用较小模型实现了1.09的测试BPC。 在enwik8上，我们的较小模型实现了0.99的测试BPC，并且优于参数数量相当的最先进模型。 我们的较大模型获得了0.97的测试BPC，与具有2 × 个参数的压缩Transformer不相上下。 我们的结果始终优于Longformer，后者是在更长的序列上训练的，具有5个阶段和48个GPU内存。 在图 3 中，我们展示了我们的模型比全注意力模型在内存和计算方面效率更高。