[论文解读] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing

Hanrui Wang, Zhanghao Wu+5 more|arXiv (Cornell University)|May 28, 2020

Advanced Neural Network Applications参考文献 70被引用 21

一句话总结

HAT 提出了一种硬件感知的神经架构搜索框架，用于为特定硬件平台定制高效 Transformer 模型。通过训练参数共享的 SuperTransformer 并结合真实硬件延迟反馈的进化搜索，HAT 在树莓派 4 上实现了相较于 Transformer-Big 3 倍的推理速度提升和 3.7 倍更小的模型尺寸，同时将搜索成本降低至先前工作的 12,041 倍，且无准确率损失。

ABSTRACT

Transformers are ubiquitous in Natural Language Processing (NLP) tasks, but they are difficult to be deployed on hardware due to the intensive computation. To enable low-latency inference on resource-constrained hardware platforms, we propose to design Hardware-Aware Transformers (HAT) with neural architecture search. We first construct a large design space with $ extit{arbitrary encoder-decoder attention}$ and $ extit{heterogeneous layers}$. Then we train a $ extit{SuperTransformer}$ that covers all candidates in the design space, and efficiently produces many $ extit{SubTransformers}$ with weight sharing. Finally, we perform an evolutionary search with a hardware latency constraint to find a specialized $ extit{SubTransformer}$ dedicated to run fast on the target hardware. Extensive experiments on four machine translation tasks demonstrate that HAT can discover efficient models for different hardware (CPU, GPU, IoT device). When running WMT'14 translation task on Raspberry Pi-4, HAT can achieve $ extbf{3} imes$ speedup, $ extbf{3.7} imes$ smaller size over baseline Transformer; $ extbf{2.7} imes$ speedup, $ extbf{3.6} imes$ smaller size over Evolved Transformer with $ extbf{12,041} imes$ less search cost and no performance loss. HAT code is https://github.com/mit-han-lab/hardware-aware-transformers.git

研究动机与目标

• 为解决在资源受限的硬件平台（如 CPU、GPU 和物联网设备）上部署计算密集型 Transformer 模型的挑战。
• 克服以 FLOPs 作为延迟代理的局限性，因为 FLOPs 无法准确反映不同硬件平台之间的实际性能差异。
• 降低现有神经架构搜索方法在 Transformer 模型上，尤其是针对硬件特化的模型，所面临的极高搜索成本。
• 实现在不牺牲模型性能的前提下，高效、准确且硬件特化的 Transformer 模型。

提出的方法

• 构建一个包含任意编码器-解码器注意力机制和异构 Transformer 层的大规模搜索空间，以支持架构多样性。
• 训练一个在所有子网络（SubTransformers）之间共享参数的 SuperTransformer，从而实现对所有候选架构的高效联合训练。
• 使用带有实时硬件延迟反馈的进化搜索，为每个目标硬件平台识别最优的 SubTransformer。
• 集成一个基于实测延迟数据训练的延迟预测器，以加速搜索过程并确保硬件感知的优化。
• 利用 SuperTransformer 的性能代理避免对每个 SubTransformer 从头开始重新训练，从而大幅降低搜索成本。
• 将 HAT 与量化和知识蒸馏等模型压缩技术结合，进一步提升效率。

实验结果

研究问题

• RQ1 硬件感知的神经架构搜索能否在 CPU、GPU 和边缘设备等多样化硬件平台上显著降低推理延迟？
• RQ2 用真实硬件延迟反馈替代 FLOPs 作为优化目标，是否能生成性能更优、更高效的 Transformer 模型，相比基于 FLOPs 的优化方法？
• RQ3 参数共享的 SuperTransformer 是否能有效近似大规模架构空间中子网络的性能，从而实现低成本的架构搜索？
• RQ4 任意编码器-解码器注意力机制与异构层的结合，如何提升不同硬件平台上的模型效率与性能？
• RQ5 与 Evolved Transformer 等基于 NAS 的先前方法相比，HAT 是否能在显著降低搜索成本的前提下实现更优的效率增益？

主要发现

• 在树莓派 4 上，HAT 在 WMT’14 En-De 翻译任务中实现了相较于 Transformer-Big 3 倍的速度提升和 3.7 倍更小的模型尺寸，且无性能损失。
• 与 Evolved Transformer 相比，HAT 在相同任务上将搜索成本降低了 12,041 倍，同时实现了 2.7 倍的速度提升和 3.6 倍更小的模型尺寸。
• 在 GPU 上训练的 HAT 模型在 ARM CPU 上表现不佳，反之亦然，证实了由于延迟因素的差异，硬件特定优化至关重要。
• 在 WMT’14 En-Fr 任务中，采用 4 位量化后，HAT 实现了模型尺寸 25 倍的缩减，BLEU 分数仅比全精度基线下降 0.1。
• HAT 的 8 位量化模型甚至在 WMT’14 En-Fr 任务中比其全精度版本高出 0.1 BLEU，证明了搜索过程的有效性。
• 基于真实硬件测量数据训练的延迟预测器能够提供准确且快速的反馈，使进化搜索既高效又具备硬件感知能力。