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

摘要: Transformer虽然效果好，但是却对资源要求比较高，很难在端设备上运行。在传统的Transformer中，每个block中都有Multi-head Attention和全连接层，其中，随着序列长度N的增大，全连接层的计算量是线性增长，而attent ...

人工智能学习离不开实践的验证，推荐大家可以多在FlyAI-AI竞赛服务平台多参加训练和竞赛，以此来提升自己的能力。FlyAI是为AI开发者提供数据竞赛并支持GPU离线训练的一站式服务平台。每周免费提供项目开源算法样例，支持算法能力变现以及快速的迭代算法模型。

最近要开始使用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.

 

更多精彩内容请访问FlyAI-AI竞赛服务平台；为AI开发者提供数据竞赛并支持GPU离线训练的一站式服务平台；每周免费提供项目开源算法样例，支持算法能力变现以及快速的迭代算法模型。

挑战者，都在FlyAI！！！