【KDD2021】用NAS实现任务无关且可动态调整尺寸的BERT压缩

不同设备往往对训练的模型有不同的要求，如模型占用内存的大小、推理时延等。为了满足设备的要求，技术人员往往需要对大规模的预训练模型进行压缩处理。但是由于存在大量复杂的场景和不同的下游任务，专门为每一种场景设计压缩过的BERT模型，需要投入大量的人力和计算资源。

对此，微软亚洲研究院的研究员们提出了 NAS-BERT 技术（相关论文已发表在 KDD 2021 大会上），可以直接在上游预训练任务中进行压缩训练，使压缩模型不再依赖于下游的具体任务。并且 NAS-BERT 只需通过一次搜索就可以直接产生不同尺寸的模型，进而满足不同设备的要求。所以即使是资源受限的设备，也可以动态选择合适的架构。相较于传统方法，NAS-BERT 极大地提高了训练效率。

论文链接：https://arxiv.org/abs/2105.14444

如表1所示，之前大部分的压缩工作都是针对具体任务本身而设计方法实现的。比如，DisilBERT 等工作会直接在上游模型进行压缩学习，对下游任务无关，但不能适应各种不同受限资源的场景。DynaBERT 则可以在下游任务中通过一次训练产生多种不同大小的模型，以达到动态调整尺寸的目的。而 NAS-BERT 与其他几种工作不同，它可以同时做到任务无关和动态调整尺寸，并且达到优越的性能。

表1：之前的 BERT 压缩工作

为了让 NAS-BERT 实现上述目的，研究员们直接在上游预训练阶段对 NAS-BERT 进行了架构搜索（NAS）与训练，并且一次性搜索出各种大小的模型，以方便各种场景的部署。通过尝试架构层面的组合方式，NAS-BERT 不仅能够探索模型本身的潜力，还可以搜索出更加优越的架构。

具体而言，NAS-BERT 首先在上游预训练任务上训练一个超网络。其中，超网络中的一个有向无环图就是一个架构，不同的架构在超网络中共享参数来降低训练成本。然后研究员们再通过自动架构搜索技术去训练这个超网络。由于超网络中包括各种尺寸的模型，因此可以覆盖到各种不同的模型大小。最后，在用超网络去评价每个架构的性能并选择出优越的架构。

然而在 BERT 预训练任务上做自动架构搜索是具有挑战性的。其主要原因有：1. BERT 预训练本身收敛慢并且需要巨大的计算资源；2. 巨大的搜索空间会进一步导致训练超网络变得困难。为了解决这两个问题，NAS-BERT 主要采用了模块化搜索和逐渐缩小搜索空间的技术。同时，为了在缩小搜索空间的基础上能够给出各种不同的架构，NAS-BERT 还使用了分桶搜索把搜索空间分成众多不同的桶，不同桶中的架构有着不同的模型大小和延迟，桶中的架构相互竞争，从而选出最优越的架构。并且不同桶的架构通过权重共享，以降低架构搜索的代价。通过使用这些技术，NAS-BERT 才得以快速地搜索出各种不同的优越架构。