Adabert:基于Automl的BERT自适应压缩

论文链接：https://arxiv.org/pdf/2001.04246v1.pdf

代码链接：https://github.com/alibaba/EasyTransfer/tree/master/scripts/knowledge_distillation

随着BERT等预训练语言模型在各项任务上都取得了STOA效果，BERT这类模型已经成为 NLP 深度迁移学习管道中的重要组成部分。但 BERT 并不是完美无瑕的，这类模型仍然存在以下两个问题：模型参数量太大，训练推理速度慢。目前主要有一下几种方法将BERT模型压缩为小模型：

1、架构改进：将原有的架构改进为更小/更快的架构，例如，将 RNN 替换为 Transformer 或 CNN，ALBERT替代BERT等；使用需要较少计算的层等。当然也可以采用其他优化，例如从学习率和策略、预热步数，较大的批处理大小等；

2、模型压缩：通常使用量化和剪枝来完成，从而能够在架构不变的情况下减少计算总量；

3、知识蒸馏：训练一个小的模型，使得其在相应任务行为上能够逼近大的模型的效果，如DistillBERT，BERT-PKD，TinyBERT等。

一、方案概述

而AdaBERT 的核心思想是把压缩后的小模型结构当做一组可学习的参数，这样不同任务压缩后应该使用什么模型结构就可以交由算法自动地搜索出来。Adabert将压缩后小模型放在知识蒸馏视角下，将搜索目标融入到学生模型的损失函数中。即Adabert主要是利用可微神经网络搜索DARTS自动搜索蒸馏后的小模型。

如下图所示，作者综合考虑了两种不同的损失函数：一是面向任务的知识蒸馏损失，为模型搜索过程提供线索；二是模型效率感知型损失，为模型搜索提供约束。这两个损失项使得 AdaBERT 能在压缩后的模型效率和有效性之间搜索一个合适的 trade-off。

数学形式上，把在目标数据集上 fine-tune 的 BERT 模型记为  ，模型结构的搜索空间记为  ，把搜索到的最优模型结构记为  ，那么整个模型的损失函数就是：

其中，wα是结构α的可训练网络权重，LCE，LKD和LE分别是目标任务，面向任务的知识蒸馏和效率的损失。 具体来说，LCE 是相对于目标数据Dt中标签的交叉熵损失，LKD是每个cell对应的output logits与teacher提供的logits之间的交叉熵损失， 是面向任务的知识蒸馏损失，为找到适用于任务的模型结构提供引导，而LE 是惩罚参数更大推理速度更慢的op以及cell个数的效率感知损失，目的是鼓励更轻量化的神经架构。 γ和β是用于平衡这些损失项的超参数。

二、搜索空间

在搜索过程中首先定义如下所示的神经网络架构：

图2 搜索空间

如上图所示，图中每个cell主要由两个输入节点、一个输出节点和中间节点构成。其中，中间节点的数量是可选的。每条进入中间节点的边有10种候选算子，分别为：cnn3, cnn5, cnn7, dilated_cnn3, dilated_cnn5, dilated_cnn7, avg_pool, max_pool, identity, zero。我们的目的就是学习架构参数，表明每条边我们应该选取哪种算子，同时表明应该保留哪些边（每个中间节点只有两条入边）。

在搜索cell参数的同时，还要搜索堆叠层数K。K对于在模型性能和效率之间至关重要，因为较大的K会有较高的模型容量，但推理速度较慢。

三、知识蒸馏损失

Adabert利用probe分类器从teacher模型中逐层分解有用的任务知识，然后将知识提取到压缩模型中。 具体来说，首先冻结teacher模型的参数，并根据groundtruth任务标签为每个隐藏层训练一个Softmax probe分类器。 总共有J个分类器，第j个分类器的分类logits可以看作是从第j层学到的知识。 给定输入实例m，将Cj,mT表示为teacher模型的第j层的隐藏表示，将Cj,mT表示为经过压缩的student模型的第i层上的隐藏状态，提取的有用的任务知识(分类logits)为：

   

其中，PjT是第j个teacher模型的probe探针分类器，PiS是在压缩模型的第i层上的可训练probe。

对于不同的下游任务，由于BERT中共每一层的作用都不一样。所以将所有层的分解知识集合地组合为：

其中，M表示共M个训练样本，ym是第m个样本的标签，wi,m是每个样本依据教师模型探针PjT的负交叉熵损失(loss)设置的归一后权重。用loss归一化的的动机是对于探针模型的预测结果越准确的层（loss 越小），这些层训练的越好或者越适合目标 task，就会被赋予更高的权重

四、效率损失

为了将原始BERT模型压缩为有效的压缩模型，Adabert将模型效率纳入损失函数中。主要包括参数大小和推理时间。 具体来说，对于搜索的结构αc和搜索堆叠层数K，将等式(1)中的效率感知损失定义为：

   

其中,Kmax是预定义的最大层数， oi,j是需要搜索的模型operation。SIZE(·)和FLOPs(·)是参数大小和每个操作的浮点操作(FLOP)的数量。