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

SuperTransformer 母网络训练
为了得到高效模型，构建一个足够大的搜索空间至关重要。然而，完整训练搜索空间中的子网络来比较精确度的开销过大。之前的 SOTA 模型 The Evolved Transformer (So et al. 2019) 就是利用这种方法进行搜索。去年 ACL2019 中， Strubell et al. (2019) 指出了 Evolved Transformer 搜索排放的二氧化碳接近五辆汽车在整个使用寿命中的二氧化碳排放总量，造成了巨大的环境负担。 MIT Technology Review 也对这一点进行了报导。
因此，为了可以实现环保 AI（Green AI），研究者采用 SuperTransformer 母网络来提供精确度的近似。母网络可以快速测试搜索空间中的任意子网络精度，而无需子网络训练。 SuperTransformer 母网络是搜索空间中最大的模型，并且通过权重共享（Pham et al. 2018）包含了搜索空间中全部网络。所有的子网络 SubTransformer 共享他们在母网络中的共同部分。例如，所有子网络 SubTransformer 共享弹性词向量和隐藏层权重的前部，共享注意力层中的 Query、Key 和 Value ，也共享编 / 解码器的共同前几层。
在训练过程中，研究者均匀地对在搜索空间中的所有子网络进行采样，得到梯度并且更新子网络对应部分的权重。所有的子网络获得了等量的训练并且均可以独立地完成任务。在实际中，母网络 SuperTransformer 所需的训练时间与单个基线 Transformer 相近，所以开销很低。

本文插图
图 6: 母网络 SuperTransformer 权重共享。
完成训练后的母网络是一个快速且精确的模型精度近似器（Performance Proxy）。给定一个子网络 SubTransformer ，可以从母网络继承这部分权重，并且快速地测试得到 validation loss 。这样就无需完整训练子网络。通过下图可以看到，近似器是非常精准的。子网络的 validation loss 越低，最终的完整训练得到的 BLEU 越高。

本文插图
图 7: 母网络 SuperTransformer 提供精确的子网络 SubTransformer 精度近似。
进化搜索（Evolutionary Search）子网络
给定一个目标硬件上的运行时间限制，研究者使用进化搜索算法来搜索得到子网络 SubTransformer 。

本文插图
图 8: 进化搜索（Evolutionary Search）子网络。
进化搜索引擎利用 SuperTransformer 来获得子网络的 validation loss ，并利用一个运行时间预测器来得到子网络在目标硬件上的运行时间。在每一步中，他们只将运行时间短于时间限制的子网络加入到种群（Population）中。具体而言，他们采用种群大小 125 ，母种群大小 25 ，重组 (Crossover) 种群大小 50 ，突变 (Mutation) 种群大小 50 ， 0.3 突变几率。
有两种方法可以获取硬件运行时间。第一，线上测试法，在搜索的过程中即时测试子网络的运行时间。然而，这种方法有数个缺点。硬件每次运行的时间受很多因素限制，很不稳定，例如散热的影响。对于每个模型都需要运行数百次取平均来得到较精确的时间。另外，在进化搜索中，每一步都需要测试种群中的很多个子模型的时间，因此线上测试非常耗时，会成为整个搜索过程的瓶颈。
因此，在 HAT 中，研究者使用线下方法，训练一个预测器来快速且精确地给出子模型的运行时间。具体来说，他们将子网络的架构编码成一个向量，并且采集一个 [网络架构，运行时间] 数据对的数据集。然后训练一个多层感知器模型（MLP）来回归子网络的运行时间。对于每个硬件，他们采集 2000 个数据点，然后训练一个三层 MLP ，每层维度 400 ，使用 ReLU 为激活函数。