[Transformer]HAT: Hardware-Aware Transformers for Efficient Natural Language Processing

HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
Paper
Code

Abstract

Transformer在NLP领域被广泛应用，但是由于其SA的密集计算使其很难在硬件上部署，为了实现在硬件平台上进行低延时推断，本文建议借助NAS设计硬件感知的Transformer（HAT,Hardware-Aware Transformer）。

本文首先构造了一个包含任意编解码注意力和异构层的搜索空间；其次本文训练了一个SuperTransformer包含了搜索空间中所有候选变压器，通过权重共享可以获得SubTransformer；最终本文使用进化搜索，搜索约束是硬件延迟从而找到适合特定硬件的目标Transformer结构。


本文在四类机器翻译任务上进行搜索，搜索到适合不同硬件平台（CPU,GPU,IoT）的特定Transforemr结构。比如在进行WMT‘14翻译任务时HAT加速比是3X，比急转Transformer规模小3.7x；与Evolved Transformer相比则是搜索成本减少了12041x，没有精度下降。

Section I Introduction

Transformer被广泛用于自然语言处理任务中，通过在encoder/decoder中堆叠多个相同的注意力模块，可以比CNN或RNN的性能获得进一步的提升。
但是由于Transformer高昂的计算成本使其难以部署到硬件，尤其是移动终端。比如为了翻译30个单词长的句子，Transformer-Big需要执行13G FLOPs，耗时20s。这种延时会极大影响用户体验，因此需要设计硬件感知的Transformer。



Fig1展示了HAT的搜索流程买寿险本文会训练一个SuperTransformer包含无数子网络；然后通过进化算法通过对硬件时延进行反馈搜索到适用于特定硬件的SubTransformer。
评价Transformer的效率有以下两方面的限制：


（1）FLOPs并不能反映Transformer的延迟。
尽管FLOPs用来作为评价Transformer效率的指标，但是相同FLOPs的模型其实际延迟可能不同，参见FIG2.


(2)不同的硬件平台适合不同的Transformer架构。


如Table 1所示，对于某硬件足有的模型可能对另一个硬件是次优的。比如嵌入维度对树莓派的延迟有显著影响，但几乎对GPU没有任何影响。


Fig 2展示了不同硬件上不同Transformer的延迟，可以看到hidden dim&lembedding dim对树莓派延迟有很大影响，但是对CPU影响次之，对GPU则影响不大。因此需要考虑硬件延迟的反馈来针对不同硬件平台设计不同的模型。




本文的HAT会将延迟反馈直接加入到HAT的设计循环中，这样就不用使用FLOPs作为代理，并且可以针对不同硬件搜索特定的模型。

HAT首先会设计一个大的搜索空间，里面包含编解码注意力和异构Transformer层。传统的Transformer在encoder-decoder之间会有bottleneck层。

Arbitrary encoder-decoder attention打破了平静不分，允许所有的decoder layer与encoder layer进行交互，而不是只关注encoder最后一层的输出，这样可以利用到low-level信息。

受Fig 2中实验现象的启发，本文引入了Transformer异构层，允许不同层具有适合不同硬件的不同结构。


为了在搜索空间中进行低成本的搜索，本文会首先训练一个Super Transformer，Super Transformer会包含所有的sub-Transformer，它们通过权重共享继承SuperTransformer的权重。随后会同时优化所有的SubTransformer、SubTransformer的性能可以作为从头训练的网络框架的近似值。

与传统NAS不同的地方在于，本文的计算成本只有SuperTransformer的训练成本这样就可以用它来评估搜索空间中的所有模型。

最后以硬件延迟作为约束，使用进化搜索寻找最优的SubTransformer。实验结果表明本文的HAT可以与量化、知识蒸馏等压缩技术相结合。
本文在WMT机器翻译任务上测试了HAT的效果，硬件平台分别是树莓派ARM CPU,Intel Xeon CPU，NvidiaXpGPU.y与prior work相比HAT加速了3倍，网络模型小了3.7x，搜索成本则比Evolved Transformer小了12041x.
HAT贡献总结如下：

（1）硬件友好。
据我们所知这是第一次将硬件反馈加入到模型设计中，针对不同硬件减少NLP模型的延迟，而不是FLOPs作为代理的评估指标。对于不同的硬件平台会探索不同的网络结构；

（2）在大规模搜索空间中以较低成本进行网络结构搜索。
本文提出arbitrary encoder-decoder attention来打破信息瓶颈，使用heterogeneous layer来调整每一层的规模。通过权重共享的方式来减少搜索成本；

（3）设计思路。
本文基于实验结果得到以下结论：使用多个encoder layer有助于帮助decoder处理结果，GPU更适合较浅较宽的网络模型而ARM CPU则更适合更深更细的网络。

Section II Proposed Approaches

Fig 3展示了HAT的整个流程。首先训练一个SuperTransformer，然后对于特定的硬件平台会收集（子网络结构，时延）这一指标对，然后训练一个延迟预测器。最后使用进化搜索，以延迟作为约束找到针对特定硬件平台的最优网络结构。

Part 1 Design Space

本文通过打破了传统Transformer中的界限设计了一个超大的搜索空间，打破了以下局限：



（1）decoder只关注encoder最后一层；



（2）所有层都是相同的




Arbitrary Encoder-Decoder Attention



不同的编码层可以提取不同尺度的特征，一般来说decoder只利用了encoder的最后一层，这就形成了information bottleneck，因为所有decoder只能从最后一层较高的抽象层次进行学习，忽略了low-level信息。本文打破了这一瓶颈，提出了任意编码-解码注意力，用于学习编码器和解码器之间最合适的连接方式。每一层decoder层可以选择多个encoder layer，来自不同层的key和value会级联后送入decoder进行处理。这一过程也十分搞笑因为没有使用额外的参数，延迟开销也可以忽略不计。
比如当每个decoder层选择两个encoder layer，在GPU上的延迟仅增加0.4%，几乎没有增加，这也证明了模型可以学习到不同抽象层次的注意力。



Hetergeneous Transformer Layers




之前Transformer中所有层都是一样的，在HAT中每一层则是不一样的，比如head数目不同，hidden dims不同，或者embedding dims不同；注意力层中不同的attention head用于捕获不同的依赖关系，但是也有研究证明许多attention head是高度冗余的，本文通过弹性设置每一个注意力模块的head数目来避免冗余。
在FFN层会将输入特征映射到更多维度，一般会升维到2或4倍但这并不是最优的，不同层需要不同的映射维度，这是取决于输入的特征。因此本文将隐藏层的维度也设置为弹性的。




此外decoder和encoder中的embedding dims也是弹性的，但是在encoder/decoder内部是保持一致的。当前Q,K,V的长度也可设置为弹性的但本文并未做尝试，将其作为未来工作的一部分。





Part 2 SuperTransformer

为了获得高性能的模型，一个较大的搜索空间是至关重要的，但是评估所有模型的性能，对比他们的BLEU分数是不可能的。因此本文提出使用超网-SuperTransformer来对性能进行近似，在不完全训练模型的前提下作为模型性能的近似。

SuperTransformer是搜索空间中通过权重共享能获得的最大模型，搜索空间的每一个模型都是SuperTransformer的一部分。比如在搜索embedding dims时所有的子模型都共享FC层的权重等；搜索head number时所有的Q,K,V是共享的。搜索某一层的layer number时则共享前面的层。

在SuperTransformer的训练中，会均匀采样SubTransformer，然后更新对应权重，实际应用时只需要训练SuperTransformer作为基线模型，这样训练快成本也低，然后在验证集上对对应的SubTransformer进行评估，从而获得最佳的性能代理模型。

Part 3 Evolutionary Search for SubTransformer

本文使用进化搜索以延迟作为约束来搜索得到最优的SubTransformer，有两种评估方法：

（1）在线测评
在搜索过程中对模型进行评估
;
（2）离线测评
训练一个latency predictor来预测试验
.
本文使用第二种，更为快速和准确。
在线测评的话需要对某个采样的SubTransformer进行数百次推理才能得到延迟信息，时间长达数分钟，大大减缓了搜索速度；使用离线测评会将SubTransformer的结构进行编码然后通过一个MLP立即预测延迟，经过数千个真实的延迟数据点的训练，预测器的准确性可以达到很高。
注意预测的延迟只在搜索过程中使用，最后使用的是真实测量的延迟信息。

本文使用进化搜索，参见Fig 3,搜索引擎会查询SubTransformer的延迟预测期的预测结果，然后在验证集上验证损失。搜索引擎只会向种群中添加延迟小于约束的SubTransformer；然后会从头训练搜索到的模型，从而获得最终的网络结构。

Section III Experiments

Datasets

本文在4个机器翻译数据集上进行了测试，分别是WMT’14 En-De, WMT’14 En-Fr,
WMT’19 En-De, IWSLT’14 De-En。

Experiment Setups


本文的baseline model包括Transformer,Levenshtein Transformer,Evolved Transformer,Lite Transformer.
翻译的精度使用的是BLEU分数；对于延迟则是测量的从源语言翻译到目标语言用的时间，本文会测试300次然后移除最快和最慢的10%，计算剩余推断时间的平均值作为延迟指标。



硬件平台：


树莓派 ARM Cortex-A72 CPU



Intel Xeon E5-2640 CPU



Nvidia Xp GPU




实施细节




SuperTransformer的搜索空间包括：





embedding dim - [512,640]






hidden dim-[1024,2048,3072]






head number-[4,8]







decoder layer-[1,2,3,4,5,6]








encoder固定设置为6层，
decoder可以选择使用encoder倒数1,2,3层的结果








Q,K,V维度固定为512








整个搜索空间包含10^15可能的SubTransformer








Hardware-Aware Evolutionary Search Setups









latency predictor的输入是SubTransformer的特征向量，包含：layer number,embed dim,hidden dim,attention head,attends layer
会将SubTransformer，latency组成输入pair，作为一个三层MLP的输入，隐藏层维度是400，使用ReLU激活。









三层的模型会比一层的更精确，并且超过三层以后精度也不会有进一步提升。
借助latency predictor，本文会执行30次进化算法，种群规模为125，parents:mutation:crossover = 25:50:50

Section IV Results

Part 1 HAT Performance

Table 8,Fig 7,8展示了HAT与其他框架在四类翻译任务上的性能对比。

本文设置了不同的延迟约束，获得不同硬件平台上的HAT模型，可以看到搜索到的HAT结构始终比基线模型性能高出一大截。可以看到在ARM cpu上搜索到的HAT结构在达到相同精度的情况下推理速度提升了3x，模型规模下降了3.7x.
CPU上搜索到的HAT获得了2x的加速；GPU上搜索到的结果可以看到，通过尺寸的缩放很难进一步减少延迟，但是在这种情况下HAT依旧可以搜索到既有低延迟又有高性能的模型。



Table 2展示的是与其他Transformer结构的对比结果，依旧是加速+模型规模更小。尤其与Envolved Transformer相比，HAT总的GPU训练时间减少了12000x。
Table 3展示的则是与其他目前最新模型的对比结果，痛死对比了延迟和精度，可以看到HAT在达到最高精度前提下延迟也是最小的。

Part 2 Analysis

在HAT WMT模型中，decoder中10%使用了encoder后三层，40%使用了encoder后两层的结果，这也说明使用arbitrar encoder-decoder注意力的必要性。
本文还可视化了不同硬件平台搜索到的模型，可以看到GPU倾向于更宽更浅的网络模型，树莓派则倾向于更深更细的模型，这与本文Fig 2显示的实验结果相吻合，即GPU对网络放缩并不敏感，但是树莓派对这方面比较敏感，这一实验结果也会指引我们的设计，在GPU模型中可以减少网络层数，增加层内维度，从而提升精度减少延迟。

Part 3Ablation Study

与最大的SubTransformer相比，搜索到的HAT延迟降低了1.5x，模型尺寸减少了1.5x,这说明大的模型并不总是最佳的，本文还对比了进化搜索和随机搜索的区别，可以看到进化搜索可以找到loss更低的模型。

SubTransformer Performance Proxy

所有的子网络会以被均匀采样然后被训练，Table 5展示了使用performence proxy的必要性，可以看到通过权重继承与从头训练获得的性能十分接近，更重要的是二者的相对精度量级十分接近，说明可以借助proxy来搜索高性能的模型架构，这样可以显著降低搜索成本。

Low search cost


Fig 10展示了不同方法的搜索成本对应的CO2排放，可以看到本文HAT比Envolved Transformer降低了12041x，虽然二者均使用了进化搜索，但是关键的区别在于Envolved Transformer需要训练所有的模型然后将他们按性能排序才能进行选择；而HAT是在SuperTransformer中一起训练所有模型 ，在将他们按性能排序进行选择。HAT最终的性能表示proxy足够精确来代表子模型的精度了。



Finetuning Inherited SubTransformers


本文会对每一个搜索到的子网络再从头进行训练；本文也可以选择继承超网权重后进行微调，Table 6的结果显示继承+微调的方式可以取得和从训练相近的精度，但是训练成本却会降低4x。因为SupreTransformer的训练成本可以分摊在所有SubTransformer的搜索中。


Quantization Friendly


HAT与其他网络压缩技术是完全正交的，比如量化。本文使用k-means量化进一步减少模型尺寸。结合4bit量化后模型尺寸减少了25x，BLEU损失可以忽略不计，并且8bit模型的BLEU比全精度模型还要高0.1充分说明了搜索到的HAT的鲁棒性。


Knowledge Distillation Friendly



HAT值关注如何搜索到高效的网络结构，而KD关注的是如何更好的训练指定结构，因此二者也可以结合。通过蒸馏token-level的信息，教师模型BLEU为28.8，参数量为49M，student 模型则包含30M参数量，精度从25.8提升到了26.1。






Section V Related Work

Transformer

Transformer最开始用于序列建模任务，通过堆叠多个block可以组件容量很大的模型但是延迟也很高。因此近期的一个研究方向就是提升Transformer的效率，比如使用卷积替换注意力模块，或者通过一次传递encoder多层来训练更深的Transformer模型。
HAT则使用架构搜索的思路。

另一研究趋势是使用自回归的方式来截断decoding的迭代次数，这样虽然减少了延迟但是也降低了精度。
还有的探究如何将encoder的结果经过线性组合后作为decoder的输入。但是HAT是直接将encoder的输出级联后作为dencoder的输入，这样可以更好的保留low-level信息。

还有的就是提出使用多分支的Lite Transforemer但是使用的评价指标是FLOPs，并不十分精确。本文提出的放大可以很好的与KD，量化结合，都可以进一步优化硬件延迟。


Neural Architecture Search


在计算机视觉领域NAS备受关注，主要通过强化学习、进化搜索等还有的基于DARTS。MNasNet，ProxyLessNAS,FBNet等则是在优化过程中考虑了硬件约束。而在降低NAS搜索成本方面也有诸多研究，如基于超网的方法通过使用一个代理来计算子网络的性能，在NLP领域NAS的应用还有许多探索空间，比如2019年Envolved Transformer提出在模型规模限制下进行Transformer架构的搜索，但是其搜索成本十分高昂（250 GPU years），无法进行硬件部署。除此之外没有对硬件延迟做特定考虑，而不用硬件具有不同的架构和特性，来自硬件的延迟反馈对于搜索高效的模型是十分重要的。

Section VI Conclusion

本文提出硬件感知的Transformer架构（Hardware-Aware Transformer）来搜索适合不同硬件平台的Transformer架构，本文在足够大的搜索空间中执行带有特定硬件延迟约束的搜索任务，通过共享SuperTransformer的权重从而在较低搜索成本的前提下获得精度、效率复合要求的Transformer模型。

本文希望HAT可以为Transformer的现实应用开辟一条新的探索途径。