HAT: Hardware-Aware Transformers for Efficient Natural Language Processing翻译

摘要

Transformer在自然语言处理（NLP）任务中无处不在，但是由于需要大量的计算，因此很难将其部署在硬件上。为了在资源受限的硬件平台上实现低延迟推理，我们建议使用神经结构搜索设计Hardware-Aware Transformer（HAT）。我们首先构造一个具有任意编码器-解码器注意力模块和异构层的大型设计空间。然后，我们训练一个涵盖设计空间中所有候选对象的SuperTransformer，并有效地产生许多具有权重共享的SubTransformer。最后，我们在硬件延迟约束下执行进化搜索，以找到n能够在目标硬件上快速运行的专用SubTransformer。在四个机器翻译任务上的大量实验表明，HAT可以发现不同硬件（CPU，GPU，IoT设备）的有效模型。在Raspberry Pi-4上运行WMT’14翻译任务时，HAT可以实现3倍的加速，是基线Transformer的3.7倍；与Evolved Transformer相比，速度提高了2.7倍，尺寸缩小了3.6倍，搜索成本降低了12,041倍，并且没有性能损失。HAT已在https://github.com/mit-han-lab/hardware-aware-transformers.git上开源。

1.介绍

Transformer已被广泛用于自然语言处理任务中。通过将多个相同的编码器/解码器层与注意力模块堆叠在一起，与以前的卷积或循环神经网络模型相比，它可以显着提高性能。
然而，由于高昂的计算成本，在移动设备上部署Transformer具有挑战性。例如，为了翻译只有30个单词的句子，Transformer-Big模型需要执行13G FLOP，并且在Raspberry Pi上花费20秒。如此长的延迟会损害边缘设备上的用户体验。 因此，我们需要硬件效率高的Transformer（图1）。
评估Transformer的效率时有两个常见的陷阱。 （1）FLOP不能反映测量的延迟。尽管在现有技术中FLOP被用作效率的度量标准，但这并不是一个好的延迟代理。如图2（右）所示，具有相同FLOP的模型可能会导致测量延迟有很大差异。（2）不同的硬件偏好不同的Transformer体系结构。如表1所示，在一种硬件上优化的Transformer模型对于另一种硬件而言是次优的，因为延迟受不同硬件平台上不同因素的影响。 例如，嵌入维度对Raspberry Pi延迟有重大影响，但几乎不影响GPU延迟（图2）。


受神经结构搜索（NAS）成功的启发，我们提出去搜索Hardware-Aware Transformers （HAT），直接将延迟反馈纳入设计循环。这样，我们就不需要FLOP作为延迟代理，并且可以搜索各种硬件的专用模型。
我们首先构造一个具有任意编码器-解码器注意力模块和异构Transformer层的大型搜索空间。传统Transformers在编码器和解码器之间存在信息瓶颈。任意编码器-解码器的注意力模块打破了瓶颈，允许所有解码器层参与多个不同的编码器层，而不仅仅是最后一层。因此，来自编码器的低级信息也可以被解码器使用。受图2的启发，我们引入了异构的Transformer层，以允许不同的层具有适应各种硬件的不同体系结构。
为了在如此大的设计空间中执行低成本搜索，我们首先训练一个Transformer超级网络-SuperTransformer，其中包含许多共享权重的SubTransformer。 我们通过优化SuperTransformer中统一采样的SubTransformer来同时训练所有SubTransformer。具有从SuperTransformer继承的权重的SubTransformer的性能可以为从头开始训练的不同体系结构提供良好的相对性能近似值。与传统的NAS不同，我们只需要一次性支付SuperTransformer的训练成本，就可以使用它评估设计空间中的所有模型。最后，我们进行进化搜索以在硬件延迟约束下找到最佳的SubTransformer。实验表明，HAT可以自然地与模型压缩技术（例如量化和知识蒸馏）结合使用。
我们通过在Raspberry Pi ARM CPU，Intel Xeon CPU和Nvidia TITAN Xp GPU上的WMT’14 En-De，WMT’14 En-Fr，WMT’19 En-De和IWSLT’14 De-En任务评估HAT。与以前的工作相比，在不损失准确率的条件下，HAT的速度提高了3倍，尺寸缩小了3.7倍。HAT的搜索成本降低了12041倍，以2.7倍的加速和3.6倍的较小尺寸胜过了Evolved Transformer。与Levenshtein和Lite Transformers相比，它还实现了1.9倍的加速，而BLEU得分没有损失。通过4位量化，HAT可以进一步缩小模型尺寸25倍。
HAT具有三方面的贡献：
（1）Hardware-Aware和专有化。据我们所知，我们是第一个直接将硬件反馈纳入模型设计中，以减少目标硬件的NLP模型延迟时间，而不是依赖于代理信号（FLOP）。对于不同的硬件平台，探索了用于低延迟推理的专有模型。
（2）具有大设计空间的低成本神经结构搜索。我们提出了任意编码器-解码器的注意力模块，以打破信息瓶颈；和异构层，以使不同的层更改其容量。经过训练的共享权重的SuperTransformer，可以低成本搜索有效的模型。
（3）设计见解。基于搜索结果，我们揭示了一些设计见解：引入多个编码器层对解码器有利；GPU倾向于浅和宽模型，而ARM CPU则倾向于深和薄模型。

2.提出的方法

HAT框架的概述如图3所示。我们首先训练一个具有较大设计空间的SuperTransformer。然后，对于给定的硬件平台，我们收集了不同模型（SubTransformer体系结构，测量的延迟）对的数据集，并训练了延迟预测器。最后，我们进行带有延迟约束的进化搜索，以找到专门针对目标硬件的有效模型。

2.1 设计空间

我们通过打破Transformer设计中的两个约定来构建较大的设计空间：（1）所有解码器层仅使用最后一个编码器层； （2）所有层的超参数均相同。
（1）任意编码器-解码器注意力模块

不同的编码器层在不同的抽象级别上提取特征。按照惯例，所有解码器层仅使用最后一个编码器层。它形成一个信息瓶颈，迫使所有解码器层仅从高抽象级别学习，而忽略低级别信息。为了克服瓶颈，我们建议使用“任意编码器-解码器注意模块”来学习编码器和解码器之间最合适的连接。每个解码器层可以选择多个编码器层。来自编码器层的 $key$和 $value$向量在句子长度维度上串联起来（图4），并馈送到编码器解码器交叉注意力模块。该机制之所以有效是因为它不引入其他参数。延迟开销也可以忽略不计。例如，在每个解码器层都使用两个编码器层的情况下，基于Nvidia TITAN Xp GPU的Transformer-Base的延迟几乎不会增加0.4％。通过允许关注不同的抽象级别来提高模型容量。
（2）异构Transformer层
以前的Transformer对所有层都重复了一种架构。相反，在HAT中，不同的层是异构的，具有不同数量的头，隐藏层维度和嵌入层维度。在注意力层中，使用不同的头来捕获各种依赖性。但是，Voita et al. (2019)显示具有许多头是多余的。因此，我们使注意力头数具有弹性，以便每个注意力模块可以确定其必要的头数。
在FFN层中，将输入特征投影为更高的维度（隐藏层维度），然后带入激活层。传统上，将隐藏层维度设置为嵌入层的2倍或4倍，但这是次优的，因为不同的层根据特征提取的难度需要不同的容量。因此，我们使隐藏层维度具有弹性。
此外，我们还支持编码器和解码器的弹性嵌入维度，但在编码器/解码器内部是一致的。编码器和解码器层的数量也很灵活，可以学习适当级别的特征编码和解码。注意力模块中其他设计选择（例如Q，K，V向量的长度）可以自然地合并到我们的框架中，我们将留作以后的工作。

2.2 SuperTransformer

为了找到高性能的模型，拥有较大的设计空间至关重要。但是，训练所有模型并比较其BLEU分数是不可行的。因此，我们提出了SuperTransformer，这是一种用于性能近似的超级网络，它可以在不完全训练模型的情况下判断模型的性能。SuperTransformer是搜索空间中具有权重共享的最大模型。搜索空间中的每个模型（一个SubTransformer）都是SuperTransformer的一部分。所有SubTransformer均承担其共有部分的权重。对于弹性嵌入维度，所有SubTransformer共享最长的词嵌入的前部和相应的FC层权重。如图5所示，对于隐藏的弹性FFN隐藏层维度，FC权重的前部是共享的。对于注意力模块中的弹性头数，整个Q，K，V向量（长度在我们的设计空间中是固定的）通过划分头数部分来共享。弹性层号使所有SubTransformer共享前几层。
在SuperTransformer训练期间，将对所有可能的SubTransformer进行统一采样，并更新相应的权重。实际上，SuperTransformer只需要按照与基线Transformer模型相同的步骤进行训练，这是快速且低成本的。训练后，我们即可通过在验证集上评估相应的SubTransformer来获得设计空间中采样模型的性能代理。

2.3 对subTransformer执行进化搜索

给定延迟要求，我们执行进化搜索以找到满意的SubTransformer。有两种方法可以评估SubTransformer的硬件延迟：（1）在线测量，其中我们在搜索过程中测量模型。（2）离线，我们在其中训练延迟预测器以提供延迟。我们在此应用离线方法，因为它快速准确。对于在线方法，单个采样的SubTransformer需要数百次推理才能获得准确的延迟，该延迟持续数分钟并减慢搜索速度。对于离线方法，我们将SubTransformer的体系结构编码为特征向量，并使用多层感知器（MLP）即时预测其延迟。经过数千的真实延迟时间数据的训练，预测器可以产生很高的准确性（图6）。请注意，预测的延迟仅在搜索过程中使用，我们在实验部分报告了实际测得的延迟。与为每个硬件推导封闭形式的延迟模型相比，延迟预测器方法更通用，更快捷。
我们使用进化算法来进行搜索过程。如图3所示，搜索引擎在延迟预测器中查询SubTransformer延迟，并在验证集上验证损失。该引擎仅向设计空间添加延迟小于硬件约束的SubTransformers。然后，我们从头开始训练搜索的模型以获得最终性能。