LSRA: 轻量级Transformer,注意力长短搭配

最近要开始使用Transformer去做一些事情了,特地把与此相关的知识点记录下来,构建相关的、完整的知识结构体系。

以下是要写的文章,本文是这个系列的第二十四篇,内容较为深入,需要学习基础的同学点击链接进入文章列表查看基础知识相关文章。

Overall

Transformer虽然效果好,但是却对资源要求比较高,很难在端设备上运行。参考文献[1]提出了一种长短attention结合(Long-Short Range Attention, LSRA)的方式,有效的将一些计算节省掉,加速了模型的运行,使得模型可以在端设备上快速运行。

Transformer的计算瓶颈

在传统的Transformer中,每个block中都有Multi-head Attention和全连接层,其中,随着序列长度N的增大,全连接层的计算量是线性增长,而attention的计算量则是平方增长(如果不使用之前提的Performer类似算法的话)。类比的,卷积操作也和序列长度是线性关系。

所以,当序列长度比较长的时候,attention占绝大多数的计算量。所以之前的方法是将通道数变少,然后做注意力计算,然后再放大。如下图的最上一部分所表示的那样。

如果维度是d,序列长度是N,那么注意力层的计算量(Multi-Add)是O(4Nd2+N2d),而FFN的计算量(Multi-Add)是O(2x4Nd2)。

而对于翻译任务来说,一般长度在20-30左右,比d要小很多,此时理论上来说FFN占用的计算量比较多,而不是attention。如下图所示,柱状图和上面的三个算法一一对应,对于Base算法,做了通道先减后增之后,FFN占的计算量占了大部分,attention计算反而是少数。但FFN这块并不能捕捉上下文信息,更多的是非线性变化。所以,设计了第二种方法,即,把attention计算的通道数的先减后增去掉,得到的计算量中Attention就占了大多数了。

长短搭配

在翻译任务中,如果使用传统的Transformer的话,会发现attention权重的分布如下,下图中颜色越深,权重越大。

从上图的权重分布中,我们可以看到,attention的很大一部分都是在做局部的attention,每个词主要是需要attend到周边的两个词语,其他长依赖比较少。

这就激发了一种想法,那就是把主要的局部attention剥离出来,由卷积完成,Multi-head attention则主要负责长依赖。如下图:

做了这个改进之后,有两点好处:

  • 卷积比attention要快,计算量提升2.5倍。
  • 多头注意力专注于长依赖,得到的注意力更加均匀。如下图:

注意,这里使用的卷积并不是普通的卷积,而是depth-wise的卷积加上线性层。具体我们会在后面的文章里讲,这里先挖个坑。

实验

在翻译任务上的结果:

和Evolved Transformer相比:

进一步压缩:

总结

论文的idea比较直观,主要是长依赖和短依赖的剥离,并引入卷积来捕捉短依赖,总体思想和Transformer之自适应宽度注意力有点类似。这篇文章中发现低层次上的注意力都比较短,层次越高,注意力的所关注的依赖越长。

所以感觉这里可以更进一步,在低层次上甚至可以不用Multi-Head Attention?

(0)

相关推荐