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

编者按：如今，基于 Transformer 的大规模预训练语言模型，如 BERT、XLNE、RoBERTa 和 GPT-3 等，已经在很多自然语言处理任务中都取得了十分惊人的效果。但是巨大的模型尺寸，使其在众多不同的下游任务中进行部署时非常困难。而且由于存在大量复杂的场景以及不同的下游任务，单独为不同场景设计一种压缩过的 BERT 模型既耗时又耗力。

因此，微软亚洲研究院的研究员们针对这些问题展开了研究，并提出了 NAS-BERT 技术。相关研究论文“NAS-BERT：Task-Agnostic and Adaptive-Size BERT Compression with Neural Architechture Search”已被跨学科数据科学会议 KDD 2021 收录。欢迎感兴趣的读者点击阅读原文，查看论文全文。

不同设备往往对训练的模型有不同的要求，如模型占用内存的大小、推理时延等。为了满足设备的要求，技术人员往往需要对大规模的预训练模型进行压缩处理。但是由于存在大量复杂的场景和不同的下游任务，专门为每一种场景设计压缩过的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 才得以快速地搜索出各种不同的优越架构。

搜索空间的架构设计

NAS-BERT 的架构设计（如图1所示），首先使用了一个预训练的 BERT 模型作为教师模型，然后用教师模型指导超网络的搜索。为了降低搜索空间，NAS-BERT 把学生模型（超网络）分成了几个不同的搜索模块，然后每个搜索模块会进行单独的训练。超网络模块的监督信号由教师模型的相应模块给出。换言之，就是给定一批数据，可以得到教师模型的每个模块输入输出的隐藏表征，再用这些输入输出的隐藏表征去训练搜索模块。这样，不同的搜索模块可以分开搜索且并行训练，指数级地降低了搜索空间。

对于每一个搜索模块，网络设计如图1(a)所示，网络的每一层包括了所有可能的搜索的候选操作（例如卷积、注意力网络）。一个架构就是一条从底部到最上层的单向路径，不同的架构（路径）共享权重以降低训练所需要的资源。

图1：NAS-BERT 方法框架图

为了使搜索的架构多样化，研究员们分别采用了三种不同的操作方法：注意力网络 MHA，前馈网络 FFN 和可分离卷积 CONV。MHA 和 FFN 是 Transformer 模型中的操作方法，采用这两种方法，可以使 Transformer 容纳在搜索空间中。由于 CONV 已经在很多自然语言任务中取得了不错的效果，因此研究员们也将其加入了搜索空间，以探索卷积和其他操作组合的潜力（操作设计请见图2）。同时，为了使搜索的架构具有各种不同的模型大小，研究员们让每一种操作都有128/192/256/384/512这5种不同的隐藏层大小，如表2所示。与此同时，还加入了 Identity（无操作），使其能搜索出各种不同层的架构。

表2：操作集合的设计

图2：卷积操作的设计

通过采用这种设计，可以使一个6层的搜索模块中有超过20万种可能的架构。如果采用4个搜索模块，架构的组合方式将超过10^20种。为了训练这个超网络，每次得到一批数据，就要随机采样一个架构进行训练，进而监督和学习教师模块的隐藏表征。由于搜索模块采样出来的架构输入输出的隐藏层大小可能和教师模块不一致，所以可以直接引入一个可学习的线性层进行转换，并和超网络一起训练。