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

paper: HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
code: https://github.com/mit-han-lab/hardware-aware-transformers

1 介绍

Transformer网络取得了非常大的成功并具有广泛的应用。 然而由于它巨大的计算代价，使得部署到手机等端侧设备面临巨大的挑战。
在评估Transformer网络的效率时有两个常见的陷进： 1） FLOPs不能反映真实的时延； 2）不同硬件偏好不同的Transformer结构。

上图表明， 在一个硬件上表现好的网络在另一个硬件上表现反而不好。

上图表明： 1）FLOPs和lateny并不完全是线性关系； 2）不用硬件的影响因素不同。

受NAS成功的启发， 我们提出了搜索硬件感知的Transformer网络， 通过把latency的反馈纳入搜索流程， 我们可以搜索针对多种硬件友好的Transformer网络。

作者首先设计了一个巨大的搜索空间， 包含任意的encoder-decoder连接。传统的Transformer网络在encoder和decoder之间存在信息瓶颈，任意的encoder-decoder允许decoder可以连接多个不同的encoder而不仅仅是最后一个从而打破了这一瓶颈。

为了在这个巨大的搜索空间上进行低成本的搜索， 作者首先训练了一个包含许多subTransformer的superTransformer网络。

最后作者通过进化搜索去寻找满足latency约束的最优的subTransformer网络。 实验结果表明， HAT天然地可以与模型压缩技术如量化、蒸馏等融合。

2 方法

方案总览：

2.1 搜索空间

核心是2点：

• decoder可以任意连接多个不同的encoder
• 不同的层结构可以是不同的， 而不是像传统的Transformer网络那样每层都是相同的进行重复

2.2 superTransformer

为了搜索出性能更好的网络， 搜索空间需要足够大。 然而训练所有的模型然后评估他们的指标是不可能的。 因此作者采用了superTransformer的方法， 所有的subTransformer共享公共部分的权重。 在训练superTransformer时， 所有的子网进行均匀采样， 对应子网的权重被更新。 最终的superTransformer只需要训练跟原始的Transformer相同的步数。

2.3 进化搜索

搜索方法就是常规的进化算法， 这个不多说。 有一点需要关注的是，latency指标的评估， 作者采用了基于预测的方法。 预测器用的是很简单的MLP， 提取了10维的特征向量(这10维特征是：encoder嵌入维度，encoder层数，encoder平均隐层数，encoder平均自注意力头数， decoder嵌入维度，decoder层数，decoder平均隐层数，decoder平均自注意力头数，平均encoder-decoder注意力头数，平均decoder连接到encoder的数量 )，用了2000个训练样本， 最终的预测精度相当不错。

3 实验

细节部分不多介绍， 感兴趣的话可以参考论文。
简单描述一下搜索空间设置：

• embedding_dim: [512, 640]
• hidden_dim: [1024,2048,.3072]
• mutihead_num： [4,8]
• decoder layer num: [1,2,3,4,5,6]
• decoder connect encoder: [-1,1,2] # -1 means last one, 1 means last two, 2 means last three

以上是WMT数据集的设置， 对IWSLT数据集的设置略有不同。 在WMT数据集上， SuperTransformer总共训练40K 个step， 在IWSLT数据集上总共训练50K个step.

4 结果

4.1 结果对比

WMT’14 数据集上的对比：

WMT’19 和IWSLT 数据集上的对比

其他方面的对比：

4.2 结果分析

几点发现：

• 大约有10% 的decoder连接了3个encoder, 40% 连接了2个encoder
• 不同硬件对结构偏好不同： GPU偏好“宽”而“浅”的网络，树莓派手机偏好“窄”而“深”的网络。
• 更大的网络部不一定有更好的效果