Press "Enter" to skip to content

Longformer: 局部Attention和全局attention的混搭

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

 

Overall

 

Bert模型虽然很强大,但双向attention它的时间和空间复杂度呈N^2的趋势增长,所以最初的Bert模型能够处理的最长长度是512。

 

而在这一限制的基础上,如果想处理较长的序列,就需要用妥协的方式:

 

直接截断成512长度的。这点普遍用于文本分类问题。

 

截成多个长度为512的序列段(这些序列段可以互相overlapping),每个都输入给Bert获得输出,然后将多段的输出拼接起来。

 

两个阶段去解决问题,一般用于Question-Answer问题,第一个阶段去选择相关文档,第二个阶段去找到对应的answer。

 

无论哪种方式,毫无疑问都会带来损失。如果能直接处理长序列就好了。

 

关于这一点,我们在Reformer和自适应宽度注意力这两篇中各自讲述了办法去解决:

 

Reformer使用局部敏感哈希来解决性能问题,动机就在于attention起作用的在于top-N而不是全部。

 

自适应宽度使用一种动态窗口的方法来解决,动机在于attention有可能只attend最近的一些context。

 

而今天的这篇论文[1],用了一种更加直接的方式去对attention进行改造。那就是局部attention和全局attention的结合,局部attention用来捕捉局部信息,一般用于底层,这点和自适应宽度其实有点像。全局attention则捕捉全局信息,用于高层,目的在于保持住所有的信息。除此之外,论文还提供了改造后的attention的C++实现优化,使之相对于pytorch中的naive实现有了很大的提升。

 

下图中有个对比,可以看到,在计算速度上,Longformer与Full attention持平,但Full attention在超过一定长度后因为内存问题就无法运行了,论文提供的实现要比pytorch的原始实现快6倍。而在内存上,Longformer则是线性增长的。

 

 

滑动窗口Attention

 

论文的核心就在于局部Attention的设计,在这里采用的是滑动窗口来做,滑动窗口的大小为w,那幺每个位置只attend前后w/2个位置。如下图b所示。

 

 

因为模型都是多层叠加的,所以层级越高,attend的视野域就越广。如果w=3,那幺第一层只能注意3个位置,但到第二层能注意到第一层输出的三个位置,换算到第一层的输入,就是5个位置。所以随着层级越高,理论上每个位置注意到的区域就越大,所能存储的信息就越接近全局attention时的状态。

 

旁白君:这点和卷积神经网络很像。

 

另外,每一层的w其实可以不同,鉴于越高层需要的全局信息越多,可以在层级较高的时候把w调大。

 

因为w远小于长度,所以有了滑动窗口,内存占用就从l^2变成了l乘以w,也就是线性。

 

滑动窗口+空洞Attention

 

上面的滑动窗口很类似于卷积,那幺相应的,我们还可以像卷积一样加空洞,如下图c所示,这里有个参数d,意为空洞的大小。

 

 

空洞可以帮助attention在不增大内存占用的同时,增大视野域d倍。

 

全局Attention

 

在现在Bert架构中,只靠上面两个局部attention是不够的,因为储存的信息毕竟有限制,为了解决这个问题,所以全局attention应运而生。

 

这里的全局attention并不是所有位置attend所有位置,而是选中一些位置让它们之间去做两两的attention。而这些位置的选择,则与具体的问题相关。例如,对于文本分类问题而言,[CLS]这个特殊token会被当做所有信息的聚合点,因而这个位置肯定要被选中。而对于QA问题而言,所有的question的token上要去做全局attention。

 

旁白君: 其实这也是一种妥协,放弃了任务上的通用性。

 

QKV线性映射

 

回顾一下Attention的计算,对于一个序列的embedding,我们需要让它经过三个矩阵Q,K,V分别转化为q,k,v,然后q和k计算相似度,得到的权重再去和v做组合。

 

而有了全局attention和局部attention的区分后,我们也需要将这两种attention对应的QKV矩阵区分开,所以有两套QKV矩阵。

 

CUDA实现

 

在Tensorflow和pytorch的原始实现中,并没有用于计算滑动窗口attention的专门的实现,因为这个实现需要实现矩阵乘积且只要对角线的位置的非0值需要存在内存里。

 

而如果用for循环又异常的慢。所以用high-level的python struct描述了这种算法,并基于TVM生成了可以在GPU上编译的代码。在最上图可以看到,比naive实现快了6倍。

 

Attention的设置

 

正如上面提到的,在底层的时候使用较小的w,而在高层的时候w变大。这个参数设置需要做参数搜索,来平衡性能和效果。

 

另外,在底层不适用空洞attention,因为底层需要学习直接的局部信息。而在较高的层次会使用空洞attention,不过只限定在2个头上。

 

实验效果

 

在字符级的语言模型上,可以看到Longformer比之前的算法都要好。

 

 

而在消融实验中,可以看到递增w的策略和空洞attention都能带来提升。

 

 

为了节省训练的时间,用Roberta的参数来初始化Longformer,当然,因为结构的不同,需要做一些变动。得到的结果如下,可以看到,loss降了约8%~10%。

 

 

思考

 

勤思考,多提问是Engineer的良好品德。

 

 

用Roberta的checkpoint来初始化Longformer,需要做哪些变动?应该如何考虑?

 

 

参考

 

[1]. Beltagy, Iz, Matthew E. Peters, and Arman Cohan. “: The long-document transformer.” arXiv preprint arXiv:2004.05150 (2020).

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注