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？

 

参考文献

[1]. Wu, Z., Liu, Z., Lin, J., Lin, Y., & Han, S. (2020). Lite transformer with long-short range attention. arXiv preprint arXiv:2004.11886.

 

声明：文章收集于网络，版权归原作者所有，为传播信息而发，如有侵权，请联系小编删除，谢谢！