一次训练、多端部署——华为提出宽度和深度可伸缩的动态预训练模型DynaBERT

 华为诺亚方舟实验室语音语义团队提出了一种新型动态BERT模型：DynaBERT，该模型可以选择不同深度和宽度的子网络进行推理，达到“一次训练，多端部署”的效果。实验表明，在同等性能约束（包括参数量、FLOPs、时延）下，DynaBERT的最大子网络与BERT-BASE性能相当，而较小宽度和深度的子网络则显著优于现有的BERT压缩方法。 论文发表在：https://arxiv.org/abs/2004.04037  。

图1：宽度和深度可伸缩的DynaBERT训练过程。

研究背景

近来，基于Transformer结构的预训练语言模型（如BERT和RoBERTa）在自然语言处理任务上取得了显著成果。但是，这些模型参数量巨大，限制了其在存储、计算、功耗等性能有限的边缘设备上的部署。将BERT部署到这些设备的困难体现在两个方面：（1）各种终端设备性能各不一样，怎么给不同性能的设备部署适合自身性能的模型；（2）同一个终端设备在不同的情况下可接受的运行模型的存储、计算和功耗也不一样。例如，在手机上，如果同时运行着大量占用存储、计算资源的其他程序时，可分配给BERT模型的资源就会相应较少。为了解决这个问题，我们提出了一种宽度和深度可伸缩的动态预训练模型DynaBERT。该模型可以根据不同设备的硬件性能部署不同宽度和深度的子网络。并且一旦某个设备部署了BERT模型，该设备也可以根据自身资源条件动态加载已部署模型的一部分进行推理。

图2：将DynaBERT部署到性能不同的设备的示意图。

DynaBERT的训练过程如下：首先训练宽度可变的DynaBERTw，然后再训练允许宽度和深度都可以伸缩的DynaBERT。在训练DynaBERTw时，我们首先根据注意头和神经元的重要性重构每个Transformer层中的权重连接，以确保最重要的注意力头和神经元被更多的子网络利用。然后，我们在DynaBERTw中利用知识蒸馏技术将最大宽度学到的知识迁移到宽度更小的子网络中。DynaBERTw训练后，我们再次利用宽度深度两个方向上的知识蒸馏将最大深度的知识迁移到深度更小的子网络中，同时不遗忘宽度方向学到的可伸缩性。我们对BERT-BASE和RoBERTa-BASE进行了可伸缩训练，并且在在各种效率约束（参数量、FLOPs、推理速度）下和GLUE任务中广泛评估了我们提出的DynaBERT和DynaRoBERTa的有效性。

基于MHA中的注意力头和FFN中的神经元的重要性重构网络

在训练宽度可伸缩的DynaBERTw之前，我们需要对BERT模型的宽度进行定义。相比于用规则的卷积层堆积起来的卷积神经网络，BERT模型由多Transformer层堆叠而成，且每个Transformer模块中的多头注意力机制（MHA）和全连接层（FFN）分别在序列和特征两个方向做变换。因此，BERT模型的宽度不能像卷积神经网络那样简单地定义成每层卷积核的个数。

对每个层，我们将多头注意力机制MHA中的计算拆分成每个注意力头计算然后再相加，所以这MHA的宽度可以用注意力头的个数决定。全连接网络FFN的宽度可以用中间层的神经元的个数决定。因此，对于一个Transformer层，宽度可以通过MHA的注意力头和FFN中间层的神经元的个数决定。为了让重要的注意力头和神经元被更多的子网络共用，在进行宽度可伸缩网络训练之前，我们将MHA中的注意力头和FFN中间层的神经元在宽度方向上按重要性从高到低排列（图3）。

图3：利用MHA的注意力头和FFN中间层的神经元的重要性重构网络连接。

宽度可伸缩的DynaBERTw训练方法　

在宽度可伸缩的DynaBERTw训练中，我们首先使用一个已经按照图3重构好的BERT模型当作老师模型和初始化DynaBERTw，然后采用了知识蒸馏的方法来将最大宽度模型学到的知识迁移到宽度较小的子网络中去。以分类任务为例，知识蒸馏背后的基本思想是首先训练大型模型，然后通过logits和中间层特征表达的拟合，将其学习的知识迁移到小模型。具体地，对于某种宽度的DynaBERTw子网络，它的目标函数包括了三项，这三项分别使得该子网络的logits，词嵌入和隐层状态逼近老师模型。为了提升蒸馏的效果，我们在训练过程中使用了数据增强技术来进一步提升蒸馏效果。

图4：利用知识蒸馏训练宽度可伸缩的DynaBERTw。

宽度和深度同时可伸缩的DynaBERT训练方法

在训练DynaBERTw之后，我们进一步使用知识蒸馏来训练宽度和深度的同时可伸缩的DynaBERT。具体地，我们使用已经训练好的的DynaBERTw作为老师模型，并且初始化DynaBERT。然后将知识从最大深度的老师模型迁移到同等宽度不同深度的学生子网络。在DynaBERT训练过程中，为了避免忘记在宽度方向所学的可伸缩性，在每次训练迭代中我们依然会使用不同的宽度。具体地，对于一定宽度和深度的学生子网络，它的目标函数依然包括三部分。这三个部分分别使logits词嵌入和隐层状态逼近同等宽度但是深度最大的老师模型。

图5：利用知识蒸馏训练宽度和深度同时可伸缩的DynaBERT。

实验结果

  我们在GLUE benchmark上进行了实验。DynaBERT和DynaRoBERTa模型的最大尺寸与BERT-BASE和RoBERTa-BASE模型相同。训练中使用的宽度比例是[1, 0.75, 0.5, 0.25]，深度比例是[1.0, 0.75, 0.5]。因此，DynaBERT和DynaRoBERTa都可以在12种不同宽度和深度的子网络上进行推理。

表1：DynaBERT和DynaRoBERTa不同子网络在GLUE上实验结果。

表1是在GLUE dev集上的结果。可以看到，对于大多数任务，DynaBERT或DynaRoBERT能在更小的深度或者宽度达到与BERT-BASE和RoBERTa-BASE相似或者更高的精度，表明原始BERT或RoBERTa模型中存在冗余。另一观察结果是相同压缩比下，在宽度方向压缩往往比在深度方向压缩具有更高的精度。

图6：Kirin810 ARM CPU上时延和精度的对比。

图6是DynaBERT的比较和DynaRoBERTa和其他BERT压缩方法在麒麟810 ARM CPU上时延和精度的对比（其他性能约束之下的对比如参数量、FLOPs、NVIDIA GPU时延可以参考论文）。可以看出，在同样的时延DynaBERT的最大子网络与BERT-BASE性能相当，而在较小的宽度和深度上显著优于其他BERT压缩方法。