ACL2020| 模型压缩25倍,MIT韩松组提出高效适配不同硬件的HAT模型

运用注意力机制的 Transformer 模型近几年在 NLP 领域获得了广泛应用。然而，由于参数量和计算量巨大， Transformer 模型难以在存储和算力有限的边缘硬件设备上高效部署。为了解决 Transformer 的低效问题，来自 MIT 的研究人员提出了 HAT: Hardware-Aware Transformers ，针对不同的硬件设备的特性，为每个硬件搜索出一个高效的 Transformer 模型，从而在保持精确度的前提下大幅降低内存消耗。在同样的精度下，相比于基线 Transformer ，HAT 可以获得 3 倍加速， 3.7 倍模型压缩。
该论文已被自然语言处理顶会 ACL 2020 收录。此外， HAT 的所有代码和模型已经在 GitHub 上开源，作者也将在 7 月 8 日 / 9 日的 ACL 大会上线上宣讲他们的工作。

本文插图

论文链接：https://arxiv.org/abs/2005.14187
GitHub：https://github.com/mit-han-lab/hardware-aware-transformers
B站介绍: https://www.bilibili.com/video/BV1mt4y197FL/

近年来，自然语言处理领域受 Transformer 的驱动获得了快速发展， Transformer 模型也被广泛应用于多种任务，例如机器翻译、聊天机器人、文本摘要等等。然而， Transformer 的高准确性需要非常高的模型参数量和计算量来支撑，这对于受到存储大小、算力和电池容量限制的边缘计算设备（比如手机和物联网硬件）来说是很大的挑战。例如，如果在树莓派上使用 Transformer-Big 模型，它需要运行 20 秒才可以完成一个 30 词长度的句子的翻译，这对于很多需要实时对话和反馈的场景来说是无法接受的。
为了获得更高效和快速的 Transformer 模型， MIT 等机构的研究者提出了 HAT: Hardware-Aware Transformers ，借助神经网络搜索（NAS）技术，在搜索过程中加入硬件反馈，来对每一个硬件平台设计一个专用的高效 Transformer 网络结构。
在搜索空间上，他们提出打破之前设计 Transformer 的两点陈规。首先，传统 Transformer 的所有解码层都从最后一层编码层获得输入，这样形成了一个从编码器到解码器的信息瓶颈，对小模型尤其不友好；对此，研究者提出“任意编码器 - 解码器注意力”（Arbitrary Encoder-Decoder Attention），允许解码器获得多个和任意某些个编码器的输出。从而使得编码器不同的抽象层的信息都可以被解码器获取和使用。另外，传统 Transformer 的所有层都有相同的网络结构，研究者提出“异构层”（Heterogenous Layers）来使得每层都可以有不同的隐藏层维度（Hidden Dim）和注意力头数 (Head Number) ，编 / 解码器也可以有不同的词向量长度 (Embedding Dim) 和层数等等。
在搜索算法上，为了能够减少搜索开销，实现环保 AI ，他们采用权重共享的方法来训练出一个母网络 SuperTransformer 来涵盖在搜索空间中的全部可能模型，这样一来，其中的每个子网络 SubTransformer 可以直接继承母网络中对应部分的权重，快速估计子网络的精确度。之后，研究者使用进化搜索（Evolutionary Search），利用预先设定好的在目标硬件上的运行时间来作为限制，在母网络中搜索出一个满足运行时间并且精度够高的子网络，这就是他们想要得到的为目标硬件设计的专用 SubTransformer 网络。
研究者在四个机器翻译任务（WMT‘14 EN-De（英语 - 德语）、WMT’14 EN-Fr（英语 - 法语）、WMT’19 EN-DE（英语 - 德语）和 IWSLT‘14 De-En（德语 - 英语）），以及三个不同的硬件平台（树莓派、Intel CPU 和 Nvidia GPU）上验证了论文所提方法的有效性。
在同样的精度下，相比于基线 Transformer ，HAT 可以获得 3 倍加速， 3.7 倍模型压缩；相比于 Evolved Transformer ， HAT 有 2.7 倍的加速和 3.6 倍的模型压缩，并且将搜索开销降到了 1.2 万分之一。同时， HAT 的方法也与其他模型压缩技术兼容，例如，研究者将搜索得到的 HAT 模型进行了 4-bit 压缩，进一步获得了仅为基线 Transformer 1/25 大小的模型。