摘要
基于Transformer的模型由于其自注意力机制的操作，无法处理长序列，因为其计算复杂度随序列长度呈二次方增长。为了解决这个问题，我们提出了Longformer，其注意力机制的计算复杂度随着序列长度线性增长，使得处理成千上万标记或更长的文档变得更加容易。Longformer的注意力机制可以作为标准自注意力的替代，并结合了局部窗口注意力和任务驱动的全局注意力。继之前关于长序列Transformer的研究后，我们在字符级语言建模任务上评估了Longformer，并在text8和enwik8数据集上取得了最先进的结果。与大多数之前的工作不同，我们还对Longformer进行了预训练，并在多个下游任务上进行了微调。我们的预训练Longformer在长文档任务上 consistently 超越了RoBERTa，并在WikiHop和TriviaQA上创下了新的最先进结果。最后，我们介绍了Longformer-Encoder-Decoder（LED），它是Longformer的一种变体，用于支持长文档生成式序列到序列任务，并在arXiv摘要数据集上证明了其有效性。
1 引言
Transformer（Vaswani等，2017）在广泛的自然语言处理任务中取得了最先进的成果，包括生成式语言建模（Dai等，2019；Radford等，2019）和判别式语言理解（Devlin等，2019）。这种成功部分归功于自注意力组件，它使得网络能够捕捉到整个序列的上下文信息。然而，尽管功能强大，自注意力的内存和计算需求随着序列长度的增长呈二次方增长，这使得处理长序列变得不可行（或者非常昂贵）。
为了解决这个限制，我们提出了Longformer，一种修改过的Transformer架构，采用了自注意力操作，其计算复杂度随着序列长度线性增长，使其能够处理长文档（见图1）。这对于自然语言任务，如长文档分类、问答（QA）和共指消解，具有优势，因为现有方法通常将长上下文划分或缩短为更小的序列，这些序列通常符合BERT风格预训练模型的512标记限制。这种划分可能会导致重要的跨部分信息丢失，而为了解决这个问题，现有方法通常依赖复杂的架构来处理这些交互。另一方面，我们提出的Longformer能够通过多层注意力构建整个上下文的表征，从而减少了对任务特定架构的需求。
最近的研究已经解决了Transformer在长序列上的计算低效问题（参见表1）。然而，它们主要集中在自回归语言建模（LM）上，而将长文档Transformer应用于转移学习设置下的文档级NLP任务（Dai和Le，2015；Peters等，2018；Howard和Ruder，2018；Devlin等，2019）仍然没有得到充分探索。我们填补了这一空白，并展示了Longformer的注意力机制可以作为预训练Transformer中自注意力机制的替代，从而在一系列文档NLP任务中取得了进展。
Longformer的注意力机制结合了基于窗口的局部上下文自注意力和任务驱动的全局注意力，后者编码了关于任务的归纳偏置。通过消融实验和控制试验，我们表明这两种注意力类型是至关重要的------局部注意力主要用于构建上下文表征，而全局注意力则使Longformer能够构建完整序列的表征，以便进行预测。
我们首先在自回归字符级语言建模任务上评估Longformer，使用基于窗口和新的膨胀注意力模式的组合，使得该模型能够在现代GPU上处理最多32K字符的序列。我们在text8和enwik8基准数据集上取得了最先进的结果，展示了Longformer在长文档建模中的有效性。
接着，为了评估Longformer替代现有预训练模型中完整自注意力操作的能力，我们使用掩蔽语言建模（MLM）目标对其进行预训练，继续从RoBERTa（Liu等，2019）发布的检查点开始。预训练后，我们通过微调将其应用于下游语言任务，并展示Longformer在广泛的文档级自然语言任务中始终超越RoBERTa，包括文本分类、问答和共指消解，在其中两个数据集上取得了最先进的结果。
最后，我们介绍了Longformer的一个变体，它不是采用仅编码器的Transformer架构，而是采用了类似于原始Transformer模型（Vaswani等，2017）的编码器-解码器架构，旨在进行序列到序列（seq2seq）学习（Sutskever等，2014）。我们称这种模型为Longformer-Encoder-Decoder（LED），它在编码器网络上使用Longformer的高效注意力模式，使其能够解决长文档的seq2seq任务，如摘要生成。我们在arXiv摘要数据集（Cohan等，2018）上展示了LED的有效性。
2 相关工作
长文档Transformer的相关工作总结见表1。已有两种类型的自注意力方法被探索。第一种是自左向右（ltr）方法，它将文档分块处理，从左到右进行处理。虽然这种模型在自回归语言建模中取得了成功，但它不适用于转移学习方法，尤其是对于需要双向上下文的任务。
我们的工作属于另一种常见的方法，它定义了某种稀疏注意力模式，并避免了计算完整的二次方注意力矩阵乘法。与我们注意力模式最相似的模型是Sparse Transformer（Child等，2019），它使用了一种膨胀滑动窗口的块形式，块大小为8x8，采用了BlockSparse（Gray等，2017）提供的方案。我们的实现（§3）也包括一个自定义CUDA内核，但它比BlockSparse更灵活且更易于维护，因为BlockSparse是用C++实现的，且针对特定版本的TensorFlow设计。我们还引入了额外的任务驱动的全局注意力模式，适用于常见的NLP任务（§3），并证明它们对于在转移学习环境中取得良好性能是至关重要的。
一些模型尝试了除自回归语言建模之外的其他任务，这向前迈出了一步，因为可以说，集中在语言建模作为主要评估标准，导致了只能在有限场景下适用的模型的开发。BPTransformer（Ye等，2019）在机器翻译（MT）上进行了评估，但没有探索预训练-微调的设置。Blockwise attention（Qiu等，2019）对他们的模型进行了预训练，并在问答（QA）任务上进行了评估。然而，这些评估存在局限性，因为它们没有包括语言建模任务，并且问答数据集中的文档相对较短，因此该模型在长文档任务上的有效性尚未得到充分探讨。
任务特定的长文档模型
为了绕过像BERT这样的预训练Transformer模型的512标记限制，已经开发了许多任务特定的方法。最简单的方法是直接截断文档，这种方法通常用于分类任务（Xie等，2019）。另一种方法是将文档分割成长度为512的块（可以重叠），分别处理每个块，然后使用任务特定的模型将激活结果合并（Joshi等，2019）。第三种方法常用于多跳和开放域问答任务，采用两阶段模型，其中第一阶段检索相关文档，第二阶段从中提取答案（Clark和Gardner，2017；Chen等，2017）。所有这些方法都由于截断或两阶段方法中的级联错误而导致信息丢失。与此不同，Longformer能够处理长序列而不需要截断或分块，这使我们能够采用更简单的方法，将可用的上下文连接起来并一次性处理。
一些同时期的工作也探讨了类似Longformer的思路，使用局部+全局注意力机制，并对长文档自然语言任务进行了预训练。特别地，ETC（Ainslie等，2020）使用了类似的局部+全局注意力机制，而不是完全的自注意力，以将Transformer扩展到长文档。与Longformer不同，ETC使用了相对位置嵌入（我们仅在自回归语言模型设置中使用过），引入了一个额外的训练目标（CPC损失）用于预训练，并以稍微不同的方式配置了全局注意力。它在包括阅读理解和分类在内的几个任务上展示了强大的结果。GMAT（Gupta和Berant，2020）使用了类似的思路，将输入中的少量全局位置作为全局记忆。BigBird（Zaheer等，2020）是ETC的扩展，并在包括摘要生成在内的额外任务上进行了评估。重要的是，BigBird通过理论分析表明，稀疏Transformer是序列函数的通用逼近器，并保留了完整自注意力的这些特性。
3 Longformer
原始的Transformer模型具有一个自注意力组件，其时间和内存复杂度为 O ( n 2 ) O(n^2) O(n2)，其中 n n n是输入序列的长度。为了解决这个问题，我们根据"注意力模式"稀疏化完整的自注意力矩阵，该模式指定了输入位置对之间的注意力关系。与完整自注意力不同，我们提出的注意力模式随着输入序列的增加而线性扩展，使其在处理更长序列时更为高效。本节将讨论该注意力模式的设计和实现。
3.1 注意力模式
滑动窗口
考虑到局部上下文的重要性（Kovaleva等，2019），我们的注意力模式采用了一个固定大小的窗口，围绕每个标记进行注意力计算。通过使用多个堆叠的窗口注意力层，产生了一个较大的感受野，其中顶层能够访问所有输入位置，并具有构建能够跨整个输入信息进行整合的表征的能力，类似于卷积神经网络（CNN）（Wu等，2019）。给定一个固定的窗口大小 w w w，每个标记会关注自己左右各 1 2 w \frac12w 21w个标记（见图2b）。这种模式的计算复杂度为 O ( n × w ) O(n \times w) O(n×w)，它随着输入序列长度 n n n线性扩展。在具有 ℓ \ell ℓ层的Transformer中，顶层的感受野大小为 ℓ × w \ell \times w ℓ×w（假设每层的窗口大小 w w w是固定的）。根据应用场景，可能需要在每一层使用不同的 w w w值，以平衡效率和模型表示能力（§4.1）。
膨胀滑动窗口
为了在不增加计算量的情况下进一步增加感受野，滑动窗口可以进行"膨胀"。这类似于膨胀卷积神经网络（van den Oord等，2016），其中窗口大小为膨胀因子 d d d（见图2c）。假设所有层的 d d d和 w w w是固定的，那么感受野的大小为 ℓ × d × w \ell \times d \times w ℓ×d×w，即使对于较小的 d d d值，也可以达到数万个标记。
在多头注意力中，每个注意力头计算不同的注意力得分。我们发现，在每个头使用不同的膨胀配置能够提高性能，因为没有膨胀的头可以专注于局部上下文，而有膨胀的头则专注于更长的上下文。
全局注意力
在最先进的BERT风格的自然语言任务模型中，最优的输入表示会根据任务的不同而有所变化。例如，在掩蔽语言建模（MLM）任务中，模型使用局部上下文来预测被掩蔽的词；而在分类任务中，模型将整个序列的表示聚合成一个特殊的标记（在BERT中是[CLS]标记）。对于问答（QA）任务，问题和文档会连接在一起，允许模型通过自注意力将问题与文档进行比较。
在我们的模型中，窗口化和膨胀的注意力模式不足以学习任务特定的表示。因此，我们在一些预先选定的输入位置上添加了"全局注意力"。重要的是，我们使这种注意力操作是对称的：也就是说，一个具有全局注意力的标记会关注序列中的所有标记，而序列中的所有标记也会关注它。图2d展示了一个带有全局注意力的滑动窗口注意力示例，在自定义位置的几个标记上应用了全局注意力。例如，在分类任务中，全局注意力用于[CLS]标记，而在问答任务中，全局注意力则应用于所有问题标记。由于这些标记的数量相对于 n n n较少，并且与 n n n独立，因此局部和全局注意力的组合复杂度仍为 O ( n ) O(n) O(n)。虽然指定全局注意力是任务特定的，但它为模型的注意力添加了归纳偏置，是一种比现有任务特定方法（那些使用复杂架构将信息跨小输入块组合起来）更简单的方法。
全局注意力的线性投影
回顾一下，在给定线性投影 Q , K , V Q, K, V Q,K,V的情况下，Transformer模型（Vaswani等，2017）计算注意力得分的方式如下：
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V (1) \mathrm { A t t e n t i o n } ( Q , K , V ) = \mathrm { s o f t m a x } \left( { \frac { Q K ^ { T } } { \sqrt { d _ { k } } } } \right) V \tag{1} Attention(Q,K,V)=softmax(dk QKT)V(1)
我们使用两组投影， Q s , K s , V s Q_s, K_s, V_s Qs,Ks,Vs来计算滑动窗口注意力的注意力得分， Q g , K g , V g Q_{g}, K_{g}, V_{g} Qg,Kg,Vg来计算全局注意力的注意力得分。这些额外的投影提供了灵活性，以建模不同类型的注意力，我们证明这对下游任务的最佳性能至关重要。 Q g , K g , V g Q_g, K_g, V_g Qg,Kg,Vg都被初始化为与 Q s , K s , V s Q_s, K_s, V_s Qs,Ks,Vs匹配的值。
3.2 实现
在常规的Transformer中，注意力得分的计算如式(1)所示。昂贵的操作是矩阵乘法 Q K T QK^T QKT，因为 Q Q Q和 K K K都具有 n n n（序列长度）个投影。对于Longformer，膨胀滑动窗口注意力只计算 Q K T QK^T QKT的固定数量的对角线。如图1所示，这使得内存使用线性增加，而不是完整自注意力的二次方增长。然而，实施这一点需要一种带状矩阵乘法形式，而现有的深度学习库（如PyTorch/Tensorflow）并不支持这种操作。图1比较了三种不同实现方式的性能：loop是一个内存高效的PyTorch实现，支持膨胀，但由于速度过慢，仅用于测试；chunks只支持非膨胀情况，用于预训练/微调设置；cuda是我们完全实现并高度优化的自定义CUDA内核，使用TVM（Chen等，2018）实现，并用于语言建模实验（更多细节请参见附录A）。
4 自回归语言建模
自回归或自左向右语言建模可以粗略地定义为估计给定输入序列中前一个标记/字符的概率分布。这个任务被认为是自然语言处理中的基础任务之一，最近的关于使用Transformer建模长序列的研究工作也将这一任务作为其主要评估任务（Dai等，2019；Rae等，2020；Sukhbaatar等，2019）。类似地，我们开发并评估了我们的模型在自回归语言建模上的表现。
4.1 注意力模式
对于自回归语言建模，我们使用膨胀滑动窗口注意力。根据Sukhbaatar等（2019）的做法，我们在不同的层之间使用不同的窗口大小。具体而言，我们在低层使用较小的窗口大小，在高层使用较大的窗口大小。这样可以让顶层学习整个序列的更高级别表示，同时让低层捕捉局部信息。此外，这在效率（较小的窗口大小由于较少的非零值计算量较小）和性能（较大的窗口大小有更丰富的表示能力，通常会带来性能的提升）之间提供了平衡。
我们不在低层使用膨胀滑动窗口，以最大化它们学习和利用即时局部上下文的能力。对于高层，我们只在2个头上使用少量递增的膨胀。这使得模型能够直接关注远程标记，同时不牺牲局部上下文。
4.2 实验设置
为了与先前的工作进行比较，我们专注于字符级语言建模（text8和enwik8；Mahoney，2009）。
训练 理想情况下，我们希望在现代GPU内存中能够容纳的最大窗口大小和序列长度上训练模型。然而，我们发现模型需要大量的梯度更新来首先学习局部上下文，然后再学习利用更长的上下文。为了适应这一点，我们采用了分阶段的训练程序，在多个训练阶段中逐渐增加注意力窗口大小和序列长度。具体而言，在第一阶段，我们从较短的序列长度和窗口大小开始，然后在每个后续阶段，我们将窗口大小和序列长度加倍，同时将学习率减半。这样可以加快训练速度，同时将较慢的部分（最长的序列和窗口大小）推迟到最后。我们在5个阶段中训练模型，第一阶段的序列长度为2048，最后阶段的序列长度为23040（有关每个阶段的详细配置以及所有其他超参数，请参见附录B）。
评估
我们使用长度为32,256的序列进行评估。按照Dai等（2019）的方法，我们将数据集拆分成大小为32,256、步长为512的重叠序列，并报告序列中最后512个标记的性能。
4.2.1 结果
表2和表3总结了在text8和enwik8数据集上的评估结果。我们在text8和enwik8上分别使用小型模型取得了新的最先进的性能，text8的BPC为1.10，enwik8的BPC为1.00，展示了我们模型的有效性。
对于大型模型，考虑到这些实验的高成本，并参考最近的研究（Kitaev等，2020；Rae等，2020），我们仅在enwik8上进行评估。表3显示，Longformer超越了可比的TransformerXL模型，性能与可比的Sparse Transformer（Child等，2019）相当，并且与那些具有超过两倍参数量的最近模型的表现持平或略逊一筹。值得注意的是，Adaptive Span（Sukhbaatar等，2019）和Compressive Transformer（Rae等，2020）并不适合预训练-微调范式，如§2中所讨论的。
4.2.2 消融研究
为了展示我们注意力模式设计选择的重要性，我们尝试了不同的变体，并报告了它们的控制实验结果。为了使消融研究更易于管理，我们在text8数据集上使用小型模型，采用第一阶段配置训练每个配置150K步，然后报告开发集上的BPC性能。
表4的顶部展示了每层窗口大小配置的不同方式对性能的影响。我们观察到，从底层到顶层逐渐增加窗口大小能够带来最佳性能，而将其反向配置则导致性能下降，使用固定的窗口大小（其他配置的窗口大小的平均值）则会得到介于两者之间的性能。
表4的底部展示了添加膨胀的影响。相比没有膨胀，给两个头部添加一些膨胀会带来一定的性能提升。
温馨提示：
阅读全文请访问"AI深语解构 " Longformer：长文档Transformer