Longformer

论文名称

Longformer: The Long-Document Transformer

Abstract

提示

基于transformer的模型无法处理长序列，这是因为它们的自注意力操作与序列长度成二次方比例。

为了解决这一局限性，我们引入了Longformer，它的注意力机制与序列长度成线性比例，使得它可以很容易地处理成千上万的标记或更长的文件。

Longformer的注意力机制是对标准的自注意力的直接替代，并结合了局部窗口注意力和任务驱动的全局注意力。

按照之前关于长序列transformer的工作，我们对Longformer在字符级语言建模上进行了评估，并在text8和enwik8上取得了最先进的结果。

与大多数先前的工作不同，我们还对Longformer进行了预训练，并在各种下游任务中对其进行了微调。我们预训练的Longformer在长文档任务上的表现一直优于RoBERTa，并在WikiHop和TriviaQA上创造了最先进的结果。最后，我们介绍了Longformer-Encoder-Decoder（LED），这是一个Longformer的变体，用于支持长文档生成序列到序列的任务，并在arXiv总结数据集上展示了其有效性。

Introduction

Transformer（Vaswani等人，2017）在广泛的自然语言任务中取得了最先进的成果，包括生成性语言建模（Dai等人，2019；Radford等人，2019）和判别性语言理解（Devlin等人，2019）。这种成功部分归功于自注意力机制，它使网络能够从整个序列中捕获上下文信息。虽然功能强大，但自注意力机制消耗的资源随着序列长度的增长而平方倍增长，使得处理长序列不可行（或非常昂贵）。

为了解决这个问题，我们提出了Longformer，一个改进的Transformer架构，它的自注意力操作与序列长度成线性关系，使其能够通用于处理长文档（图1）。这是自然语言任务的一个优势，如长文档分类、问题回答（QA）和共指消解（coreference resolution），在这些任务中，现有的方法将长上下文分割或将其缩短为较小的序列，使得这些序列在BERT等预训练模型的经典512token限制之内。这种分割有可能导致重要的跨分区信息的损失，为了缓解这个问题，现有的方法往往依靠复杂的架构来解决这种相互作用。另一方面，我们提出的Longformer能够使用多层注意力建立整个上下文的上下文表征，减少对特定任务架构的需求。

最近的工作已经解决了Transformer在长序列上的计算效率低下问题（见表1）。然而，他们主要关注的是自回归语言建模（LM），而文档级NLP任务在迁移学习中的应用（Dai和Le，2015；Peters等人，2018；Howard和Ruder，2018；Devlin等人，2019）基本上没有得到探索。我们解决了这个问题，并表明Longformer的注意力机制可以作为预训练的Transformers中自注意力机制的替代品，并在一套文档NLP任务中的取得较好效果。

Longformer的注意机制是一个窗口式的局部上下文自注意力和一个终端任务驱动的全局注意力的组合，全局注意力对任务的归纳偏置进行编码。通过消融和控制试验，我们表明这两种注意力类型都是必不可少的--局部注意力主要用于建立上下文表征，而全局注意力则使Longformer能够建立完整的序列表征以进行预测。

我们首先评估了Longformer在自回归字符级语言建模中使用窗口化和新的扩张注意力模式的组合，使该模型能够在现代GPU上处理多达32K字符的序列。我们在text8和enwik8基准数据集上取得了最先进的结果，证明了Longformer在长文档建模上的有效性。

然后，为了评估Longformer取代现有预训练模型的完全的注意力机制的能力，我们用掩蔽语言建模（MLM）目标对其进行预训练，延续RoBERTa（Liu等人，2019）发布的checkpoint。在预训练后，我们通过微调将其应用于下游语言任务，并证明Longformer在广泛的文档级自然语言任务中始终优于RoBERTa，包括文本分类、QA和共指消解，在其中两个数据集上取得了最先进的结果。

我们最后介绍了Longformer的一个变体，它不是只采用编码器的Transformer架构，而是采用类似于原始Transformer模型的编码器-解码器架构（Vaswani等人，2017），它旨在用于序列到序列（seq2seq）学习（Sutskever等人，2014）。我们称这个模型为Longformer-Encoder-Decoder（LED），它在编码器网络上使用Longformer的高效注意力模式，使其能够解决长文档seq2seq任务，如总结。我们在arXiv总结数据集上证明了LED的有效性（Cohan等人，2018）。

Related Work

Long-Document Transformers

Tab. 1总结了最近关于长文档的工作。两种类型的自注意力的方法已经被探索过了。第一种是从左到右（ltr）的方法，它以从左到右的块状方式处理文件。虽然这种模型在自回归语言建模中取得了成功，但它们不适合用于迁移学习的方法，因为这种方法的任务受益于双向的上下文。

我们的工作属于另一种一般的方法，即定义某种形式的稀疏注意力模式，并避免计算完整的注意力矩阵乘法。与我们的注意力模式最相似的模型是Sparse Transformer（Child等人，2019），它使用BlockSparse（Gray等人，2017）提供的大小为8x8的块的扩张滑动窗口的形式。我们的实现也包括一个定制的CUDA内核，但它比BlockSparse更灵活和可维护，BlockSparse是用C++实现的，并为TensorFlow的一个特定版本设计。我们还引入了额外的任务驱动的全局注意力模式，适用于常见的NLP任务，并表明它们对于迁移学习的性能至关重要。

也有少数模型尝试了自回归语言建模以外的任务，这是一个进步，因为可以说专注于语言建模作为主要的评估，导致了适用性有限的模型的发展。BPTransformer（Ye等人，2019）对机器翻译（MT）进行了评估，但没有探索pre-train finetune设置。Blockwise attention（Qiu等人，2019）预训练了他们的模型，并对问题回答（QA）进行了评估。然而，评估是有限的，因为它不包括语言建模，而且QA数据集是相对较短的文档（SQuAD的语境通常符合512的限制，而MRQA是通过放弃长文件的例子来构建的。），因此这个模型在长文档任务上的有效性仍然没有被探索。

Task-specific Models for Long Documents

许多特定任务的方法被开发出来解决像BERT这样的预训练transformer模型的512限制问题。最简单的方法只是截断文件，通常用于分类（Xie等人，2019）。另一种方法是将文件分成长度为512的小块（可能是重叠的），分别处理每个小块，然后用特定的任务模型结合激活（Joshi等人，2019）。第三种流行于多跳问答和开放领域QA任务的方法使用两阶段模型，第一阶段检索相关文件，然后传递给第二阶段进行答案提取（Clark和Gardner，2017；Chen等人，2017）。所有这些方法都因两阶段方法的截断或级联错误而遭受信息损失。相比之下，Longformer可以在不截断或分块的情况下处理长序列，使我们可以采用更简单的方法，将可用的上下文串联起来，一次性处理。

一些同时代的工作探索了与Longformer类似的想法，在transformer中使用局部+全局注意力，并为长文档自然语言任务进行预训练。特别是ETC（Ainslie等人，2020）使用类似的局部+全局注意力，而不是完全的自注意力机制，以将transformer扩展到长文档。与Longformer不同的是，ETC使用相对位置嵌入（我们只用于自回归LM设置），为预训练引入了一个额外的训练目标（CPC损失），并以稍微不同的方式配置了全局注意力。它在包括阅读理解和分类在内的几个任务上显示了强大的结果。GMAT（Gupta和Berant，2020）使用了一个类似的想法，即输入中的少数全局位置作为全局记忆。BigBird（Zaheer等人，2020）是在ETC基础上的扩展，对包括总结在内的其他任务进行评估。重要的是，通过理论分析，BigBird表明稀疏transformers是序列函数的通用近似器，并保留了完整的自注意力的这些特性。

Longformer

原始的Transformer模型有一个自注意力模块，其时间和内存复杂度为$O(n^{2})$，其中$n$为输入序列长度。为了应对这一挑战，我们根据 "注意力模式 "对完整的自注意力矩阵进行了稀疏化，该模式指定了相互关注的输入位置对。与完整的自注意力不同，我们提出的注意力模式与输入序列成线性关系，这使得它对较长的序列很有效。本节讨论了这种注意力模式的设计和实现。

Attention Pattern

Sliding Window

鉴于局部上下文的重要性（Kovaleva等人，2019年），我们的注意力模式采用了围绕每个token的固定大小的窗口注意力。使用这种窗口式注意力的多个堆叠层会产生一个大的感受野，其中顶层可以接触到所有的输入位置，并且有能力建立包含整个输入信息的表征，类似于CNN（Wu等人，2019）。给定一个固定的窗口大小$w$，每个token都关注每侧的$\frac{1}{2}w$个令牌（图2b）。这种模式的计算复杂度为$O(n{\times}w)$,与输入序列长度$n$呈线性关系。在具有$l$层的transformer中，顶层的感受野大小为$l{\times}w$（假设所有层的$w$是固定的）。根据不同的应用，对每一层使用不同的$w$值来平衡效率和模型表示能力可能是有帮助的。

Dilated Sliding Window

为了在不增加计算量的情况下进一步增加感受野，滑动窗口可以被"扩张"。这类似于扩张的CNN（van den Oord等人，2016），其中窗口有大小为$d$的间隙（图2c）。假设所有层都有固定的$d$和$w$，感受野是$l{\times}d{\times}w$，即使$d$的值很小，也可以达到几万个标记。

在多头注意力中，每个注意力头计算不同的注意力分数。我们发现，每个头有不同的扩张配置的设置，通过允许一些没有扩张的头关注局部上下文，而其他有扩张的头关注更长的上下文，来提高性能。

Global Attention

在最先进的用于自然语言任务的BERT式模型中，最佳输入表征与语言建模不同，并因任务而异。对于掩码语言建模（MLM），模型使用局部上下文来预测masked word，而对于分类，模型将整个序列的表征汇总为一个特殊的标记（[CLS]在BERT的情况下）。对于QA，问题和文档被连接起来，允许模型通过自注意力来比较问题和文档。

在我们的情况中，窗口式和扩张式注意力不够灵活，无法学习特定任务的表征。因此，我们在少数预选的输入位置上增加了 "全局注意"。重要的是，我们使这种注意力操作是对称的：也就是说，一个具有全局注意的token会关注整个序列中的所有标记，而序列中的所有标记也会关注它。图2d显示了一个滑动窗口注意的例子，全局注意力在自定义位置的几个标记上。例如，对于分类，全局注意力用于[CLS] token，而在QA中，全局注意力提供给所有问题标记。由于这种标记的数量相对于n来说是很小的，并且与n无关，所以结合局部和全局注意的复杂性仍然是$O(n)$。虽然指定全局注意力是特定的任务，但这是一个简单的方法，可以将归纳的偏置添加到模型的注意力中，而且它比现有的特定任务的方法简单得多，这些方法使用复杂的架构来结合较小输入块的信息。

Linear Projections for Global Attention

回顾一下，给定线性投影$Q$、$K$、$V$，Transformer模型（Vaswani等人，2017）计算注意力分数的方法如下。

$$\begin{equation} \begin{aligned} \text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V \end{aligned} \end{equation}$$

我们使用两组投影，$Q_s$, $K_s$, $V_s$来计算滑动窗口注意力的分数，$Q_g$, $K_g$, $V_g$来计算全局注意力的分数。这些投影提供了对不同类型注意力建模的灵活性，我们表明这对下游任务的最佳表现至关重要。$Q_g$、$K_g$、$V_g$的初始化值都与$Q_s$、$K_s$、$V_s$相匹配。

Implementation

在常规的transformers中，注意力分数的计算与公式1相同。昂贵的操作是矩阵乘法$QK^T$，因为$Q$和$K$都有$n$个（序列长度）的投影。对于Longformer，稀释的滑动窗口注意力只计算$QK^T$的固定数目的对角线。如图1所示，这导致了内存使用量的线性增加，而完全的自注意力则是二次增加。然而，实现它需要一种带状矩阵乘法，而现有的深度学习库如PyTorch/Tensorflow并不支持。图1比较了三种不同实现方式的性能：loop是一种内存效率高的PyTorch实现，支持扩张，但速度异常缓慢，仅用于测试；chunks仅支持非扩张情况，用于预训练/调整设置；cuda是我们使用TVM（Chen等人，2018）实现的全功能高度优化的定制CUDA内核，用于语言建模实验（详情见附录A）。

Autoregressive Language Modeling

自回归或从左到右的语言建模被宽泛地定义为，给定一个标记/字符之前的标记/字符，估计这个标记/字符在输入序列中的概率分布。

这项任务被认为是自然语言的基本任务之一，最近使用transformers对长序列进行建模的前期工作都是依靠这项任务作为他们的主要评价（Dai等人，2019；Rae等人，2020；Sukhbaatar等人，2019）。同样地，我们在自回归语言建模上开发和评估我们的模型。

Attention Pattern

对于自回归语言建模，我们使用我们的扩张滑动窗口注意力。按照Sukhbaatar等人（2019）的做法，我们在各层使用不同的窗口大小。特别是，我们在低层使用小的窗口尺寸，随着我们向高层移动而增加窗口尺寸。这使得顶层可以学习整个序列的更高层次的表征，同时让低层捕捉局部信息。此外，它还在效率（由于非零值较少，较小的窗口尺寸的计算成本较低）和性能（较大的窗口尺寸具有更丰富的表示能力，通常会导致性能的提高）之间提供了平衡。

我们不对低层使用扩张的滑动窗口，以最大限度地提高其学习和利用直接的局部上下文的能力。对于较高的层，我们只在2个头上使用少量的递增扩张。这使模型有能力在不牺牲本地上下文的情况下直接注意到远处的token。

Experiment Setup

为了与之前的工作进行比较，我们把重点放在字符级的LM上（text8和enwik8；Mahoney，2009）。

Training

理想情况下，我们希望在现代GPU内存中能容纳的最大窗口大小和序列长度上训练我们的模型。然而，我们发现该模型需要大量的梯度更新来首先学习局部上下文，然后再学习利用更长的上下文。为了适应这一点，我们采用了一个分阶段的训练程序，在多个训练阶段增加注意力窗口的大小和序列的长度。特别是，在第一阶段，我们以较短的序列长度和窗口大小开始，然后在随后的每个阶段，我们将窗口大小和序列长度增加一倍，并将学习率减半。这使得训练速度很快，而将慢的部分（最长的序列和窗口大小）保留到最后。我们总共训练了5个阶段的模型，开始的序列长度为2,048，最后一个阶段的序列长度为23,040（每个阶段的详细配置和所有其他超参数见附录B）。

Evaluation

我们用长度为32,256的序列进行评估。按照Dai等人（2019）的做法，我们将数据集分成大小为32,256的重叠序列，步长为512，并报告序列上最后512个token的性能。

Results

表2和表3总结了对text8和enwik8数据集的评估结果。2和3总结了text8和enwik8数据集的评估结果。我们在text8和enwik8上使用小模型实现了新的最先进的技术，在text8和enwik8上的BPC分别为1.10和1.00，证明了我们模型的有效性。

对于大型模型，考虑到这些实验的成本很高，并遵循最近的工作（Kitaev等人，2020；Rae等人，2020），我们只在enwik8上进行评估。Tab. 3显示，Longformer的性能优于可比较的TransformerXL模型，与可比较的Sparse Transformer（Child等人，2019）的性能相匹配，并与最近那些参数数量超过两倍的模型相匹配或略逊于它们。值得注意的是，Adaptive Span（Sukhbaatar等人，2019）和Compressive Transformer（Rae等人，2020）不是很适合讨论的pretraining-finetuning范式。

Ablation Study

为了显示我们注意力模式的设计选择的重要性，我们尝试了不同的变体，并报告了它们的控制实验结果。为了使消融研究更易于管理，我们用第一阶段的配置在text8的小模型上训练每个配置150K步（有一点需要注意的是，最终性能的排序将与步骤150K时的性能不一致。然而，这个近似值节省了将每个实验运行到完成的巨大成本。），然后报告BPC在dev集上的表现。

Tab. 4展示了每层窗口大小的不同配置方式的影响。我们观察到，从底层到顶层增加窗口尺寸会导致最好的性能，以相反的方式安排它们会导致更差的性能，而使用固定的窗口尺寸（其他配置的窗口尺寸的平均值）会导致它的性能介于两者之间。表4的底部显示了加入扩张的影响。4显示了添加扩张的影响。与完全没有扩张的情况相比，给两个头添加一些扩张导致了一些改进。

Pretraining and Finetuning

目前用于许多NLP任务的最先进的系统在任务监督下对预训练的模型进行微调（例如BERT）。我们的主要动机之一是开发这样一个适用于长文档任务的模型。为此，我们在一个文档语料库上对Longformer进行了预训练，并针对六个任务对其进行了微调，包括分类、QA和共指消解。由此产生的模型可以处理长达4,096个token的序列（比BERT长8倍）。

我们用掩蔽语言建模（MLM）对Longformer进行预训练，目标是恢复序列中随机掩蔽的标记。由于MLM预训练很昂贵，我们从RoBERTa（Liu等人，2019）发布的checkpoint继续预训练，同时只做了支持Longformer注意力机制所需的最小变化。请注意，我们的注意力模式可以插入任何预训练的转化器模型中，而不需要改变模型架构。

Attention Pattern

我们使用滑动窗口注意力，窗口大小为512，因此使用的计算量与RoBERTa相同。

Position Embeddings

RoBERTa使用学习的绝对位置嵌入，最大位置为512。为了支持较长的文件，我们增加了额外的位置嵌入，以支持高达4,096个位置。为了利用RoBERTa的预训练权重，我们没有随机初始化新的位置嵌入，而是通过多次复制RoBERTa的512个位置嵌入来初始化它们，因为对BERT的注意力头的分析表明，有一种强烈的学习偏置于关注局部上下文，包括前一个或下一个标记（Clark等人，2019）。使用复制初始化，除了在分区边界外，在任何地方都保留了这种局部结构。尽管它很简单，但我们发现这是一个非常有效的方法（见表5），允许Longformer预训练通过少量的梯度更新迅速收敛。

Continued MLM Pretraining

我们使用fairseq（Ott等人，2019）对我们编制的长篇文献语料库（语料库详情见附录C）进行Longformer预训练。我们训练两种规模的模型，一个基础模型和一个大型模型。两种模型都训练了65K的梯度更新，序列长度为4096，批次大小为64（$2^{18}$个标记），最大学习率为3e-5，线性预热为500步，然后是3次方的多项式衰减。其余的超参数与RoBERTa相同。

表5显示了我们训练语料库的开发集上的BPC。第一行显示使用RoBERTa-base的BPC为1.846，这与RoBERTa论文中关于其语料库的1.880 BPC相当。这表明我们的训练语料库的分布与用于训练RoBERTa的分布相近。下面两行显示了Longformer在使用随机初始化位置嵌入和使用复制的位置嵌入进行预训练之前的性能。显著的差异表明了复制初始化的重要性，而RoBERTa BPC和初始化BPC之间相对较小的差异表明，我们的滑动窗口注意与RoBERTa权重的工作很好。下面两行显示了继续预训练的影响。训练2K步后，BPC从1.957提升到1.753，65K步后进一步下降到1.705，表明模型正在学习更好地利用滑动窗口的注意力和更长的背景。在RoBERTa-large和Longformer-large中也观察到类似的模式。

Frozen RoBERTa Weights

我们还预训练了Longformer，同时冻结了所有RoBERTa的权重，只训练新的位置嵌入。这种配置的动机是为了完美地保持RoBERTa在短文档上的性能。这种配置的BPC为1.850（低于初始化时的1.957），但高于所有权重都可训练的1.705。

Tasks

我们将Longformer应用于多个长文档任务，包括QA、共指消解和分类。Tab. 6显示了评估数据集的上下文明显长于512个字片。我们的主要目标是评估我们的注意力机制是否可以作为BERT类模型中的标准自注意力机制的替代品，并针对一个强大的基线进行对照试验。我们也有兴趣评估我们是否可以用更简单的模型取代BERT的有限上下文所需要的复杂的特定任务模型，这些模型只是将所有可用的上下文连接成一个单一的序列。

我们的基线是一个基于RoBERTa的模型，它将上下文分成尽可能长的片段，将每个片段单独通过RoBERTa，并将激活的部分连接起来进行进一步处理。对于QA任务，我们还将问题串联到每个片段中，这样RoBERTa就可以根据问题来决定它的上下文表述。Longformer变体将RoBERTa的自我注意机制替换为我们在预训练期间使用的窗口注意力，以及一个任务激励的全局注意力。全局注意力使用额外的线性投射。

Document Classification

我们在IMDB（Maas等人，2011）和Hyperpartisan新闻检测（Kiesel等人，2019）数据集上进行评估。 IMDB是一个由电影评论组成的标准情感分类数据集。虽然这个数据集中的大多数文件都很短，但其中约有13.6%的文件大于512个字（表6）。Hyperpartisan中的文档相对较长，而且它很小，只有645个文档，这使它成为Longformer适应有限数据能力的一个很好的测试。我们在[CLS]标记上使用全局关注。