论文标题：

Longformer: The Long-Document Transformer

论文作者：

Iz Beltagy (AllenAI), Matthew E. Peters (AllenAI), Arman Cohan (AllenAI)

论文链接：

https://arxiv.org/pdf/2004.05150.pdf

代码链接：

https://github.com/allenai/longformer

尽管Self-Attention可以建模长距离文本，但由于时空复杂度都为平方级，从而它无法一次处理较长的文本。

BERT的处理方法是，将文本截断为512的长度，如此一来，虽然模型可以分段处理，但数据的预处理也是相当麻烦的。

本文采用局部自注意力和全局自注意力结合的方式（或简称为稀疏注意力），同时为其进行了CUDA优化，从而使得模型最大能够容纳长度上万的文本，同时还能实现更好的结果。

自注意力机制的缺陷

众所周知，Transformer 的 Self-Attention 有一大优点一大缺点：优点是它能在理论上建模任意长的文本，缺点是实践中过长的文本会导致内存爆炸。

这是因为，自注意力的时空复杂度都是平方级的，文本越长，空间消耗就显著更大，所以在实际使用的时候，通常把句子长度截短到512，然后分段送入模型。

这种方法虽然可以在一定程度上解决内存不足的问题，但是：

• 它回避了Self-Attention的固有缺陷，只是一种暂缓之计；

• 这种方法非常麻烦，需要预处理数据集；

• 很多任务都包含大量长文本，实际需要要求我们能够一次性处理它们。

基于以上现状和自注意力的固有缺陷，本文提出结合局部自注意力和全局自注意力，在提高效率的同时又不失效果。

和之前的许多工作相比，本文没有在具体技术上有太大创新，但是本文首次将这种稀疏自注意力应用到了三种包含长文本的任务上：问答、共指消解和文档分类，并证明了其有效性。

此外，本文还基于TVM实现了稀疏自注意力的CUDA优化并开源，造福广大研究者！

总的来说，本文有如下贡献：

• 通过大量实验（语言模型、问答、共指消解、文档分类）证明了稀疏自注意力的有效性，甚至优于原版自注意力；

• 基于TVM实现了稀疏自注意力的CUDA优化，加快了实际应用的使用速度；

• 开源代码，便于复现论文结果。

稀疏自注意力

在此之前已经有许多关于稀疏自注意力的工作，本文使用简单的CNN局部自注意力实现，如下图所示。

图1是原版自注意力（每个字符都能关注其他所有字符），图2是基于CNN的局部自注意力（以窗口的形式移动，每个窗口内的字符可以互相关注），图3是基于Dilated CNN的局部自注意力（跨越式滑动窗口），图4是结合局部自注意力和全局自注意力的稀疏自注意力。

其中，“全局”的意思是，让小部分字符既能关注其他所有字符，又能让其他所有字符关注它。如果这些字符的数量很小，那么整体也可以看成是O(n)复杂度的。

在具体使用的时候，本文设置两套映射矩阵Q,K,V，一套用于CNN局部自注意力Q_s,K_s,V_s，另一套用于全局注意力Q_g,K_g,V_g。

同时，在模型底层使用较小的CNN窗口大小，以建模局部信息；在高层使用较大的CNN大小，以扩大感受野。

使用图4的稀疏自注意力，可以看到，有大量的“白色方块”，表示不需要关注，而随着文本长度的增加，这种白色方块的数量会呈平方级增加，所以实际上我们需要的“绿色方块”数量是很少的。

为了让稀疏自注意力适应当前深度学习框架，本文还基于TVM实现了对其的CUDA优化，极大提升了实际使用的效率。

下图是使用CUDA优化的稀疏自注意力（Longformer）、原版自注意力（Full Self-attention）和未使用CUDA优化的稀疏自注意力（Naive Pytorch）在时间复杂度和空间复杂度上的比较。

可以看到，LongFormer能够取得最佳的时空复杂度。

实验

本文在语言模型、问答、共指消解和文档分类任务上进行实验，验证LongFormer在长文本上的效果。

语言模型任务

语言模型使用数据集text8和enwik8，按阶段训练LongFormer。在测评时，在长度为32256长的文本上进行，其他设置详见原文。

下表是实验结果。可以看到，在参数量相近的情况下，Longformer取得了略好的结果，但由于一次可容纳的最大文本长度大大增加了，测试的时间也会有所缩短（表中未给出）。

下表是LongFormer中局部自注意力窗口大小随模型深度变化情况带来的最终效果，可以看到，由底层到高层降低感受野的方法效果最差，而由底层到高层增加感受野的方法效果最好；并且为2个头使用Dilated的方法可以更好一些。

问答、共指消解与文档分类任务

在进行这三个任务之前，LongFormer首先需要预训练。为此，模型在RoBERTa上继续预训练，预训练语料主要包含了一些长文本。

在预训练之后，再在具体任务上微调。这里略去预训练细节，有兴趣的读者可以参考原文Section 5。

问答的数据集有WikiHop, TriviaQA和HotpotQA，共指消解的数据集有OntoNotes，文档分类的数据集有IMDB和Hyperpartisan。

问答的三个数据集都有很长的文本长度，RoBERTa需要对其切分，而LongFormer直接把一个问答任务中的问题和所有文档拼接起来，不需要再切分。

对另外两个任务也是同样的操作。基线模型是RoBERTa。

下面是在这三个任务上总的结果。

可以看到，LongFormer在这些任务上都或多或少超过了基线模型，尤其是在需要长距离上下文的数据集上，如WikiHop和Hyperpartisan。

下表是LongFormer在HotPotQA上的表现，当前该数据集表现最佳的模型是多阶段+图网络，而从表中可以看到，单阶段的LongFormer也能有69.5的F1值，二阶段的模型有71.4的F1值，已经比较接近当前最佳结果。

最后来看看因子分析实验，数据集为WikiHop。从表中可以看到，对LongFormer影响最大的是全局注意力、局部注意力和两套映射矩阵。

值得注意的是，当LongFormer和RoBERTa同样采用512的长度和原版自注意力时，效果还不仅RoBERTa，尽管LongFormer在更多的语料上预训练了，这说明LongFormer的效果增强并不来自于预训练。

小结

本文提出了一种综合的稀疏自注意力模型——LongFormer，利用局部自注意力和全局自注意力在大幅降低空间复杂度的情况下提高模型效果。

LongFormer可以直接吞进长文本，而不需要像BERT类模型一样对其强行截短，这就有利于那些需要建模长距离语义的任务。

在语言模型、问答、共指消解和文档分类任务上，LongFormer展现出了它强壮的一面，也为稀疏自注意力的研究带来新的可能性，我们期待它能在更多任务上取得好的表现。

????

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」，小助手将把你带入 PaperWeekly 的交流群里。