Lite Transformer with Long-Short range attention

Lite Transformer with Long-Short range attention
2020年 arxiv文章

Abstract

Transformer在自然语言处理（例如机器翻译、问答）中变得无所不在；然而，它需要大量的计算来实现高性能，这使得它不适合受到硬件资源和电池严格限制的移动应用。本文提出了一种高效的移动NLP体系结构——Lite Transformer，以方便在边缘设备上部署移动NLP应用。关键的原语是长短距离注意 (LSRA)，其中一组头部专用于局部上下文建模（通过卷积），而另一组专用于长距离关系建模（通过注意）。这种专门化在三个成熟的语言任务（机器翻译、抽象摘要和语言建模）上给朴素transformer带来了一致的改进。在资源受限（500M/100M MAC）下，Lite Transformer 在 WMT’14 英语-法语上的性能分别比 Transformer 高 1.2/1.7 BLEU。Lite transformer使transformer基础模型的计算量减少了2.5 ×，BLEU评分降低了0.3。结合修剪和量化，进一步将Lite Transformer的模型尺寸压缩18.2倍。对于语言建模，Lite Transformer 在大约 500M MACs 时实现的 perplexity 比 Transformer 低 1.8。值得注意的是，用于移动NLP设置时，Lite transformer的性能比基于AutoML的演进transformer高出0.5个BLEU，而不需要花费超过250个GPU年的昂贵的体系结构搜索。代码已在https://github.com/mit-han-lab/lite-transformer提供。

Introduction

Transformer (Vaswani et al., 2017) 因其高训练效率和卓越的远距离依赖捕获能力而被广泛用于自然语言处理。 在它们之上，现代最先进的模型，例如 BERT（Devlin 等人，2019），能够从未标记的文本中学习强大的语言表示，甚至在具有挑战性的问答任务中超越人类的表现 。

然而，良好的性能是以高昂的计算成本为代价的。例如，单个transformer模型需要超过10G的Mult-Adds才能翻译一个只有30个单词的句子。如此极高的计算资源要求超出了智能手机和物联网等许多边缘设备的能力。因此，设计专门用于边缘实时NLP应用的高效快速transformer架构至关重要。自动网络结构搜索（Zoph&Le，2017；So等人，2019）是高精度模型设计的选择，但巨大的搜索成本（GPU小时和二氧化碳排放）引发了严重的环境问题（Strubell等人，2019），如图1b所示。

Fig.1 左：最近 NLP 模型的规模迅速增长，并在很大程度上超出了移动限制。 右图：基于 AutoML 的 NLP 模型的搜索成本令人望而却步，其二氧化碳排放量几乎是汽车平均寿命排放量的 5 倍。

在本文中，我们专注于移动设备的有效推理，其中 MultAdd(MultAdd的意思是不是乘加) 的总数被限制在 500M 以下。减少transformer计算的一种直接方法是直接缩小嵌入大小。虽然它可以有效地减少模型大小和计算量，但同时也削弱了模型捕捉长短距离关系的能力。为此，我们系统地研究了 Transformer 的计算分解，并观察到计算（Mult-Adds）由前馈网络（FFN）主导。我们发现流行的瓶颈结构transformer块效率不高。然后，我们提出了一种新颖的长短程注意（LSRA）原语。 LSRA 权衡 FFN 中的计算以获得更广泛的注意力层。它扩展了瓶颈，为注意力层引入更多的依赖捕获能力，然后缩小嵌入大小以减少总计算量，同时保持相同的性能。 LSRA 没有一个用于“一般”信息的模块，而是专门用于对长距离和短距离上下文进行建模。受吴(2019b)等人的启发 ，LSRA 在并行分支中引入卷积以捕获局部依赖关系，以便注意力分支可以专注于全局上下文捕获。通过堆叠这个原语，我们为移动 NLP 应用程序构建了 Lite Transformer。

大量实验表明，我们的 Lite Transformer 模型在机器翻译、抽象摘要和语言建模三个语言任务上提供了显著改进。对于机器翻译，在 IWSLT 2014 German-English 上，它在 100M Mult-Adds 下比 Transformer 高 3.1 BLEU；在 WMT 2014 英语-德语上，它在 500M Mult-Adds下超过transformer 0.4 BLEU，在 100M Mult-Adds下超过 1.2 BLEU；在 WMT 2014 English-French 上，它还实现了对transformer的一致改进：在 500M Mult-Adds下 1.2 BLEU，在 100M Mult-Adds下 1.7 BLEU。此外，结合通用模型压缩技术 (修剪和量化)，我们的Lite transformer可以实现18.2 × 模型大小的压缩。对于摘要任务，在CNN-DailyMail上，它将transformer基础模型的计算量减少了2.4倍。对于语言建模，它的perplexity低于500M Mult-add左右的transformer 1.8。

在我们设计见解的指导下，我们手动设计的Lite Transformer比基于AutoML的进化型Transformer（So等人，2019年）的BLEU高0.5，这需要250 GPU以上的时间来搜索，在其生命周期中排放的碳量相当于五辆汽车（见图1b）。这表明AutoML并不是万能的：仔细的分析和设计见解（即消除瓶颈、专门化的头部）可以有效地削减搜索空间并提高样本效率。

本文的贡献有四个方面:
1、我们系统地分析了现代神经网络中常用的计算瓶颈结构，发现如果使用FLOPs作为评价指标，瓶颈设计对于一维注意不是最优的。
2、我们提出了一个专门的多分支特征提取器，Long-Short Range Attention (LSRA)，作为我们transformer的基本构建块，其中卷积有助于捕获局部上下文，注意力集中在全局上下文上。
3、我们基于我们的 LSRA 构建 Lite Transformer。在移动计算资源限制（500M Mult-Adds）下，我们的 Lite Transformer 展示了对三个广泛使用的机器翻译数据集的一致改进。通过对其他任务的额外实验，