具备长短距离注意力模块的Lite Transformer

Transformer在自然语言处理(例如，机器翻译、问答)中已经变得无处不在;然而，它需要大量的计算才能实现高性能，这使得它不适合受硬件资源和电池严格限制的移动应用程序。在本文中，提出了一个高效的移动NLP架构，Lite Transformer，以方便在边缘设备上部署移动NLP应用程序。关键的原语是长短距离注意(LSRA)，其中一组头专门进行局部上下文建模(通过卷积)，而另一组头专门进行远距离关系建模(通过注意)。这种专门化在三个已建立的语言任务(机器翻译、抽象摘要和语言建模)上带来了相对于普通Transformer的一致改进。

1. 引言

Transformer 由于其高训练效率和捕获远程依赖关系的优越能力，被广泛应用于自然语言处理。在此基础上，现代最先进的模型，如BERT，能够从未标记的文本中学习强大的语言表示，甚至在具有挑战性的问答任务上超越人类的表现。

然而，良好的性能是以高计算成本为代价的。例如，一个变压器模型需要超过10G的multi - add才能翻译一个只有30个单词的句子。如此高的计算资源需求超出了智能手机和物联网等许多边缘设备的能力。因此，为边缘实时NLP应用设计高效、快速的Transformer体系结构具有重要意义。自动神经结构搜索是高精度模型设计的一种选择，但巨大的搜索成本(GPU小时和二氧化碳排放)引发了严重的环境问题。

这篇论文中，主要关注如何在移动设备上进行高效推理，并将模型的总计算量（Mult-Adds）限制在500M以下。减少Transformer计算量的直接方法是缩小嵌入层的尺寸，虽然这能够有效地降低模型的大小和计算量，但同时也削弱了模型捕捉长短距离关系的能力。为了解决这个问题，系统地研究了Transformer的计算分解，发现计算量（Mult-Adds）主要集中在前馈神经网络（FFN）部分。现有的瓶颈结构Transformer模块并不高效。为此，提出了一种新颖的长短范围注意力（LSRA）。LSRA通过扩大注意力层来平衡FFN的计算量，拉伸瓶颈以引入更多的依赖性捕捉能力，同时缩小嵌入尺寸，以在保持相同性能的前提下降低总计算量。与其使用一个模块处理“通用”信息，LSRA采用专门的头部来建模长距离和短距离的上下文关系。受Wu等人（2019b）启发，LSRA在并行分支中引入卷积来捕捉局部依赖性，从而让注意力分支专注于全局上下文的捕捉。通过堆叠这种原语，构建了适用于移动NLP应用的Lite Transformer模型。

1.1 主要贡献

1、本文系统地分析了现代神经网络中常用的计算瓶颈结构，发现以FLOPs为优值的瓶颈设计对于一维注意力并不是最优的。

2、提出了一个专门的多分支特征提取器，长短距离注意(LSRA)，作为Transformer的基本构建块，其中卷积有助于捕获局部上下文，而注意力集中在全局上下文。

3、基于LSRA构建Lite Transformer。在移动计算资源约束下(500M multi - add)，Lite Transformer在三个广泛使用的机器翻译数据集上展示了连贯的改进。通过对其他任务的额外实验，Lite Transformer对于多语言应用程序是高效的。

4、即使与基于AutoML搜索的Evolved Transformer相比，Lite Transformer在移动设置下的WMT En-De数据集上的BLEU得分也高出0.5，节省了20000倍的二氧化碳排放设计成本。

2. 相关工作

递归神经网络（RNN）在各种序列建模任务中占主导地位，由于RNN的时间依赖性，它不容易在整个序列上并行化。卷积是对局部上下文信息建模的理想原语;然而，它缺乏捕捉远距离关系的能力，而这在许多序列建模任务中是至关重要的。

2.1 Transformers

注意力能够通过两两相关来捕获全局上下文信息。

3. 瓶颈对一维注意力有效吗？

注意力机制已经广泛应用于各种领域，包括一维（如语言处理）、二维（如图像识别）以及三维（如视频识别）任务。它通过计算所有输入元素之间的成对点积来建模短期和长期关系。尽管这种方法非常有效，但它也引入了大量的计算开销。

假设输入到注意力层的元素数量（例如，语言处理中词元的长度，图像中的像素数量等）为 NNN，特征的维度（即通道数）为 d，那么计算点积所需的计算量为 $N^2d$ 。对于图像和视频来说，N通常非常大。例如，在一个视频网络中（Wang et al, 2018），中间特征图包含16帧，每帧的分辨率为112×112，这就导致 $N=2\times 10^5$ 。卷积层和全连接层的计算量相对于 N 呈线性增长，而注意力层的计算量相对于 N 呈二次方增长。因此，当 N 很大时，注意力模块的计算量将迅速膨胀，导致计算负担过重。

为了解决这个难题，一种常见的做法是在应用注意力机制之前，先通过线性投影层减少通道数量 d，然后再增加维度（如图2所示）。在原始的Transformer设计中，注意力模块中的通道维度比FFN层小4倍。类似地，在非局部视频网络中，在应用非局部注意力模块之前，通道数量先减少一半。这种做法可以将计算量减少16倍或4倍。然而，这也削弱了注意力层在较小特征维度下捕捉上下文信息的能力。对于语言处理任务来说，情况可能更糟糕，因为注意力是捕捉上下文的主要模块（不同于图像和视频任务，卷积是主要的信息捕捉手段）。

对于翻译等任务来说，输入序列的长度 N 通常较短，一般在20到30之间。一个Transformer块由一个注意力模块（或者在解码器中有两个）和一个前馈神经网络（FFN）组成。对于注意力层，计算量为 $O(4Nd^2+N^2d)$ ；对于FFN层，计算量为 $O(2\times 4Nd^2)$ 。在 N 较小时，瓶颈设计是否在计算量和准确性之间取得了良好的平衡是值得怀疑的。为了验证这一观点，首先在图2中分析了Transformer的计算分解。令人惊讶的是，对于原始的Transformer（在图中标为“Base”），FFN层实际上消耗了大部分计算资源。这并不理想，因为FFN本身并不能执行任何上下文捕捉。总的来说，由于 N 较小，瓶颈设计无法显著减少一维注意力中的计算量，而计算量减少的有限收益又被FFN层的大计算量进一步抵消。此外，由于注意力层的维度较小，这也削弱了其作为Transformer中主要上下文捕捉单元的能力。

因此，认为瓶颈设计对于一维注意力并不是最优的。设计了一种“扁平化”的Transformer块版本，这种设计不再减少和增加通道维度。在新的设计中，注意力部分现在占据了扁平化Transformer模型中的主要计算量，为进一步优化留下了更大的空间。

4. 长短范围注意力(LSRA)

通过一个卷积分支来专注于局部特征提取，该分支有效地对局部性进行建模，这样注意力分支就可以专注于全局特征提取。

Kovaleva等人（2019）和Clark等人（2020）对BERT模型（双向Transformer）中不同层的注意力权重进行了可视化。如图b所示，权重 w 说明了源句和目标句中单词之间的关系（自注意力机制的情况相同）。

当权重 $w_{ij}$ 较大时（颜色较深），表示源句中的第 i 个单词对目标句中的第 j 个单词更加关注。注意力图通常具有强烈的模式：稀疏和对角线分布。它们分别表示某些特定单词之间的关系：稀疏代表长期信息，而对角线代表小邻域内的相关性。将前者称为“全局”关系，后者称为“局部”关系。

在翻译任务中，注意力模块必须同时捕捉全局和局部上下文，这需要较大的模型容量。相比之下，通用的设计并不是最优的。以硬件设计为例，通用硬件如CPU的效率不如专用硬件。在这里，应该专门化全局和局部上下文的捕捉。当模型容量较大时，冗余是可以容忍的，甚至可能带来更好的性能。然而，在移动应用中，由于计算和功耗的限制，模型需要更加高效，因此专门化的上下文捕捉更为重要。为了解决这一问题，提出了一种更为专门化的架构——长短范围注意力（LSRA），分别捕捉全局和局部上下文。

如图a所示，LSRA模块采用了双分支设计。左侧分支捕捉全局上下文，而右侧分支则建模局部上下文。没有将整个输入送入两个分支，而是沿通道维度将其分为两部分，并在后续的前馈神经网络（FFN）层中混合。这样可以将总体计算量减少一半。左侧分支是一个标准的注意力模块，但通道维度减少了一半。对于右侧局部分支，一种自然的想法是对序列应用卷积操作。通过滑动窗口，模块可以轻松覆盖对角线分布。为了进一步减少计算量，用轻量版卷积替代了普通卷积，由线性层和深度卷积组成。通过将注意力模块和卷积模块并排放置，鼓励它们从全局和局部两个不同的角度理解句子，从而使架构受益于专门化设计并提高效率。

为了更好地理解这一点，在图3中可视化了完全训练的基础Transformer和Lite Transformer的同一层的平均注意力权重。可以很容易地看出，与基础Transformer试图同时建模全局和局部上下文不同，LSRA中的注意力模块只关注全局上下文的捕捉（没有对角线模式），而将局部上下文的捕捉留给了卷积分支。

5. 实验设置

5.3 架构

在本架构中，首先将基础Transformer模型中的瓶颈设计扁平化，然后用长短范围注意力（LSRA）替代自注意力机制。具体来说，使用两个专门化的模块：一个是注意力分支，另一个是卷积分支。卷积分支的输入和输出都通过全连接层进行转换，并且在WMT模型中，卷积核的大小是通过全连接层从输入动态计算得到的。对于编码器和解码器，卷积核的大小设置为[3, 5, 7, 31×3]，每个模块的头数为4（基础Transformer模型中头数的一半）。在总结生成任务中，模型配置与WMT模型相同。对于语言建模任务，卷积分支的卷积核大小设置为[15, 15, 31×4, 63×6]。