Longformer详解——从Self-Attention说开去

1.Longformer的应用场景

为了理解Longformer的原理，我们最好首先从为何需要使用Longformer开始说起。（这里默认各位已经对Self Attention等基础知识有一定的了解）

我们以一个简单的场景为例：

在这个例子中，我们共有六个Token，每个Token的维度是768维，当我们对这几个Token进行处理，得到的Q/K/V如下所示：



有了Q/K/V之后，下一步就可以计算每个Token对应的Attention：

注意，Q乘以K的转置可以很形象地表示为下面这张图：

换言之，我们在Self Attention中，通过前面这一系列操作，得到的是一个(Token_size, Token_size)的矩阵。这里我们只有六个Token，一旦我们需要处理较长的文本，那么就会带来巨大的计算量。

于是，为了解决这一困境，并使得Transformer能够用于较长的文本，Longformer的作者提出了下面这几种特殊的attention机制，这便是Longformer的核心内容。

2.Longformer的核心内容

书接上文，Longformer的核心机制其实就是上面提到的三种attention窗口：

下面我们对这三种窗口进行一一介绍：

a.滑动窗口注意力

这种注意力机制不妨类比为卷积核，每一个卷积核的大小就是我们的感受野。总之，如果我们仔细观察的话，不难发现，这种方法的实质就是设定了一个固定大小的窗口\(ω\)，它规定了序列中的每个Token都只能够看到\(ω\)个相邻的Token，其左右两侧各能看到\(\frac{1}{2}ω\)个Token（图中的例子，\(ω=6\))，这样一来，时间复杂度就由原来的\(O(n×n)\)变为了\(O(n×ω)\)，当然，\(ω<<n\)。当然，我们无需担心这种设定无法建立整个序列的语义信息，因为transformer模型结构本身是层层叠加的结构，模型高层相比底层具有更宽广的感受野，自然能够能够看到更多的信息，因此它有能力去建模融合全部序列信息的全局表示，就行CNN那样。

通过这种设定Longformer能够在建模质量和效率之间进行一个比较好的折中。

b.膨胀滑动窗口注意力

这种注意力机制和上一种有什么区别呢？

答案很明显，它膨胀了。

好，我们来看下一个注意力机制（开玩笑的:D）

这种注意力窗口其实是为了弥补上一种窗口所造成的长度缺失，在不增加时间复杂度的前提下，我们会在两个Token之间设置一个大小为\(d\)的空隙。那么感受野的范围就可以扩展到\(d × ω\)。（以图中为例，\(d = 1, ω=6\))。

作者在文中提到，在进行Multi-Head Self-Attention时，在某些Head上不设置Dilated Sliding Window以让模型聚焦在局部上下文，在某些Head上设置Dilated Sliding Window以让模型聚焦在更长的上下文序列，这样能够提高模型表现。

c.全局注意力

这是重头戏。

首先，这是一个颇具迷惑性的名字，这里的全局只是部分的全局，我们在这里只设定某些特定的Token能够看到其余所有的Token，而对于另外一些不是太重要的Token，我们还是采取滑动窗口注意力。注意！这里“特定的Token” 会随着任务的不同而变化，打个比方，对于分类任务，带有全局视角的token就是“CLS”；而对于QA任务，那么带有全局视角的token就是Question中所包含的这些token。

之所以需要提出这种注意力机制，是因为这之前的两种变体无法完全适应task-specific的任务。

换言之，这种注意力机制的实质就是，将特殊位置的token“暴露”给其它的token。（等等，黑暗森林！？三体人狂喜）

Longformer的实际效果

a.Longformer在减少计算量这一方面的效果

Longformer的作者在论文中对三种Longformer的实现方式与Self-attention进行了对比：



Longformer-loop是一种Pytorch实现，它支持膨胀滑动窗口

Longformer-chunks不支持膨胀滑动窗口，但计算速度很快

Longformer-cuda则是作者使用TVM实现的CUDA内核方法

左图是计算时间的对比，右图是计算所需要的内存量对比，本着男左女右，女士优先的原则，我们先看右边这幅图。可以看出，在序列长度较长的情况下，Longformer确实可以做到减少计算量的效果。但是，左图中的计算时间也可以看出，尽管Longformer减少了计算所需的内存，但计算时间并没有比一般的Self-Attention快多少，尤其是Longformer-loop，即使序列长度不长，但它还是给出了一个令人叹为观止的计算时间，这可太“酷”了。

b.Longformer在不同数据集上的实验结果

b.1.在text8与enwiki8两个字符级任务上的实验

（这里的测试指标是BPC，BPC越小，就代表性能越好）

此外，在这个实验中，作者还使用了两种不同大小的模型：

small model:12 layers, 512 hidden size

large model:30 layers, 512 hidden size

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

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

b.2.Pretraining and Finetuning

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

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

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

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

而下表是Longformer在具体的下游任务上的训练结果与RoBERTa的一个对比：



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

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

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

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

当然，结果仅供参考，具体任务具体分析，实践才能出真知。

最后附上我本人使用的Longformer中文预训练模型地址

https://huggingface.co/ValkyriaLenneth/longformer_zh