Longformer：超越RoBERTa，为长文档而生的预训练模型

作者 | ChrisCao、小轶

编辑 | 丛 末

传统Tranformer-based模型在处理长文本时有着天然的劣势。因为传统模型采用的是“全连接”型的attention机制，即每一个token都要与其他所有token进行交互。其attention复杂度高达 。此前的解决办法是将长文切分为若干个较短的text span，然后逐个处理。这就导致不同的text span之间无法进行交互，因而必然存在大量information loss。当然，我们也可以通过添加一些其他机制来加强这种text span之间的交互。但这种新增机制实现起来通常比较复杂，而且往往是task-specific的，通用性不强。

原文链接：https://arxiv.org/pdf/2004.05150.pdf

Github：https://github.com/allenai/longformer

本文提出的一种可高效处理长文本的升级版Transformer——Longformer，改进了Transformer的传统attention机制：对于每一个token，只对固定窗口大小的附近token计算local attention，并结合具体任务，计算少量的global attention。该方法的优点包括：

• 复杂度低，将attention机制的复杂度降至
• 通用性强，可用于各类文档级任务
• 部署容易，作者在cuda内核上直接实现了Longformer的attention pattern，并提供了开源代码。

Longformer在两个字符级语言建模任务上都取得了SOTA的效果。并且，作者用Longformer的attention方法继续预训练RoBERTa。训练得到的语言模型在多个长文档任务上进行fine-tune后，性能全面超越Roberta。该预训练模型也已开源，大家可以很方便地直接应用于自己的任务。

1

模型

作者共提出了三种新的attention pattern，来降低传统self-attention的复杂度，分别是滑窗机制、膨胀滑窗机制、融合全局信息的滑窗机制。下图展示了传统attention与这三种attention pattern的示意图。接下来将为大家分别讲解。

1. 滑窗机制（Sliding window）： 对于每一个token，只对其附近的w个token计算attention计算复杂度与文本序列长度 成线性关系，为 。作者认为，根据应用任务的不同可以对Transformer每一层施以不同的窗口大小 ，对模型表示能力可能有潜在帮助。

读到这里的，大家可能和我一样，误认为这个窗口 应该比较小，估计在16~64这个量级。但看到实验部分会发现，作者在具体实现的时候，设置的窗口大小 为512，和Bert的Input限制完全一样。所以，大家不要存有“Longformer比Bert还要更轻量”的错觉。

2. 膨胀滑窗机制（Dilated sliding window）： 在对每一个进行token编码时，普通滑窗机制只能考虑到长度为 的上下文。作者进一步提出膨胀滑窗机制，在不增加计算负荷的前提下，拓宽模型“视场”。其做法借鉴了空洞卷积的思想[1]。如下图所示，在滑动窗口中，被attend到的两个相邻token之间会存在大小为d的间隙。当transformer的层数为l时，则视场范围可达到 。实验表明，由于考虑了更加全面的上下文信息，膨胀滑窗机制比普通的滑窗机制表现更佳。

3. 融合全局信息的滑窗机制（Global+sliding window）： 我们知道Bert一类的语言模型在应用于具体任务时，实现方式略有不同。比如，对于文本分类任务，我们会在文本序列前添加[CLS]这一特殊token；而对于QA类任务，则会将问题与文本进行拼接后输入。在Longformer中，作者也希望能够根据具体任务的不同，在local attention的基础上添加少量的global attention。比如，在分类任务上就会在[CLS]处添加一个global attention，而在QA任务上会对question中的所有token添加global attention。如下图所示，对于添加了global attention的token，我们对其编码时要对整个序列做attention。并且，编码其他所有token时，也都要attend到它。

2

实验

1、自定义CUDA内核

由于现有的深度学习库中并没有能直接实现膨胀滑窗机制的接口，为此作者直接自定义了CUDA内核操作，用于实现Longformer的attention pattern[2]。如下图所示，Longformer的内存消耗与文本长度成线性关系（红线）。用自定义CUDA来实现Longformer，相比于用Naive Pytorch来实现（蓝线），运行速度加快了六倍。

2、Longformer在字符级别任务上的表现

作者在text8和enwik8两个字符级任务上对Longformer进行实验。实验中，模型每一层采用了不同的窗口大小：底层使用较小的滑窗，以建模局部信息；在高层使用较大的滑窗，以扩大感受野。训练时，理想状况下当然是希望使用GPU所能承受的最大的window size和sequence len。但为了加快训练速度，作者采用的是一种阶段式的训练方式：在学习更长的上下文之前，先学好局部的上下文。在第一阶段，先设置较短的序列长度和窗口大小。在后续阶段，window size和sequence length增加一倍，学习率减半 实验结果如下图所示，Longformer在这两个数据集上皆达到了SOTA效果（注：测试指标为BPC，bits-per-character；BPC越小，性能越优）。

作者通过实验，对滑窗机制的设置进行了进一步的讨论。如下表所示：

• 表中第一组实验（前三行）讨论的是：如果transformer的不同层采用不同窗口大小，是否可以提高性能？实验结果表明，由底层至高层递增窗口大小，可提升性能；递减则反而性能降低。
• 第二组实验（后两行）是对膨胀滑窗机制的消融实验，证明了增加间隙后的滑窗机制，性能可以有小幅度提升

3、Longformer用于预训练

1）MLM Pretraining

作者采用Longformer的方法在以下四个文档级语料上进行预训练，从而得到适于文档级NLP任务的语言模型。作者并没有完全从头预训练一个随机初始化的模型，而是以RoBERTa为基础，采用MLM(masked language modeling)的方法继续预训练。预训练时，每一层都采用固定的大小为512的滑动窗口，暂不添加global attention。为支持长文本，论文作者把position embedding扩展到了4096个。

预训练结束后，在多个文档级任务上再进一步对预训练模型做fine-tuning。fine-tuning时会根据任务增加global attention：共设置两套映射矩阵，一套用于局部自注意力，另一套用于全局注意力。实验表明，Longformer全面超越了RoBERTa的表现。

2）消融实验

为了证明Longformer更优异的性能并不是因为对RoBERTa额外的预训练带来的，作者做了一组消融实验。采用了与RoBERTa完全相同的序列长度和attention机制，在继续预训练后并没有在文档级任务上取得更优的性能。

参考文献：

[1] Dilated CNN：CV中常用的图像编码方式，可拓宽模型编码时的感受野。paper原文链接：https://arxiv.org/abs/1511.07122。另可参考知乎讨论“如何理解空洞卷积（dilated convolution）？”，链接https://www.zhihu.com/question/54149221。

[2] TVM: 关于如何自定义CUDA内核：这里作者使用了TVM (tensor virtual machine)(tvm.apache.org)，2018年由华盛顿大学的SAMPL组贡献的开源项目。TVM为不同的深度学习框架和硬件平台实现了统一的编译栈，从而实现将不同框架的深度学习模型到硬件平台的快速部署。

---

招 聘

AI 科技评论希望能够招聘 科技编辑/记者

办公地点：北京/深圳

职务：以跟踪学术热点、人物专访为主

工作内容：

1、关注学术领域热点事件，并及时跟踪报道；

2、采访人工智能领域学者或研发人员；

3、参加各种人工智能学术会议，并做会议内容报道。

要求：

1、热爱人工智能学术研究内容，擅长与学者或企业工程人员打交道；

2、有一定的理工科背景，对人工智能技术有所了解者更佳；

3、英语能力强（工作内容涉及大量英文资料）；

4、学习能力强，对人工智能前沿技术有一定的了解，并能够逐渐形成自己的观点。

感兴趣者，可将简历发送到邮箱：jiangbaoshang@yanxishe.com

点

击"阅读原文"，直达“KDD交流小组”了解更多会议信息。