Transformer 压缩方法 Lite Transformer（2048Lite）

Transformer 是目前 NLP 领域使用很广泛的模型，可用在翻译、文本摘要，问答系统等。Transformer 具有优秀的效果，但是需要大量的算力，使其不易部署在移动端。2020 年 4 月 MIT 和上海交大提出了 Lite Transformer，资源受限的情况下 Lite Transformer 在多个 NLP 任务上均取得比 Transformer 更好的结果。

1.前言

Transformer 是 Google 在 2017 年提出的模型，不熟悉的童鞋可以看下之前的文章。Transformer 采用了 Self-Attention 结构，在捕获长距离依赖关系上表现很好，但是其需要大量的计算，在翻译具有 30 个单词的句子时需要 10G MACs (MACs 指 Mult-Adds，用于统计计算量，pytorch 的模型可以用 torchstat 计算 MACs 值)。

大量的计算使 Transformer 不容易部署在移动端，因此 MIT 和上海交大在 2020 年 4 月提出了 Lite Transformer，《LITE TRANSFORMER WITH LONG-SHORT RANGE ATTENTION》。修改了 Transformer 的结构，使模型更适用于资源受限的情况下。

作者认为 Transformer 的 Self-Attention 是有冗余的，这在计算量充足是可以提供更好的性能，但是在移动端，更应该考虑的是模型的高效性。因此作者提出了 LSRA (Long Short Range Attention)，用于替代 Self-Attention。LSRA 将输入分为两个部分，一部分使用原来的自注意力机制提取全局信息，另一部分使用卷积核提取局部信息，LSRA 可以提高模型的计算效率，在更少的计算量得到较好的效果。下图是不同模型参数量的区别。

不同模型的参数量

论文中的实验结果显示，(500M/100M MACs) 的计算量下，Lite Transformer 在 WMT‘14 English-French 翻译任务上比 Transformer 的 BLEU 值高 (1.2/1.7)。如果 Lite Transformer 只使用了 Transformer 40% 的计算量，其 BLEU 值也仅比 Transformer 低 0.3。Lite Transformer 还可以结合一些剪枝和量化的方法，使模型压缩到 Transformer 的 5%。在 500 MACs 的计算量下进行语言建模任务，Lite Transformer 的困惑度比 Transformer 低 1.8。

2.Transformer 结构的问题

Transformer 的计算主要分为两个部分：Attention 和 FFN (Feedforward network)。作者在论文中提到，对于长度为 N 的序列，Attention 的 MACs 与 N 平方成正比，FFN 的 MACs 与 N 成正比。

Transformer 原始的减少计算量的方法是通过一种先降维再升维的方法。首先使用 FFN 降低输入的维度 d (在原 Transformer 论文中 FFN 中间层的维度是 2048，而 Attention 接收的输入维度是 512)，然后 Attention 的输出再经过一层 FFN 升维。Lite Transformer 作者把这一过程称为瓶颈方式 (bottleneck)，如下图 Base 模型所示。

作者认为 Transformer 这种 bottleneck 的方式有两个问题：

把降维后的特征传入 Attention，降低了模型捕捉上下文关系的能力，因为模型中主要是 Attention 负责上下文关系。
Transformer 中 Attention 的计算量为 O(4Ndd NNd)，FFN 的计算量为 O(2×4Ndd)。当 N 比较小的时候，这种 bottleneck 设计未必有利于计算量的减少，如下图的 Base 模型所示，FFN 用到的计算量比 Attention 的还多。

因此作者认为 bottleneck 结构不适合用于 Transformer 的优化，作者设计了一种扁平结构 (flattened)，即输入和输出维度是一致的，无需降维和升维。如上面图中的 flattened 模型，此时计算量主要集中在 Attention 部分，而 Attention 部分有很大的优化空间，作者提出了 LSRA 优化 Attention 部分。

3.LSRA

有研究发现 Transformer 的 Attention 包含比较多的冗余，当计算量足够时，这些冗余可以带来性能的提升，但是却不适合用于计算量不足的时候。

LSRA

因此作者提出了一种 LSRA (Long Short Range Attention) 结构，如下图所示。LSRA 包含两个分支，Self-Attention 用来捕获长距离信息，卷积用来捕获局部信息。两个部分得到的 Attention 矩阵如下图所示，可以看到卷积得到的 Attention 几乎是延着对角线的 (即关注局部内容)，而 self Attention 的矩阵相对稀疏 (关注长距离内容)。

卷积和Self-Attention的Attention矩阵

LSRA 左侧的 Self-Attention 和原来的一样，但是其输入维度少了一半。右边的卷积部分，为了更好的减少计算量，作者采用了《Pay Less Attention with Lightweight and Dynamic Convolutions》中的一种轻量级的卷积结构。

4.实验结果

对比 Lite Transformer 和 Transformer

下面是在资源受限的情况下，对比 Transformer 和 Lite Transformer，可以看到在计算量比较小的时候，Lite Transformer 可以取得比 Transformer 更好的效果。

下图是在 WMT En-Fr 翻译数据集上对比 Transformer 和 Lite Transformer，可以看到在计算量一致的情况下，Lite Transformer 的 BLEU 值均比 Transformer 高。Transformer 需要更多的计算量才可以达到和 Lite Transformer 相同的性能，例如图中要取得 39.7 的 BLEU 值，Transformer 需要 2.5 倍的计算量。

下图是在语言建模任务上的对比，类似地，在同样计算量下，Lite Transformer 的困惑度 PPL 要比 Transformer 低。

对比 Lite Transformer 和 Evolved Transformer

Evolved Transformer 是一种搜索 Transformer 结构的方法，可以在一个大的搜索空间中搜索出合适的 Transformer 结构。Lite Transformer 的作者同样对比了 Evolved Transformer 搜索得到的结构，结果如下所示。结果显示还是 Lite Transformer 效果更好，而且 Evolved Transformer 花费更大，也会排出大量的二氧化碳。

Lite Transformer 压缩

Lite Transformer 可以使用一些比较常见的压缩方法进行压缩，如剪枝和量化。作者用了《Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding》中的量化方法把网络的权重压缩到 8 bit；也用《 Learning both weights and connections for efficient neural network》中的剪枝策略。

最终可以把模型压缩到原始 Transformer 的 95%，而 BLEU 值只下降了 0.4，如下图所示。