Hardware-Aware Transformers (HAT)——Transformer模型的NAS策略

Paper地址: https://arxiv.org/abs/2005.14187

Github: GitHub - mit-han-lab/hardware-aware-transformers: [ACL 2020] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing

Introduction

对Transformer进行网络结构搜索时，会面临两个关键问题，且主要跟实际硬件平台特性有关：

• FLOPS不能完全代表Transformer模型的计算速度/实际执行效率：

• 不同硬件平台上，Transformer最适配的网络结构不同;

Method

• 搜索空间：Encoder-decoder Attention的任意连接方式，以及Transformer layer内部结构（Attention head数目、Width等）；
• 超网络预训练：训练了一个SuperTransformer；
• 执行速度预测：训练了一个MLP，用于预测不同结构、在指定硬件平台上的执行速度，作为Cost model或Predicter；
• 搜索策略：在超网络预训练之后，结合资源约束，通过进化搜索算法（EA: Evolutional Algorithm）寻找最优子网络；

算法步骤

• 超网络预训练与延迟预估：

1. 第一步：按照One-shot NAS的思想，在超网络中采样子网络进行训练，从而预训练得到一个权重共享的超网络（SuperTransformer）;
2. 第二步：在超网络中采样子网络，并在指定硬件平台上实际测速；
3. 第三步：根据第2步得到的数据（精度、速度），训练一个Latency Predictor；

• 进化算法搜索（EA）；
• 选择符合约束条件的、最优子网络模型，然后微调或重新裸训；

搜索空间

• Encoder-decoder attention的任意连接方式；
• 不同Encoder的Key tensor和Value tensor，在Seq_length维度上执行Concat，作为Encoder-decoder attention的Key和Value；

• 层内超参：
  • Attention Head数目；
  • Hidden size in FFN；
  • Embedding Dimension；
  • Layer数目、或网络深度；

精度预测

• 在SuperTransformer超网络上采样子网络，子网络的权重取自超网络的前面/左边部分，然后进行初始化，再进行预训练；
• 不同维度执行均匀采样；

速度预测

• 从训练好的超网络中，采样出大量的子网络，直接评估精度、测量硬件平台上的执行速度，然后基于这些数据训练Latency predictor；

Experiment