1.摘要

在这篇论文中，我们提出了一种高效的移动NLP架构——Lite Transformer，以在边缘设备上部署移动NLP应用。Transformer已经成为自然语言处理（例如机器翻译、问答系统）中无处不在的技术，但要实现高性能需要大量计算资源，这使得它在硬件资源和电池方面受限的移动应用中并不适用。Lite Transformer的关键基元是长短范围注意力（Long-Short Range Attention，LSRA），其中一组头部专注于局部上下文建模（通过卷积），而另一组头部专注于远距离关系建模（通过注意力机制）。这种专业化在机器翻译、抽象摘要和语言建模这三个广为人知的语言任务上都带来了持续的改进。在受限的资源条件下（500M/100M MACs），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的条件下比Transformer的困惑度低1.8。值得注意的是，在不需要耗费超过250个GPU年的昂贵架构搜索的情况下，Lite Transformer在移动NLP设置中比基于AutoML的Evolved Transformer高出0.5个BLEU分数。

2.引言

Transformer（Vaswani等，2017）由于其高效的训练能力和在捕捉远距离依赖关系方面的卓越能力而被广泛应用于自然语言处理。在此基础上，现代最先进的模型，如BERT（Devlin等，2019），能够从无标签文本中学习到强大的语言表示，并在具有挑战性的问答任务上甚至超过人类表现。

然而，这种出色的性能付出了巨大的计算代价。例如，一个单独的Transformer模型在翻译仅30个词的句子时需要超过10G的乘加操作。这种极高的计算资源需求超出了许多边缘设备（如智能手机和物联网设备）的能力。因此，为边缘上的实时NLP应用设计高效快速的Transformer架构具有重要意义。自动神经架构搜索（Zoph＆Le，2017; So等，2019）是高准确性模型设计的一种选择，但巨大的搜索成本（GPU小时和二氧化碳排放）引起了严重的环境担忧。

在本文中，我们专注于移动设备的高效推断，其中乘加操作的总数限制在500M以下。减少Transformer计算量的一种直接方法是缩小嵌入大小。尽管这可以有效减小模型大小和计算量，但同时也削弱了模型捕捉长距离和短距离关系的能力。为此，我们系统地研究了Transformer计算的分解情况。

本文观察到，在Transformer中，计算（乘加操作）主要由前馈网络（FFN）占据主导地位。我们发现目前主流的瓶颈结构化Transformer块并不高效。因此，我们提出了一种新颖的长短范围注意力（Long-Short Range Attention，LSRA）基元。LSRA在FFN中的计算与更宽的注意力层之间进行权衡。它扩展了瓶颈，增加了注意力层对依赖关系的捕捉能力，然后通过缩小嵌入大小来减少总的计算量，同时保持相同的性能。LSRA不再将一个模块用于“通用”信息，而是专门将头部分配给长距离和短距离上下文建模。受Wu等人（2019b）的启发，LSRA在并行分支中引入了卷积，以捕捉局部依赖关系，使注意力分支能够专注于全局上下文捕捉。通过堆叠这个基元，我们构建了适用于移动NLP应用的Lite Transformer。

大量实验证明，我们的Lite Transformer模型在机器翻译、抽象摘要和语言建模三个语言任务上相比Transformer表现出显著改进。在IWSLT 2014德英机器翻译任务上，在100M乘加操作下，比Transformer高出3.1 BLEU；在WMT 2014英德机器翻译任务上，在500M乘加操作下超过Transformer 0.4 BLEU，在100M乘加操作下超过Transformer 1.2 BLEU；在WMT 2014英法机器翻译任务上，与Transformer相比也取得了持续的改进：在500M乘加操作下提高了1.2 BLEU，在100M乘加操作下提高了1.7 BLEU。此外，结合通用的模型压缩技术（修剪和量化），我们的Lite Transformer可以实现18.2倍的模型大小压缩。在摘要任务中，在CNN-DailyMail数据集上，它将Transformer基础模型的计算量减少了2.4倍。在语言建模方面，在约500M乘加操作的条件下，其困惑度比Transformer低1.8。

基于我们的设计见解，我们手动设计的Lite Transformer在移动NLP设置下比基于AutoML的Evolved Transformer（So等，2019）高出0.5 BLEU，而Evolved Transformer需要超过250个GPU年的搜索，并在它们的寿命中产生了相当于五辆汽车的碳排放量（参见图1b）。这表明AutoML并非万能解：仔细的分析和设计见解（例如

去除瓶颈、专门化头部）可以有效减小搜索空间并提高样本效率。

本文的贡献有四个方面：

1. 我们对现代神经网络中常用的计算瓶颈结构进行了系统分析，并发现瓶颈设计在使用FLOPs作为评估指标时并不优化1-D注意力。
2. 我们提出了一种专门的多分支特征提取器，即长短范围注意力（LSRA），作为我们的Transformer的基本构建块，其中卷积有助于捕捉局部上下文，而注意力则集中在全局上下文上。
3. 我们基于LSRA构建了Lite Transformer。在移动计算资源受限（500M乘加操作）的情况下，我们的Lite Transformer在三个广泛使用的机器翻译数据集上展现了一致的改进。通过对其他任务的额外实验，Lite Transformer在多语言应用中也表现出高效性。
4. 与基于AutoML搜索的Evolved Transformer相比，我们的Lite Transformer在移动设置下在WMT En-De数据集上提供了0.5个更高的BLEU分数，节省了20000倍的设计成本（在CO2排放方面）。这提醒我们重新思考AutoML在设计成本和“绿色人工智能”方面的实用性。

3 IS BOTTLENECK EFFECTIVE FOR 1-D ATTENTION?

注意力机制已广泛应用于各种应用领域，包括1-D（语言处理（Vaswani等，2017））、2-D（图像识别）和3-D（视频识别（Wang等，2018））。它通过计算输入元素之间的成对点积来建模短期和长期关系。尽管其有效性，这种操作引入了大量计算。假设输入到注意力层的元素数量（例如语言处理中的令牌长度，图像中的像素数等）为N，特征的维度（即通道数）为d，点积所需的计算量为N^2d。对于图像和视频来说，N通常非常大。例如，视频网络（Wang等，2018）中的中间特征图具有16帧，每帧分辨率为112×112，导致N = 2 × 105。卷积和全连接层的计算量与N呈线性增长，而注意力层的计算量与N呈二次增长。随着N的增大，注意力模块的计算量很快就会变得巨大。

为了解决这个困境，一种常见的做法是在应用注意力之前使用线性投影层来减少通道数d，然后在之后增加维度（如图2所示）。在最初的Transformer设计中（Vaswani等，2017），注意力模块中的通道维度比FFN层小4倍。类似地，在非局部视频网络（Wang等，2018）中，在应用非局部注意力模块之前，通道数先减少一半。这种做法可以节省16倍或4倍的计算量。然而，这也会降低注意力层的上下文捕获能力，因为特征维度较小。对于语言处理来说，情况可能会更糟，因为注意力是上下文捕获的主要模块（不像图像和视频中的卷积层主要进行信息捕获）。

对于翻译等任务，输入序列的长度N倾向于较小，通常在20-30左右。一个Transformer块由一个注意力层（解码器有两个）和一个前馈网络（FFN）组成。对于注意力层，Mult-Adds的数量为O(4Nd^2 + N^2d)；对于FFN，Mult-Adds的数量为O(2 × 4Nd^2)。给定较小的N，可以怀疑瓶颈设计是否在1D注意力的计算和准确性之间取得了良好的权衡。为了验证这个想法，我们首先对Transformer中的计算进行了分析。令人惊讶的是，对于原始的Transformer（在图中标记为“Base”），FFN层实际上消耗了大部分计算资源。这是不可取的，因为FFN本身不能进行任何上下文捕获。总之，由于N较小，瓶颈设计无法显著减少1D注意力的计算量，而较大的FFN层进一步削弱了计算减少的有限好处。由于较小的维度，这也损害了注意力层的容量，而注意力层是Transformer中主要的上下文捕获单元。

因此，我们认为瓶颈设计对于1-D注意力不是最优的。我们相反设计了一个“扁平化”版本的Transformer块，不减少也不增加通道维度。通过这种新设计，在图2中的扁平化Transformer模型中，注意力部分占据了主要的计算量，为进一步优化留下了更大的空间。我们还测试了这种修改对WMT’14 En-Fr数据集的性能变化。我们可以在稍大的计算量下获得可比

4 LONG-SHORT RANGE ATTENTION (LSRA)

研究人员一直试图理解注意力所捕获的上下文。Kovaleva等人（2019）和Clark等人（2020）可视化了BERT中不同层的注意力权重。如图3b所示，权重w描述了源句子和目标句子之间的单词关系（自注意力也是如此）。具有较大的权重wij（较深的颜色），源句子中的第i个单词对目标句子中的第j个单词更加关注。注意力图通常具有明显的模式：稀疏和对角线。它们代表了一些特定单词之间的关系：稀疏表示长期信息的关系，对角线表示小邻域内的相关性。我们将前者称为“全局”关系，后者称为“局部”关系。

针对翻译任务，注意力模块必须同时捕捉全局和局部上下文，需要较大的容量。与专门设计相比，这不是最优的选择。以硬件设计为例，通用硬件如CPU不如专用硬件如FPGA高效。在这里，我们应该专门处理全局和局部上下文的捕捉。当模型容量相对较大时，冗余可以被容忍，并且可能提供更好的性能。然而，在移动应用中，由于计算和功耗的限制，模型应该更加高效。因此，专门的上下文捕捉更为苛刻。为了解决这个问题，我们提出了一种更专门的架构，即长短范围注意力（LSRA），它分别捕捉全局和局部上下文。

如图3a所示，我们的LSRA模块采用了两个分支的设计。左分支捕捉全局上下文，而右分支模拟局部上下文。我们将输入沿着通道维度分为两个部分，这些部分将在后续的FFN层中混合。这种做法将总体计算量减少了2倍。左分支是一个正常的注意力模块，就像Vaswani等人（2017）中的那样，只是通道维度减少了一半。对于局部关系的右分支，一个自然的想法是在序列上应用卷积。通过滑动窗口，对角线组可以很容易地被该模块覆盖。为了进一步减少计算量，我们将常规卷积替换为一个更轻量级的版本（Wu等人，2019b），由线性层和深度卷积组成。通过这种方式，我们将注意力和卷积模块并排放置，鼓励它们对句子具有不同的视角，从而使架构可以从专业化中受益，并实现更好的效率。

为了更好地理解，我们在图3中可视化了完全训练的基本transformer和我们的Lite Transformer相同层的平均注意力权重。可以轻松区分出，与试图同时建模全局和局部上下文的基本transformer不同，LSRA中的注意力模块仅关注全局上下文的捕捉（没有对角线模式），将局部上下文的捕捉交给卷积分支。