Longformer: The Long-Document Transformer

Longformer：面向长文档的Transformer

作者/机构: Iz Beltagy, Matthew E. Peters, Arman Cohan (Allen Institute for Artificial Intelligence, Seattle, WA, USA)

A1 主要贡献

本文旨在解决基于Transformer的模型因其自注意力操作（其计算量与序列长度成二次方关系）而无法处理长序列的局限性。

核心问题: 标准Transformer模型的自注意力机制的内存和计算需求随序列长度二次方增长（$O(n^2)$），这使得处理长文档（如数千个token）变得不可行或成本极高。现有的方法通常通过分割长文档来规避这个问题，但这可能导致跨分片重要信息的丢失，并且需要复杂的任务特定架构来处理分片间的交互。

研究目标: 提出一种改进的Transformer架构，其自注意力操作的计算复杂度与序列长度呈线性关系（$O(n)$），从而能够高效地处理长文档。目标是创造一个可以作为标准自注意力“即插即用”替代品的模块，并将其应用于预训练和微调的范式中，以提升各种长文档自然语言处理任务的性能。

核心创新点:
1. 提出Longformer模型: 引入了一种新的注意力机制，它结合了局部窗口化注意力（local windowed attention）和任务驱动的全局注意力（task-motivated global attention）。
* 局部窗口化注意力: 每个token只关注其周围固定大小（窗口大小为$w$）的上下文，将计算复杂度从$O(n^2)$降低到$O(n \times w)$。通过堆叠多层，模型依然能构建一个较大的感受野。
* 任务驱动的全局注意力: 在少数预先选择的token上使用全局注意力，即这些token可以关注序列中的所有其他token，同时序列中的所有token也关注它们。这种设计为模型注入了关于任务的归纳偏置（inductive bias），例如，在分类任务中为[CLS] token设置全局注意力，在问答任务中为问题的所有token设置全局注意力。
2. 多种注意力模式与实现:
* 提出了滑动窗口注意力（Sliding Window Attention）和扩张滑动窗口注意力（Dilated Sliding Window Attention），后者通过在窗口中设置空洞来在不增加计算量的情况下扩大感受野。
* 为支持不同类型的注意力，模型为局部注意力和全局注意力使用了两套独立的线性投影（$Q, K, V$），这被证明对下游任务的最佳性能至关重要。
* 开发了多种实现方式，包括一个使用TVM框架实现的自定义CUDA核，以高效地执行这种稀疏注意力计算。下图展示了不同实现的性能和内存占用对比。


图 1: 全局自注意力与Longformer不同自注意力实现的运行时间和内存对比；Longformer-loop是非向量化实现，Longformer-chunk是向量化实现，Longformer-cuda是自定义cuda核实现。Longformer的内存使用随序列长度线性扩展，不像全局自注意力机制在当前GPU上处理长序列时会耗尽内存。不同实现的运行速度不同，向量化的Longformer-chunk最快。更多细节见3.2节。
3. 在多种任务上验证有效性:
* 在自回归语言建模任务（text8和enwik8）上取得了当时的SOTA结果。
* 在预训练-微调范式中，证明了Longformer可以作为RoBERTa等预训练模型的直接替代品，并在多个长文档下游任务（包括问答、共指消解和文档分类）上持续优于基线模型，在WikiHop和TriviaQA上取得了SOTA结果。
* 提出了Longformer-Encoder-Decoder (LED)模型，将Longformer的注意力机制应用于Encoder-Decoder架构，用于处理长文档的序列到序列任务（如摘要），并在arXiv摘要数据集上取得了SOTA结果。

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

面向长文档的Transformer

现有方法的分类。处理长文档的自注意力方法主要有两类。第一类是从左到右（ltr）的方法，它按块处理文档。这类模型在自回归语言建模中很成功，但不适用于需要双向上下文的迁移学习任务。本文的工作属于第二类，即定义某种形式的稀疏注意力模式，以避免计算完整的二次方注意力矩阵。

与Sparse Transformer的比较。与本文注意力模式最相似的模型是Sparse Transformer【7，Generating long sequences with sparse transformers，2019，arXiv】，它使用了BlockSparse【18，Gpu kernels for block-sparse weights，2017】提供的一种8x8块大小的扩张滑动窗口。本文的实现（§3）也包含一个自定义CUDA核，但它比用C++实现的、为特定TensorFlow版本设计的BlockSparse更灵活、更易于维护。此外，本文还引入了适用于常见NLP任务的、由任务驱动的全局注意力模式（§3），并证明了它们对于在迁移学习设置中取得良好性能至关重要。

对其他任务的探索有限。之前的一些模型尝试了自回归语言建模以外的任务，这是一个进步，因为仅关注语言建模作为主要评估导致了模型适用性的局限。BP-Transformer【60，BP-Transformer: Modelling long-range context via binary partitioning，2019，arXiv】在机器翻译（MT）上进行了评估，但没有探索预训练-微调的设置。Blockwise attention【38，Blockwise self-attention for long document understanding，2019，arXiv】对模型进行了预训练并在问答（QA）上评估，但评估范围有限，未包括语言建模，且QA数据集的文档相对较短，因此该模型在长文档任务上的有效性仍有待探索。


表 1: 以往工作中关于适配Transformer处理长文档的总结。ltr：从左到右。

针对长文档的任务特定模型

现有方法的局限性。为了绕过BERT等预训练Transformer模型512个token的限制，许多任务特定的方法被开发出来。最简单的方法是截断文档，常用于分类任务【57，Unsupervised data augmentation for consistency training，2019，arXiv】。另一种方法是将文档分块，每块长度为512（可以重叠），然后分别处理每个块，再用一个任务特定的模型组合这些激活值【23，BERT for coreference resolution: Baselines and analysis，2019，EMNLP-IJCNLP】。第三种方法在多跳和开放域问答任务中很流行，它使用一个两阶段模型：第一阶段检索相关文档，然后传递给第二阶段进行答案抽取【8，Simple and effective multi-paragraph reading comprehension，2017，ACL；3，Reading wikipedia to answer open-domain questions，2017，ACL】。这些方法都因截断导致信息丢失或两阶段方法带来的级联错误而受损。

Longformer的优势。相比之下，Longformer可以处理长序列而无需截断或分块，这使得我们可以采用一种更简单的方法，即将所有可用上下文连接起来，在一次传递中进行处理。

同期相关工作

与同期工作的比较。一些同期的工作探索了与Longformer类似的思想，即在Transformer中使用局部+全局注意力，并为长文档自然语言任务进行预训练。特别是，ETC【1，ETC: Encoding long and structured inputs in transformers，2020，EMNLP】使用了类似的局部+全局注意力来替代全自注意力，以将Transformer扩展到长文档。与Longformer不同的是，ETC使用了相对位置嵌入（本文只在自回归语言模型设置中使用），为预训练引入了一个额外的训练目标（CPC损失），并以略有不同的方式配置全局注意力，它在包括阅读理解和分类在内的多个任务上展示了强大的结果。GMAT【20，Gmat: Global memory augmentation for transformers，2020，ArXiv】使用了类似的想法，即输入中的少数全局位置充当全局记忆。BigBird【61，Big bird: Transformers for longer sequences，2020，ArXiv】是ETC的扩展，在包括摘要在内的更多任务上进行了评估。重要的是，通过理论分析，BigBird证明了稀疏Transformer是序列函数的通用近似器，并保留了全自注意力的这些属性。

A2 方法细节

3. Longformer

核心挑战与解决方案。原始Transformer模型的自注意力组件具有$O(n^2)$的时间和内存复杂度，其中$n$是输入序列的长度。为了解决这一挑战，本文根据一个指定输入位置对之间相互关注的“注意力模式”来稀疏化完整的自注意力矩阵。与完整的自注意力不同，本文提出的注意力模式与输入序列长度呈线性扩展关系，使其能高效处理更长的序列。本节将讨论这种注意力模式的设计与实现。

3.1 注意力模式

滑动窗口注意力。鉴于局部上下文的重要性【27，Revealing the dark secrets of bert，2019，EMNLP/IJCNLP】，本文的注意力模式采用了围绕每个token的固定大小窗口注意力。使用多个堆叠的这种窗口化注意力层可以形成一个大的感受野，顶层可以访问所有输入位置，并有能力构建包含整个输入信息的表示，这与CNNs类似【56，Pay less attention with lightweight and dynamic convolutions，2019，arXiv】。给定一个固定的窗口大小$w$，每个token会关注其两侧各$\frac{1}{2}w$个token（如图2b所示）。这种模式的计算复杂度为$O(n \times w)$，随输入序列长度$n$线性增长。在一个有l层的transformer中，顶层的感受野大小为l $\times w$（假设所有层的$w$固定）。根据应用场景，为不同层使用不同的$w$值可能有助于平衡效率和模型表示能力（§4.1）。


图 2: 全局自注意力模式与我们Longformer中注意力模式配置的比较。

扩张滑动窗口。为了在不增加计算量的情况下进一步增大感受野，滑动窗口可以被“扩张”。这类似于扩张CNNs【34，Wavenet: A generative model for raw audio，2016，SSW】，其中窗口有大小为$d$的间隙（如图2c所示）。假设所有层都有固定的$d$和$w$，感受野大小为l $\times d \times w$，即使$d$的值很小，也能达到数万个token。在多头注意力中，每个注意力头计算不同的注意力分数。本文发现，为每个头设置不同的扩张配置可以提升性能，这允许一些没有扩张的头专注于局部上下文，而其他有扩张的头则专注于更长的上下文。

全局注意力。在当前最先进的BERT风格模型中，用于自然语言任务的最佳输入表示与语言建模不同，并因任务而异。对于掩码语言建模（MLM），模型使用局部上下文来预测被掩盖的词；而对于分类，模型将整个序列的表示聚合到一个特殊的token（如BERT中的[CLS]）中。对于问答（QA），问题和文档被连接起来，让模型通过自注意力来比较问题和文档。本文的窗口化和扩张注意力不够灵活，无法学习任务特定的表示。因此，本文在一些预先选择的输入位置上增加了“全局注意力”。重要的是，这种注意力操作是对称的：即一个具有全局注意力的token会关注序列中的所有token，并且序列中的所有token也会关注它。图2d展示了一个在自定义位置的少数token上带有全局注意力的滑动窗口注意力示例。例如，对于分类任务，全局注意力用于[CLS] token；而在QA任务中，全局注意力施加在所有问题token上。由于这类token的数量相对于$n$很小且独立于$n$，局部和全局注意力结合后的复杂度仍然是$O(n)$。虽然指定全局注意力是任务特定的，但这是向模型的注意力中添加归纳偏置的一种简单方法，并且比现有那些使用复杂架构来组合较小输入块信息的任务特定方法简单得多。

全局注意力的线性投影。回顾一下，给定线性投影$Q, K, V$，Transformer模型【52，Attention is all you need，2017，NIPS】计算注意力分数的方式如下：


本文使用两套投影，$Q_s, K_s, V_s$用于计算滑动窗口注意力的分数，而$Q_g, K_g, V_g$用于计算全局注意力的分数。额外的投影为模型提供了建模不同类型注意力的灵活性，本文证明这对于下游任务的最佳性能至关重要。$Q_g, K_g, V_g$都用与$Q_s, K_s, V_s$相匹配的值进行初始化。

3.2 实现

实现细节与挑战。在常规的transformer中，注意力分数如公式1计算。其中昂贵的操作是矩阵乘法$QK^T$，因为$Q$和$K$都有$n$（序列长度）个投影。对于Longformer，扩张滑动窗口注意力只计算$QK^T$的固定数量的对角线。如图1所示，这导致内存使用量呈线性增长，而全自注意力则是二次增长。然而，实现它需要一种带状矩阵乘法，而现有的深度学习库如PyTorch/Tensorflow并不支持。图1比较了三种不同实现方式的性能：loop是一种内存高效的PyTorch实现，支持扩张但速度慢到无法使用，仅用于测试；chunks只支持非扩张情况，用于预训练/微调设置；cuda是本文使用TVM【5，TVM: An automated end-to-end optimizing compiler for deep learning，2018，OSDI】实现的完全功能化且高度优化的自定义CUDA核，用于语言建模实验（更多细节见附录A）。

4. 自回归语言建模

任务定义。自回归或从左到右的语言建模被广义地定义为估计一个输入序列中给定其先前token/字符的现有token/字符的概率分布。这项任务被认为是自然语言处理的基础任务之一，并且最近使用transformer对长序列进行建模的前期工作都依赖于这项任务作为其主要评估【12，Transformer-XL: Attentive language models beyond a fixed-length context，2019，ACL；41，Compressive transformers for long-range sequence modelling，2020，ICLR；46，Adaptive attention span in transformers，2019，ACL】。同样，本文也在自回归语言建模上开发和评估我们的模型。

4.1 注意力模式

自回归语言建模的注意力配置。对于自回归语言建模，本文使用扩张滑动窗口注意力。遵循Sukhbaatar等人【46，Adaptive attention span in transformers，2019，ACL】的工作，本文在不同层使用不同的窗口大小。具体来说，在底层使用较小的窗口，并随着层数增高而增加窗口大小。这使得顶层能够学习整个序列的更高级别表示，同时让底层捕捉局部信息。此外，这也平衡了效率（较小的窗口大小因非零值较少而计算成本较低）和性能（较大的窗口大小具有更丰富的表示能力，并通常带来性能提升）。本文在底层不使用扩张滑动窗口，以最大化它们学习和利用直接局部上下文的能力。对于高层，仅在2个头上使用少量递增的扩张。这使模型能够在不牺牲局部上下文的情况下直接关注远距离的token。

4.2 实验设置

数据集与训练策略。为了与先前工作进行比较，本文专注于字符级语言建模（text8和enwik8【33，Large text compression benchmark，2009】）。理想情况下，我们希望在现代GPU内存允许的最大窗口尺寸和序列长度上训练模型。然而，我们发现模型需要大量的梯度更新来首先学习局部上下文，然后才能学习利用更长的上下文。为了适应这一点，我们采用了一种分阶段训练程序，在多个训练阶段中增加注意力窗口大小和序列长度。具体来说，在第一阶段，我们从一个短的序列长度和窗口大小开始，然后在每个后续阶段，我们将窗口大小和序列长度加倍，并减半学习率。这使得训练速度很快，同时将慢的部分（最长的序列和窗口大小）留到最后。我们在总共5个阶段上训练模型，起始序列长度为2,048，最后一个阶段的序列长度为23,040（每个阶段的详细配置及所有其他超参数见附录B）。

评估方法。我们使用长度为32,256的序列进行评估。遵循Dai等人【12，Transformer-XL: Attentive language models beyond a fixed-length context，2019，ACL】的方法，我们将数据集分割成大小为32,256、步长为512的重叠序列，并报告序列上最后512个token的性能。

5. 预训练和微调

核心动机与方法。本文的主要动机之一是为长文档任务开发一个合适的预训练模型。为此，我们在一个文档语料库上预训练了Longformer，并为六个任务（包括分类、问答和共指消解）对其进行了微调。最终的模型可以处理长达4,096个token的序列（是BERT的8倍）。我们使用掩码语言建模（MLM）目标来预训练Longformer。由于MLM预训练成本高昂，我们从RoBERTa【30，RoBERTa: A robustly optimized bert pretraining approach，2019，arXiv】发布的检查点继续预训练，只做了支持Longformer注意力机制所必需的最小改动。

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

位置嵌入。RoBERTa使用最大位置为512的学习式绝对位置嵌入。为了支持更长的文档，我们增加了额外的位置嵌入，以支持最多4,096个位置。为了利用RoBERTa的预训练权重，我们没有随机初始化新的位置嵌入，而是通过多次复制RoBERTa的512个位置嵌入来初始化它们。这么做的依据是，对BERT注意力头的分析显示出一种强烈的学习偏好，即关注局部上下文，包括前一个或后一个token【9，What does bert look at? an analysis of bert’s attention，2019，arXiv】。复制初始化除了在分区边界外，保留了这种局部结构。尽管简单，我们发现这种方法非常有效（见表5），使得Longformer预训练能够通过少量梯度更新快速收敛。

持续MLM预训练。我们在一个我们编制的长文档语料库上使用fairseq【35，fairseq: A fast, extensible toolkit for sequence modeling，2019，NAACL-HLT 2019】预训练Longformer（语料库细节见附录C）。我们训练了两种模型尺寸：一个基础模型和一个大型模型。两种模型都使用4,096的序列长度、64的批处理大小（2^18个token）进行了65K次梯度更新，最大学习率为3e-5，线性预热500步，然后是3次幂多项式衰减。其余超参数与RoBERTa相同。

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

7. Longformer-Encoder-Decoder (LED)

模型架构与动机。原始的Transformer【52，Attention is all you need，2017，NIPS】包含一个encoder-decoder架构，用于序列到序列任务（如摘要和翻译）。虽然仅有encoder的Transformer在多种NLP任务上有效，但预训练的encoder-decoder Transformer模型（例如BART【29，BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension，2020，ACL】和T5【42，Exploring the limits of transfer learning with a unified text-to-text transformer，2020，J. Mach. Learn. Res.】）在摘要等任务上取得了强大的结果。然而，这类模型无法高效地扩展到输入更长的序列到序列任务。

LED模型设计。为了便于对长序列进行序列到序列学习的建模，我们提出了一个Longformer的变体，它同时具有encoder和decoder的Transformer堆栈，但在encoder中不使用完整的自注意力，而是使用Longformer高效的局部+全局注意力模式。decoder则对整个编码后的token和先前解码的位置使用完整的自注意力。我们称这个模型为Longformer-Encoder-Decoder (LED)，它与输入长度呈线性扩展关系。

初始化与实现。由于预训练LED成本高昂，我们从BART初始化LED的参数，并在层数和隐藏层大小方面遵循BART的确切架构。唯一的区别是，为了处理更长的输入，我们将位置嵌入扩展到16K个token（BART为1K个token），并通过重复复制BART的1K个位置嵌入16次来初始化新的位置嵌入矩阵，做法与第5节中对RoBERTa的操作相同。遵循BART，我们发布了两种模型尺寸：LED-base和LED-large，它们分别在encoder和decoder堆栈中有6层和12层。

A4 实验环境

数据集:
* 自回归语言建模: text8 和 enwik8，均为包含1亿个维基百科字符的字符级语言建模数据集。
* 问答 (QA):
* WikiHop: 多跳问答，上下文段落数多，平均长度显著超过512个wordpiece。
* TriviaQA (Wikipedia setting): 开放域问答，文档较长。
* HotpotQA (distractor setting): 多跳问答，包含干扰段落。
* 共指消解: OntoNotes，文档长度变化大，但多数提及对（mention pair）距离较近。
* 文档分类:
* IMDB: 情感分类，约13.6%的文档超过512个wordpiece。
* Hyperpartisan News Detection: 新闻立场检测，文档相对较长，数据集规模小（645篇）。
* 摘要: arXiv summarization dataset，科学文献摘要，文档非常长，90分位长度为14.5K个token。


表 6: 数据集中上下文长度的平均值和95百分位数（以wordpiece为单位）。WH: WikiHop, TQA: TriviaQA, HQA: HotpotQA, ON: OntoNotes, HY: Hyperpartisan news

模型架构:
* Longformer (自回归LM): 小模型（12层，512隐藏层），大模型（30层，512隐藏层）。序列长度在分阶段训练中从2,048增加到23,040，评估时使用32,256。
* Longformer (预训练-微调): 基于RoBERTa-base和RoBERTa-large，将自注意力替换为窗口大小为512的滑动窗口注意力。最大序列长度为4,096。
* LED (摘要): 基于BART-base和BART-large架构。Encoder使用窗口大小为1,024的局部注意力和全局注意力。最大输入序列长度扩展到16K。

硬件配置:
* 自回归LM实验使用4或8块 48GB RTX8000 GPU。
* QA实验使用32GB V100 GPU。
* 分类、共指消解、HotpotQA实验使用RTX8000 GPU。

软件配置:
* 实现: 基于PyTorch，部分任务使用PyTorch-Lightning。共指消解任务为了复用代码，使用了PyTorch和Tensorflow的混合实现。
* 依赖库: fairseq用于预训练，TVM用于实现自定义CUDA核，apex用于混合精度训练。
* 基线模型: RoBERTa-base和RoBERTa-large。

A4 实验结果

自回归语言建模

  • 实验内容: 在text8和enwik8数据集上进行字符级语言建模，评估指标为BPC (bits-per-character)。
  • 实验结果:
    • 小模型在text8和enwik8上分别取得了1.10和1.00 BPC的SOTA结果（表2）。
    • 大模型在enwik8上表现优于Transformer-XL，与Sparse Transformer相当，但略逊于参数量为其两倍以上的模型（表3）。
  • 分析结论: Longformer的扩张滑动窗口注意力机制在长序列建模上是有效的。

    表 2: 小模型在text8和enwik8上的BPC

    表 3: 大模型在enwik8上的性能

自回归LM消融研究

  • 实验内容: 在text8开发集上，通过控制变量法研究注意力窗口大小配置和扩张注意力的影响。
  • 实验结果:
    • 窗口大小配置: 从底层到顶层递增窗口大小的性能最好，反之最差，固定大小居中（表4-上）。
    • 扩张注意力: 在两个注意力头上加入少量扩张，性能优于完全不使用扩张（表4-下）。
  • 分析结论: 递增的窗口大小和适度的扩张注意力是模型设计的关键。

    表 4: 上：跨层改变窗口大小。下：有/无扩张（在第一阶段的15万步时）

预训练效果

  • 实验内容: 评估从RoBERTa继续预训练Longformer的效果，指标为MLM BPC。
  • 实验结果:
    • 位置嵌入初始化: 复制RoBERTa的位置嵌入（BPC 1.957）显著优于随机初始化（BPC 3.036），表明了复制初始化的重要性（表5）。
    • 持续预训练: 经过65K步的预训练，Longformer-base的BPC从1.957降至1.705，表明模型学会了更好地利用滑动窗口注意力和长上下文（表5）。
  • 分析结论: 从现有模型继续预训练是有效的，特别是采用复制位置嵌入的初始化策略。

    表 5: RoBERTa和不同预训练Longformer配置的MLM BPC

下游任务微调

  • 实验内容: 在六个长文档任务（WikiHop, TriviaQA, HotpotQA, OntoNotes, IMDB, Hyperpartisan）上比较Longformer-base和RoBERTa-base。
  • 实验结果:
    • Longformer在所有任务上都一致地优于RoBERTa基线。在需要长上下文的任务（如WikiHop和Hyperpartisan）上性能提升尤为明显（表7）。
    • 在TriviaQA、HotpotQA等任务上提升较小，因为局部上下文已足够回答大部分问题或存在强监督信号。
    • Longformer-large在WikiHop和TriviaQA上取得了新的SOTA结果，比之前的方法有显著提升（表8）。在HotpotQA上，性能强劲，但略低于使用图神经网络（GNN）的SOTA模型（表9）。
  • 分析结论: Longformer作为RoBERTa的直接替代品，能有效利用长上下文，在多种长文档任务上提升性能。

    表 7: 在QA、共指消解和文档分类开发集上的微调结果摘要。比较了我们的Longformer-base和RoBERTa-base。TriviaQA、Hyperpartisan的指标是F1，WikiHop和IMDB使用准确率，HotpotQA是联合F1，OntoNotes是平均F1。

    表 8: 提交时（2020年5月）Longformer-large的排行榜结果。所有数字均为F1分数。

    表 9: HotpotQA在distractor设置测试集上的结果。Quark的测试结果不可用。所有数字均为F1分数。†表示同期的排行榜提交。

WikiHop消融研究

  • 实验内容: 在WikiHop开发集上，对Longformer-base的各个设计选择进行消融实验。
  • 实验结果:
    • 更长的序列长度（4096 vs 512）、使用全局注意力、为全局注意力使用独立的投影矩阵、进行MLM预训练以及更长的训练时间，都对性能有正面影响。
    • 当配置与RoBERTa-base相同时（序列长度512，n²注意力），Longformer性能略低于RoBERTa-base，证明性能提升并非来自额外的预训练，而是长上下文处理能力。
  • 分析结论: Longformer的各项设计，特别是长序列处理和全局注意力，都是其在长文档任务上取得成功的关键。

    表 10: WikiHop开发集消融实验

LED摘要任务

  • 实验内容: 在arXiv长文档摘要数据集上评估LED模型，指标为ROUGE。
  • 实验结果:
    • 仅从BART初始化而未经额外预训练的LED-large (16K序列长度) 在arXiv数据集上取得了SOTA结果，ROUGE-1/2/L分别为46.90/18.84/43.08，略微超过了经过专门摘要预训练的BigBird模型（表11）。
    • 增加输入序列长度能显著提升摘要质量（图3）。
  • 分析结论: LED能够有效处理超长文档的序列到序列任务。处理更长输入的能力是其性能优越的关键，即使没有进行针对性的预训练。

    表 11: Longformer-Encoder-Decoder (LED)在arXiv数据集上的摘要结果。从左到右的指标是ROUGE-1, ROUGE-2和ROUGE-L。

    图 3: 在arXiv验证集上改变输入大小时LED的ROUGE-1和ROUGE-2分数。

A5 结论

本文提出了Longformer，一个基于Transformer的模型，能够可扩展地处理长文档，并简化了一系列文档级NLP任务的执行，避免了对长输入进行分块/截断或使用复杂架构来组合信息的需要。Longformer采用了一种结合了局部和全局信息的注意力模式，同时其计算复杂度与序列长度呈线性关系。

主要成果:
1. 在字符级语言建模任务text8和enwik8上取得了SOTA结果。
2. 经过预训练的Longformer在长文档任务上持续优于RoBERTa，并在WikiHop和TriviaQA上刷新了SOTA记录。
3. 提出了LED，一个Longformer的encoder-decoder变体，用于序列到序列任务，并在arXiv长文档摘要任务上取得了SOTA结果。

未来工作:
作者计划研究其他的预训练目标（特别是针对LED），增加序列长度，并探索其他可能受益于该模型的任务。

A6 附录

A 实现细节

实现方法对比。实现Longformer的扩张滑动窗口注意力需要一种带状矩阵乘法，而这在PyTorch/Tensorflow等现有深度学习库中没有直接支持。图1比较了三种实现方式的运行时间和内存占用。
* Longformer-loop：这是一个朴素的实现，通过循环单独计算每个对角线。它内存效率高，因为它只计算非零值，但速度非常慢，不适合实验，仅用于测试。
* Longformer-chunks：此实现仅支持非扩张的情况。它将Q和K矩阵分块成大小为$w$、重叠为$\frac{1}{2}w$的块，然后将这些块相乘，并掩码掉不需要的对角线。这种方法计算效率很高，因为它利用了PyTorch的单个矩阵乘法操作，但内存消耗是一个完美优化实现的两倍，因为它计算了一些零值。由于其计算效率，此实现最适合预训练/微调场景。
* Longformer-cuda：这是一个使用TVM【5，TVM: An automated end-to-end optimizing compiler for deep learning，2018，OSDI】实现的自定义CUDA核。它是一个功能完整的注意力实现（不像Longformer-chunks那样受限），内存效率最高，并且速度与高度优化的全自注意力相当。主要用于自回归语言建模实验，因其内存效率高（允许处理最长序列）且支持扩张（字符级LM实验需要）。

Tensor Virtual Machine (TVM)。我们使用TVM【5，TVM: An automated end-to-end optimizing compiler for deep learning，2018，OSDI】构建了自定义CUDA核，它是一个深度学习编译器堆栈，能将高级函数描述编译为优化的设备特定代码。我们用高级Python结构描述了带状矩阵乘法，然后TVM生成相应的CUDA代码并为GPU进行编译。

B 字符级LM超参数

数据集与模型基础。我们在text8和enwik8上进行评估，这两个数据集都包含1亿个维基百科字符，按90M/5M/5M的比例划分为训练/开发/测试集。我们的模型只规定了自注意力组件的工作方式，与transformer模型的其他设计选择无关。我们的实现基于Transformer-XL【12，Transformer-XL: Attentive language models beyond a fixed-length context，2019，ACL】的代码，但禁用了记忆机制。我们使用与Dai等人【12】相同的正弦权重相对位置嵌入。我们使用了两种模型尺寸：一个小的（12层，512隐藏尺寸）和一个大的（30层，512隐藏尺寸）。

训练技术。我们使用apex进行混合精度训练（16位和32位浮点数）以减少内存消耗并加速训练，但注意力计算保持在fp32以避免数值不稳定。我们使用梯度检查点【6，Training deep nets with sublinear memory cost，2016，arXiv】来减少内存使用，并在48GB的RTX8000 GPU上运行实验。所有超参数和阶段配置都列在表12中。我们的CUDA核支持自回归模式，其中每个token只关注前一个窗口的token。我们的实现还包括一个与扩张滑动窗口注意力兼容的相对位置嵌入版本。

实验耗时与超参数搜索。小模型实验在4个RTX8000 GPU上运行了16天。大模型实验在8个RTX8000 GPU上运行了13天。我们的超参数搜索主要是在text8上运行配置15万步，类似于表4中的消融研究。我们试验了绝对位置嵌入和学习位置嵌入，dropout值（小模型[0.1, 0.2]，大模型[0.1, 0.4]），pre-layernorm和post-layernorm【58，On layer normalization in the transformer architecture，2020，arXiv】，第一阶段的学习率（LR）值[2.5e-5, 5e-4, 1e-4]以及常数和余弦LR调度，以及不同的扩张配置。表12中报告的每阶段梯度更新次数是通过运行每个阶段直到验证BPC不再改善来确定的。


表 12: 字符级语言建模最佳性能模型的超参数

C 预训练数据

数据构成。为了让模型在预训练中学习长距离依赖，我们构建了一个长文档语料库。其中一些数据源也包含在原始RoBERTa预训练中，包括Books语料库【63，Aligning books and movies: Towards story-like visual explanations by watching movies and reading books，2015，ICCV】和英文维基百科。我们额外加入了三分之一的Realnews数据集【62，Defending against neural fake news，2019，NeurIPS】子集（文档长度超过1200个token）以及三分之一的Stories语料库【48，A simple method for commonsense reasoning，2018，arXiv】。我们的目标是混合长短文档，既让模型学习长距离依赖，又不会忘记原始RoBERTa预训练中的信息。预训练数据的统计信息如表13所示。


表 13: 预训练数据

D 任务特定模型细节

通用设置。所有的问答和分类模型都使用PyTorch-Lightning实现。除Hyperpartisan news（我们随机划分为80/10/10的训练/开发/测试集）外，我们都使用官方的数据集划分。


表 14: 问答模型的超参数。所有模型都使用类似的调度器，包含线性预热和衰减。

WikiHop。输入格式为[q] question [/q] [ent] candidate1 [/ent] ...，后面接上用</s>分隔的上下文。对问题和候选答案序列使用全局注意力。预测时，在每个[ent] token上附加一个线性层输出logit，然后对每个候选答案的所有logit求平均，通过softmax后用交叉熵损失进行训练。

TriviaQA。输入格式为[s] question [/s] document [/s]。对所有问题token使用全局注意力。预测时，增加一个层来预测答案的起始和结束位置，使用Clark和Gardner (2017)【8】的损失函数，该函数类似OR逻辑，模型只需正确预测一个答案片段即可。

HotpotQA。采用两阶段模型。第一阶段，将问题和所有10个段落拼接成一个长上下文输入，使用全局注意力关注问题、段落标题和句子token，预测相关段落。第二阶段，在筛选出的相关段落（最多5个）上进行答案抽取和证据句识别。模型以多任务方式联合训练，预测相关段落、证据句、答案片段和问题类型（是/否/片段）。

共指消解。直接改编自Joshi等人 (2019)【23】的基于BERT的模型。将文档分割成最长为4096（Longformer）或384（RoBERTa）的非重叠片段，然后将激活值拼接起来送入粗到细的聚类阶段。未使用全局注意力。实现上，为了复用Tensorflow中的复杂逻辑，采用了一个PyTorch和Tensorflow混合的“超级hack”方案，在两者之间传递激活值和梯度。

文本分类。在[CLS] token上使用全局注意力，并附加一个分类头，使用二元交叉熵损失。训练设置包括Adam优化器、批量大小32、以及占总训练步数10%的线性预热和衰减。