广告行业中那些趣事系列50：一文看懂BERT知识蒸馏发展进程

导读：本文是“数据拾光者”专栏的第五十篇文章，这个系列将介绍在广告行业中自然语言处理和推荐系统实践。本篇从理论上到实践介绍了BERT知识蒸馏发展进程，对BERT知识蒸馏感兴趣并希望应用到实际业务中的小伙伴可能有帮助。

欢迎转载，转载请注明出处以及链接，更多关于自然语言处理、推荐系统优质内容请关注如下频道。

知乎专栏：数据拾光者

公众号：数据拾光者

摘要：本篇从理论上到实践介绍了BERT知识蒸馏发展进程。首先介绍了BERT类模型提升线上性能的方法以及知识蒸馏为什么有效；然后重点介绍了BERT蒸馏的主要发展进程，主要包括微调阶段蒸馏最后一层、微调阶段蒸馏中间层、预训练阶段蒸馏、预训练+微调两阶段蒸馏以及其他蒸馏方案；最后介绍了我们实际项目中BERT蒸馏实践。对BERT知识蒸馏感兴趣并希望应用到实际业务中的小伙伴可能有帮助。

下面主要按照如下思维导图进行学习分享：

01

如何提升BERT类模型的线上性能

还是按照老规矩介绍下背景，最近主要做小X语音助手安全服务模型，从任务上来看属于文本分类任务，所以主要调研并介绍下BERT知识蒸馏的发展进程。对于文本分类任务来说算法选型主要有两大类选择：如果选择FastText或者textCNN这一类传统文本分类模型，优势在于推理速度很快，可以满足线上时延要求，但是缺点是模型效果较差；如果选择当前NLP领域很火的预训练类模型BERT，优点在于模型效果很好，但是缺点是模型非常大(基础版本的BERT参数量上亿)，部署到线上推理速度较慢很难满足时延要求。如何提升BERT等预训练模型的线上推理性能就非常重要，常见的方法主要有剪枝、因式分解、权值共享、量化和知识蒸馏等等，下面是各种方法简单介绍：

• 剪枝。对模型的网络进行修剪，比如减掉多余的头(因为Transformer使用多头注意力机制)，或者直接粗暴的使用更少的Transformer层数；

• 因式分解。之前比较火的ALBERT模型使用的一个优化策略就是对embedding参数进行因式分解。因为BERT将词向量和encode输出的维度都设置为768维，而encode中包含丰富的语义信息，所以明显存储的信息量比词向量多，所以ALBERT的策略就是采用因式分解的方法把词向量映射到低维空间，这样就能大大降低参数量，最后再映射回高维的embedding向量；

• 权值共享。这也是ALBERT中使用的优化策略之一。对Transformer各层参数可视化分析发现各层参数类似，都是在[CLS]token和对角线上分配更多的注意力，通过多层之间共享参数从而达到了模型加速的目的；

• 量化。量化操作主要是以精度换速度，业界也有尝试在BERT微调阶段进行量化感知训练，使用最小的精度损失将BERT模型参数压缩了4倍。这些量化操作方案很多也是为了将模型移植到移动端进行的优化；

• 知识蒸馏。知识蒸馏是把大模型或者多个模型ensemble学到的知识想办法迁移到一个轻量级的小模型上去，线上部署这个小模型就可以了。

上述提升模型性能的方法有各自的优缺点，本篇重点介绍实际业务中经常使用的知识蒸馏方法。

02

为什么知识蒸馏有效

如果使用BERT来构造文本分类模型，模型效果比较好，但是因为参数量大导致模型部署到线上很难满足时延要求；如果使用TextCNN这一类传统的文本分类模型虽然可以满足线上时延要求，但是效果相比于BERT来说差很多。是否有一种方法可以得到效果媲美BERT的模型，同时模型的线上推理性能很好可以满足时延的要求，也就是可以很好的兼顾性能和效果？知识蒸馏可能是一种不错的方案。下面是通过BERT知识蒸馏构建又好又快的线上模型图：

图1 通过BERT知识蒸馏构造又好又快的模型

BERT知识蒸馏的核心思路是让BERT类预训练+微调两阶段模型作为老师去教学生模型TextCNN，从而线上部署TextCNN满足线上推理时延要求，同时经过蒸馏的TextCNN学生模型效果可以媲美老师BERT模型。

知识蒸馏之所以有效主要有两个原因：第一个原因是知识蒸馏的目标是让学生模型学习老师模型的泛化能力，而不仅仅是拟合训练数据；第二个原因是知识蒸馏可以提供更多的“暗”知识。这里以手写数字任务为例，常规算法任务中只能得到这条样本的标签是1，但知识蒸馏任务中主要使用的是logits，可以知道这条样本是1的概率是0.7，是7的概率为0.2，是9的概率为0.1，相当于提供了更多的“暗”知识。下面是知识蒸馏提供更多暗知识图：

图2 知识蒸馏提供更多暗知识

03

BERT蒸馏的主要发展进程

为了在实际任务中选择合适的蒸馏方案，需要深入研究近几年BERT知识蒸馏发展进程。整体来看对于BERT的蒸馏主要经历了微调阶段蒸馏最后一层、微调阶段蒸馏中间层、预训练阶段蒸馏、预训练+精调两阶段蒸馏和其他蒸馏思路，下面会分别进行介绍。

3.1 微调阶段蒸馏最后一层

论文《Distilling Task-Specific Knowledge fromBERT into Simple Neural Networks》中提出使用业务相关数据微调BERT作为老师模型，将老师学习的知识蒸馏到学生模型中，学生模型可以是简单的神经网络比如CNN或者LSTM等。下图是以BiLSTM作为学生模型进行知识蒸馏：

图3 BiLSTM作为学生模型进行知识蒸馏

论文作者先用任务相关的训练集去微调BERT-large模型得到老师模型，然后去蒸馏老师模型的最后一层，也就是用这个老师模型得到样本的logits去训练学生模型。如上图所示，学生模型是BiLSTM，其中a是输入embedding，b是双向LSTM，c和d分别是前后向的隐层向量，e和g是全连接层，f是语句向量隐层表示，h是logits输出，i是softmax激活函数，j是经过softmax的输出概率。通过简单蒸馏微调BERT的最后一层得到的BiLSTM模型效果虽然比BERT下降不少，但是可以媲美ELMO模型。这里需要说明的是，作者通过样本增强的方式使用了更多的数据集。

对应到实际业务实践中，我们蒸馏初版使用的就是BERT-base作为老师模型，TextCNN作为学生模型，线上推理速度很快，模型效果也较好。在CPU机器上单条预测耗时平均为30ms，在GPU机器上单条预测耗时平均为5ms，可以很好的满足业务需求。这里建议小伙伴们如果实际业务中需要使用蒸馏技术，蒸馏textCNN方案实现简单，模型效果较好，线上推理性能很好，可以作为蒸馏初版方案。蒸馏过程可以使用海量的业务相关的数据集作为蒸馏语料可以有效提升模型效果。

3.2 微调阶段蒸馏中间层

直接将微调后的BERT最后一层蒸馏到简单神经网络虽然方法简单，但是也存在容易过拟合的风险。论文《Patient Knowledge Distillation for BERT Model Compression》提出了蒸馏BERT中间层的思路：将12层BERT-base作为老师模型，蒸馏中间层作为学生，这里可以蒸馏6层或者3层Transformer作为学生模型。因为Transformer模型层数越多，效果越好，但是性能越慢，所以学生模型Transformer层数的选择可以均衡线上性能和模型效果之后的结果。论文作者提出了两种层选择方案，一种是隔层选择的PKD-Skip方案，比如可以选择2/4/6/8/10/12这六层Transformer；另一种是选择最后几层的PKD-Last方案，比如可以选择最后6层Transformer进行蒸馏。下面是PKD两种层选择方案图：

图4 PKD两种层选择方案

3.3 预训练阶段蒸馏

上面是使用业务相关的训练数据微调BERT之后进行蒸馏，论文《DistilBERT, a distilled version of BERT: smaller, faster, cheaperand lighter》中提出直接在预训练阶段对BERT进行蒸馏。论文作者将预训练好的BERT-base作为老师模型，采用PKD-Skip方式蒸馏出6层的Transformer作为学生模型。DistilBERT的损失函数包括MLM loss、最后一层交叉熵和隐层之间的cosine loss三部分，其中cosine loss主要用于调整老师和学生隐层向量的方向。

3.4 预训练+微调两阶段蒸馏

因为分别在微调和预训练阶段进行蒸馏都能提升效果，那么将两者联合起来，在预训练和微调阶段都进行蒸馏效果应该会更好，所以论文《TinyBERT: Distilling BERT for Natural Language Understanding》提出了预训练和微调两阶段蒸馏方法，TinyBERT两阶段蒸馏如下图所示：

图5 TinyBERT两阶段蒸馏

TinyBERT第一阶段是预训练阶段蒸馏，将BERT-base预训练模型作为老师蒸馏得到六层Transformer的学生General TinyBERT模型；第二阶段是微调阶段蒸馏，先将任务相关数据进行样本增强，然后用样本增强之后的数据微调BERT-base得到老师模型去蒸馏得到任务相关的六层TinyBERT模型。TinyBERT还提出了一种针对Transformer的蒸馏方法，不同层对应的损失函数如下：

图6 TinyBERT针对不同层的损失函数

当m=0时对输入embedding进行蒸馏；当m=M+1时对最后一层进行蒸馏；对于中间的Transformer层蒸馏损失函数主要包括隐层loss和注意力loss两部分，因为有研究表明注意力矩阵可以获取更多的知识，所以对注意力矩阵进行了蒸馏。下面是对Transformer各层进行蒸馏的介绍：

图7 对Transformer各层进行蒸馏

TinyBERT作者通过消融实验对比GD(预训练阶段蒸馏)、TD(微调阶段蒸馏)和DA(数据增强)三个流程对模型效果的影响，其中TD和DA影响显著，GD作用最小，下面是实验结果：

图8 GD、TD和DA不同阶段对TinyBERT效果的影响

作者还通过消融实验对比了不同目标函数对模型效果的影响，可以看出中间Transformer层目标函数蒸馏对模型效果影响最大，而Embedding层和最后一层影响最小，下面是实验结果：

图9 不同目标函数对蒸馏效果的影响

对应到实际业务实践中，蒸馏textCNN的优化版本就是TinyBERT方案。相比于蒸馏textCNN来说，TinyBERT效果会提升很多，但是相对来说实现也比较复杂。至于模型效果和推理速度的选择可以根据实际的需求确定蒸馏Transformer的层数。

3.5 其他蒸馏方案

其他的蒸馏方案还有通过降低每层Transformer的维度进行蒸馏。论文《MobileBERT: a Compact Task-Agnostic BERT for Resource-LimitedDevices》中提出了保留24层Transformer，通过降低每层Transformer的维度来减少参数量，提升线上推理速度。下图是MobileBERT的主要思想bottleneck机制介绍：

图10 MobileBERT的bottleneck机制

上图中a部分是常规BERT结构图，b和c部分是加入bottleneck机制的BERT-large结构图，主要区别是在Transformer前后添加了Linear线性层。b作为老师模型将embedding的维度从512扩大到1024，c作为学生模型将embedding的维度从512缩小到128，从而降低模型的参数上，提升线上性能。

04

项目实践

我们小X语音助手安全服务文本分类模型项目实践主要经过了以下几个阶段：

• baseline：使用fasttext或者textCNN模型在训练集上直接训练得到baseline模型；

• v1：使用训练集基于BERT去训练模型，然后将训练得到的模型去得到更多伪标签样本添加到训练集中，用扩展之后的训练集再去训练fasttext或者textCNN模型，相当于从样本层面优化模型；

• v2：将BERT模型作为老师模型蒸馏textCNN；

• v3：将BERT模型作为老师模型蒸馏tinyBERT。

05

总结及反思

本篇从理论上到实践介绍了BERT知识蒸馏发展进程。首先介绍了BERT类模型提升线上性能的方法以及知识蒸馏为什么有效；然后重点介绍了BERT蒸馏的主要发展进程，主要包括微调阶段蒸馏最后一层、微调阶段蒸馏中间层、预训练阶段蒸馏、预训练+微调两阶段蒸馏以及其他蒸馏方案；最后介绍了我们实际项目中BERT蒸馏实践。对BERT知识蒸馏感兴趣并希望应用到实际业务中的小伙伴可能有帮助。

最新最全的文章请关注我的微信公众号或者知乎专栏：数据拾光者。

码字不易，欢迎小伙伴们点赞和分享。