KDD 2021 | 用NAS实现任务无关且可动态调整尺寸的BERT压缩（2）

渐进架构搜索，分桶设计与模型选择

虽然分模块搜索已经大大减小了搜索空间，但是训练超网络依然需要很长的时间收敛。而且大量的模型在竞争有限的资源，将导致每个模型的评估结果都不准确，严重影响了架构搜索的排序性能。所以研究员们在训练过程中的每一轮 epoch 结束时，会把搜索空间中一些没有希望的架构裁剪掉，然后在裁剪后的搜索空间上继续训练。通过这种方式，可以将更多的计算资源给到更有希望的架构，进而得到更加准确的评估，这既减少了资源的浪费，又加快了搜索过程。

然而直接对整个搜索空间进行裁剪，会导致无法得到大量不同大小的架构。因为超网络在搜索过程中可能偏向于大模型或者收敛快的模型，所以无约束的裁剪会使得最后得到的架构大小趋于相同。基于此，研究员们提出了分桶裁剪的方法，并且采用具体的数字为例，阐述了在分桶裁剪过程中，搜索空间是如何变化的。

首先研究员们将搜索空间均匀地分成10个桶，每个桶的架构被约束在不同的参数量和延迟下。尽管不同的架构落在不同的桶里，但是在训练的时候，它们的参数仍然是共享的。在每个迭代训练结束的时候，研究员们会随机从每个桶里采样出2000个架构，同时使用超网络评价它们在验证集上的分数，并删除一半的架构。通过重复这个过程，可以不断地收缩搜索空间，直到每个桶里只剩下10个架构。这样对于一个搜索模块而言，最终就得到了100个架构。假设共有4个搜索模块，那么就有100^4种组合方式，每一种组合方式都是一个完整的架构。

由于每个搜索模块中的100个架构来自不同的桶，因此它们有着显著不同的大小和延迟。所以通过组合不同的搜索模块而得到的100^4种架构，也具有显著不同的模型大小和延迟。

为了评估每个架构的性能，研究员们会首先评价每个搜索模块中100个架构在验证集上的损失。然后用不同模块的损失直接加和作为组合架构的性能。通过这种方式，仅仅通过4*100次验证集测试，就可以粗略得到所有架构的性能。该方法对于延迟的评估也是类似的，即先测量每个搜索模块内架构的延迟，然后再用不同模块的延迟加和作为整个架构的延迟。

所以当给定任何一种关于模型大小和延迟的约束时，可以通过查表的方式快速找到性能最高的架构。这种架构本身是在上游预训练任务中挑选出来的，因此与下游任务无关。

实验结果

为了评估搜索得到的架构性能，研究员们重新在上游任务上训练了这个架构，然后在各种下游任务中测试了模型的性能。

在搜索阶段使用的教师模型是 BERT 110M 的模型。教师模型和搜索出来的架构都是在 BookCorpus+English Wikipedia (16GB) 数据上进行了预训练。为了和 BERT 的12层 Transformer 对齐，超网络包含了24个子层（每一个 Transformer 层等于一层 MHA 加一层FFN）。研究员们在大量的下游任务上评估搜索出来的架构的性能，包括 GLUE 和 SQuAD 数据集，并选择了5M，10M，30M，60M参数量的架构为例，来展示搜索出来的架构效果。

研究员们首先把 NAS-BERT 搜索出来的架构和手工设计的 BERT 架构进行了对比。从表3中可以看到，NAS-BERT 在各种模型大小及延迟条件下，都超过了手工设计的 BERT 模型。

表3：NAS-BERT 和 BERT 对比

进一步，研究员们将 NAS-BERT 和之前的模型压缩工作进行对比。之前的模型压缩工作主要采用了多种蒸馏技巧和训练方式。而 NAS-BERT 只使用了简单的两阶段蒸馏来突出其搜索出的架构优势，不使用复杂的蒸馏技术或训练方法，如注意力蒸馏，逐层替换等。从表4可以看到，NAS-BERT 超过了之前的方法，证明了架构本身的优越性。

表4：NAS-BERT 和之前的 BERT 压缩工作对比

为了证明缩小搜索空间方法的有效性，研究员们又做了一组实验——与不使用逐渐缩小搜索空间进行对比。从图3的损失函数曲线来看，使用逐渐缩小搜索空间显著加快了收敛速度。从最终搜索得到的架构来看，使用逐渐缩小搜索空间能够帮助搜索到更好的架构。

图3：逐渐缩小搜索空间的分离实验

同时，研究员们还探究了缩小搜索空间的其他方法。与从架构层面缩小搜索空间不同，研究员们尝试了从操作层面缩小搜索空间的方法（具体见论文）。从表5可以看到，从架构层面缩小搜索空间更为准确，因此能搜索到更好的架构。

表5：不同逐渐缩小搜索空间的分离试验

最后，研究员们使用了各种不同的训练方式来训练 NAS-BERT，并将其与 BERT 模型对比，以证明搜索的架构的鲁棒性。并且研究员们还尝试了在训练过程中，在上游预训练阶段或下游微调阶段或两者都有的情况下，观察 NAS-BERT 搜索的架构的鲁棒性。从表6中可以看到，NAS-BERT 在各种训练配置上，都能显著超过手工设计的 BERT 架构。

表6：不同训练方式的分离实验

表7展示了部分 NAS-BERT 搜索得到的架构。可以发现，搜索出来的架构都由不同的操作且复杂的方式组成，这证明了 NAS-BERT 能够搜索出更加新颖的新架构。

表7：部分 NAS-BERT 搜索得到的架构

微软亚洲研究院的研究员们在本篇论文中提出的 NAS-BERT，是一种用自动架构搜索方法实现任务无关且动态调整尺寸的 BERT 压缩技术。其拥有新颖的搜索空间、卷积、注意力、前馈网络以及不同的隐藏层大小。加上高效的搜索方法，NAS-BERT 可以探索出不同操作的复杂组合方式得到模型的潜力。研究员们通过大量的比较和分离实验，也证明了 NAS-BERT 搜索得到架构的有效性。