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相比：

进一步压缩：