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 \times n)$ 变为了 $O (n \times ω)$ ，当然， $ω << n$ 。当然，我们无需担心这种设定无法建立整个序列的语义信息，因为transformer模型结构本身是层层叠加的结构，模型高层相比底层具有更宽广的感受野，自然能够能够看到更多的信息，因此它有能力去建模融合全部序列信息的全局表示，就行CNN那样。

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

b.膨胀滑动窗口注意力

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

答案很明显，它膨胀了。

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

这种注意力窗口其实是为了弥补上一种窗口所造成的长度缺失，在不增加时间复杂度的前提下，我们会在两个Token之间设置一个大小为 $d$ 的空隙。那么感受野的范围就可以扩展到 $d \times ω$ 。（以图中为例， $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，即使序列长度不长，但它还是给出了一个令人叹为观止的计算时间，这可太“酷”了。