《Longformer: The Long-Document Transformer》论文笔记 - 知乎

2020/04/24 07:59
阅读数 2.6K

论文地址

https://arxiv.org/abs/2004.05150arxiv.org

The Long-Document Transformer,顾名思义,就是应用在长文本场景下的Transformer。

本文作者是这几位大大:Iz Beltagy、Matthew E. Peters、Arman Cohan

他们都来自Allen Institute for Artificial Intelligence(AI2)

背景

先介绍一下这篇论文的背景。

我们都知道Transformer架构很强,成功的部分原因在于自注意力机制,自注意力机制使网络能够从整个序列中捕获上下文信息

但是虽然自注意力机制很有效,但它所需的内存和算力会随着序列长度呈平方增长,这使得当前硬件在处理长序列的情况下不可行,或者说非常昂贵、代价很大。

现有的方法是将上下文缩短或者划分成为较小的序列,以限制这些序列在512的长度以内。但是这种划分可能导致重要的信息丢失。

那么,为了解决这个问题,作者就提出了Longformer

它的attention包括窗口化的局部上下文的self attention和由终端任务激活的全局attention。

局部attention用来建立局部的上下文表示,全局attention用来建立完整的序列表示以进行预测。

模型

前面提到,现有的注意力计算方法时间和空间复杂度都是O(n²),这使得难以在长序列上进行训练。

为了解决这个问题,作者根据一个“attention pattern”来稀疏完整的自注意力矩阵

图中第一个是全自注意力矩阵,后面三个是作者提到的稀疏后的注意力矩阵

下面我们分别来介绍这三个attention pattern。

Sliding window attention

首先介绍一下滑动窗口attention,顾名思义,就是围绕每一个token采用固定大小的窗口计算局部注意力。

假设窗口大小是w,序列长度是n,那么这个pattern的计算复杂度就是O(n*w)

为了效率起见,w相比于n应该是个比较小的数值。

我们知道transformer无论编码还是解码器都是由l层堆叠形成的。

那么作者认为,根据应用程序的不同,为每一层使用不同的w值对平衡效率和模型的表示能力可能会有所帮助。

Dilated Sliding Window

空洞滑动窗口

窗口的缝隙大小是d,假设window size是w,transformer的层数是l,那么窗口能覆盖到的接受范围就是l*d*w。

那么,在mutilhead attention中,作者设置允许一些没有空洞的head专注于局部上下文,而其他具有空洞的head专注于较长的上下文,最终发现这样的做法能提升整体的表现

图源: https://www.pianshen.com/article/4552301946/

这个方法其实参考的CNN里头的空洞卷积。

Global Attention

最后,介绍一下全局attention。

对于一些任务,比如说QA,需要通过question去找document里的答案,因此局部注意力在这种情况下就不适用。

因此,作者就在一些预先选择的位置上(图中标红)添加了全局注意力。在这些位置上,会对整个序列做attention。

那么这些位置怎么去选择呢,其实要根据具体任务决定。

作者举了几个例子,在分类任务上,这个global attention就落在[CLS]标签上,在QA任务上,就在整个问句上计算global attention。

可以看出,global attention是视具体任务而定的,换个任务可能之前的做法就不适用了,但作者认为,这仍然比现有的trunk或者shorten的做法要简单很多。

Tensor Virtual Machine (TVM)

另外,值得一提的是,作者用TVM构建了自己的自定义CUDA kernel。这使得longformer的速度与full self-attention的操作几乎一样快,而且显存占用要小得多。

实验设置

下面先来看一下实验的一些设置.

模型采用字符级别的自回归语言模型。

数据集用的是text8和enwik8 (Mahoney, 2009)。

训练:从理想状况来讲,训练时应该选用gpu所能承受的最大的window size和seq len。但是作者发现,在学习更长的上下文之前,模型需要大量的梯度更新才能首先学好局部的上下文。

因此这里采用一种阶段式的训练方式,在第一阶段,从较短的序列长度和窗口大小开始,然后在每个后续阶段,将窗口大小和序列长度增加一倍,并将学习率减半。这样可以快速训练,同时将最慢的部分放到最后。

作者这里一共训练了5个阶段,第一个阶段seq len是2048,最后一个阶段是23040(受到gpu显存的限制,不能再高了)

评估过程中,作者将数据集分为大小为32,256的重叠序列,步长为512,并报告序列上最后512个token的性能。

实验结果

作者这里使用了两种不同的模型大小,小的有12层,大的有30层。

在小的模型上,long former比其他模型都要好。这也证明了这个模型的有效性。

在大的模型上,比18层的transformerxl要好,跟第二个和第三个齐平,不如第四、五两个。但作者说明,这两个模型并不适用于pretrain-finetune的模式。


下面这个表展示了作者在不同attention pattern的配置上的实验结果。

表格的上半部分展示了,递增窗口大小的表现最好,递减窗口大小表现最差,使用固定窗口的表现介于两者之间。

表格的下半部分展示了使不使用空洞窗口的区别,可以看出对两个head增加一些空洞比完全不使用空洞窗口表现要好一些。

Pretraining and Finetuning

本文的预训练使用的同样是MLM,并且是在RoBERTa的checkpoint后继续训练。

在所有layer上使用512window size的滑动窗口attention,这么设置是为了与RoBERTa的seq len相匹配。

为了支持长文本,作者添加了额外的position embedding到4096的大小,此外,为了利用RoBERTa的权重,作者通过多次复制其512个位置嵌入来对其进行初始化。

下面这个表中展示了预训练使用的数据集。

上面两个是RoBERTa预训练同样用到的,下面两个是作者新增的,并且各选了三分之一用作预训练。

MLM Pretraining

第一行在RoBERTa base上的1.846与RoBERTa论文中的1.880相当,这表明本文的训练语料与训练RoBERTa的语料有相似的分布。

第二行展示了longformer在预训练之前使用随机初始化的position embedding的表现。下一行是在预训练之前使用从RoBERTa中拷贝的position embedding的表现。

RoBERTa初始化的表现和RoBERTa本身的表现差距较小,这也表明了滑动窗口attention是可以使用RoBERTa的权重的。

最后两行分别是训练两千步和6万5千步之后的结果,可以看出效果是越来越好,这也说明了模型正在学习如何更好地利用滑动窗口attention和更长的上下文。

Finetuning

本文finetune使用的数据集

第一行是数据集里头序列的平均长度,第二行是第95个百分位的数据的长度。

下面这个表是longformer在具体的下游任务上的训练结果与RoBERTa的一个对比。

其中QA任务中,在问题和候选答案的位置使用全局attention。

共指消解任务上没有使用全局attention。

文本分类任务中,在[CLS]标签上用了全局attention。

可以看出,所有结果都是优于RoBERTa的。


最后

Longformer的代码和预训练模型也已经公布啦!

https://github.com/allenai/longformergithub.com

展开阅读全文
打赏
0
0 收藏
分享
加载中
更多评论
打赏
0 评论
0 收藏
0
分享
返回顶部
顶部