预训练模型蒸馏 -- 宽度自适应策略蒸馏TinyBERT

预训练模型蒸馏

在前面的课程中，大家了解了自然语言处理领域中一些经典的模型，比如BERT、ERNIE等，它们在NLP任务中的强大之处是毫无疑问的，但由于预训练模型的参数较多，体积庞大，在部署时对设备的运算速度和内存大小以及能耗都有着极高的要求。但当我们处理实际的产业应用需求时，比如将深度学习模型部署到手机上时，就需要对模型进行压缩，在不影响性能的前提下使其变得体积更小、速度更快、能耗更低。本节课我们会先对预训练模型蒸馏中几个比较经典的模型进行介绍，比如：Patient-KD、DistilBERT、TinyBERT和DynaBERT。从原理和结构上对以上几个模型进行详解。然后再通过一个实验案例，带领大家使用DynaBERT训练策略中宽度自适应部分来对TinyBERT在GLUE基准数据集的QQP任务中进行蒸馏，以此进行实际效果验证。

资源

• 更多CV和NLP中的transformer模型(BERT、ERNIE、ViT、DeiT、Swin Transformer等)、深度学习资料，请参考：awesome-DeepLearning

• 了解并使用更多模型压缩相关工具，请参考：PaddleSlim

模型压缩简介

模型压缩方法主要可以分为以下四类：

• 参数修剪和量化（Parameter pruning and quantization）：用于消除对模型表现影响不大的冗余参数。早期工作表明，网络修剪和量化在降低网络复杂性和解决过拟合问题上是有效的。它可以为神经网络带来正则化效果从而提高泛化能力。参数修剪和量化可以进一步分为三类：量化和二值化，网络剪枝和结构化矩阵。量化可以看作是“量子级别的减肥”，神经网络模型的参数一般都用float32的数据表示，但如果我们将float32的数据计算精度变成int8的计算精度，则可以牺牲一点模型精度来换取更快的计算速度。而剪枝则类似“化学结构式的减肥”，将模型结构中对预测结果不重要的网络结构剪裁掉，使网络结构变得更加 ”瘦身“。比如，在每层网络，有些神经元节点的权重非常小，对模型加载信息的影响微乎其微。如果将这些权重较小的神经元删除，则既能保证模型精度不受大影响，又能减小模型大小。结构化矩阵则是用少于 $m \times n$ 个参数来描述一个 $m \times n$ 阶矩阵，以此来减少内存消耗。

• 低秩分解（Low-rank factorization）：卷积神经网络中的主要计算量在于卷积计算，而卷积计算本质上是矩阵分析问题，因此可以通过对多维矩阵进行分解的方式，用多个低秩矩阵来逼近该矩阵，比如将一个3D卷积转换为3个1D卷积，从而降低参数复杂度和运算复杂度。

• 迁移/压缩卷积滤波器（Transferred/compact convolutional filters）：通过构造特殊结构的卷积滤波器来降低存储空间、减小计算复杂度。

• 知识蒸馏（Knowledge distillation）：类似“老师教学生”，使用一个效果好的大模型指导一个小模型训练，因为大模型可以提供更多的软分类信息量，所以会训练出一个效果接近大模型的小模型。

在本节课中，我们主要讲述以知识蒸馏的方法对BERT（transformer-based）模型进行压缩。

知识蒸馏

2014年，Geoffrey Hinton在 Distilling the Knowledge in a Neural Network 中提出知识蒸馏（KD）概念：把从一个复杂的大模型（Teacher Network）上学习到的知识迁移到另一个更适合部署的小模型上（Student Network），叫知识蒸馏。

​

知识蒸馏结构

​

如上图所示，左边的教师网络是一个复杂的大模型，以它带有温度参数T的softmax输出作为软目标作为学生网络学习的软目标。学生网络在学习时，也通过带有温度参数T的softmax进行概率分布预测，与软目标计算soft loss。同时，也通过正常的训练流程获得预测的样本类别与真实的样本类别计算hard loss。最终根据 $\gamma * soft loss + (1 - \gamma) * hard loss$ 作为损失函数来训练学生网络。

其中，知识蒸馏过程中涉及到的两种标签分别是：

• 硬标签（hard target）：网络训练的目标，即分类任务中正确分类的label，正标签为1，其余标签都为0；

• 软标签（soft target）：大模型的softmax层输出的类别概率，正标签的概率最高。

论文中提出的softmax函数中增加了温度(Temperature)这个参数，其公式如下：

其中，T代表温度。原始的softmax函数就是 $T=1$ 的情况。那么，为什么要增加一个变量T呢？一个复杂高性能网络往往可以获得很好的分类效果，也就是说预测错误的概率会比预测正确的概率小得多。但是一个小的网络，可能很难达到复杂的大网络的性能。T参数的加入就是为了帮助小网络可以更好的学习到大网络对负标签的判断信息。在传统的训练过程中（仅使用hard target），所有负标签都是零，被全部统一对待。但是加入T，T越高，softmax的类概率分布会越平滑，就可以将大模型对每一个标签的学习信息都传递给小模型。

举一个例子，假设一个三分类问题，z_0 = 2.2, z_1 = 0.8, z_2 = 0.2，那么当 T = 1 时，

当 T = 10时，

可以看到，T越高，softmax的类分布概率会变得越平滑。这就使得学生网络可以学习到教师网络对负标签归纳的信息。

​

手写数字识别任务（图来源于参考文献3）

​

举个例子，如上图所示，在用MNIST数据集做手写数字识别任务时，某个输入的“2”更加类似“3”，则softmax的输出值中“3”对应的概率应该要比其他负标签类别高；而另一个“2”更类似于“7”，则这个这个样本的softmax输出值中“7”对应的概率应该比其他负标签类别高。这两个“2”对应的hard target是相同的，但是他们的soft target是不同的，soft target内蕴含着更多的信息。

Patient-KD

Patient-KD 算法综述

论文地址：Patient Knowledge Distillation for BERT Model Compression

​

图1: Vanilla KD和PKD比较

​

BERT预训练模型对资源的高需求导致其很难被应用在实际问题中，为缓解这个问题，论文中提出了Patient Knowledge Distillation（Patient KD）方法，将原始大模型压缩为同等有效的轻量级浅层网络。同时，作者对以往的知识蒸馏方法进行了调研，如图1所示，vanilla KD在QNLI和MNLI的训练集上可以很快的达到和teacher model相媲美的性能，但在测试集上则很快达到饱和。对此，作者提出一种假设，在知识蒸馏的过程中过拟合会导致泛化能力不良。为缓解这个问题，论文中提出一种“耐心”师生机制，即让Patient-KD中的学生模型从教师网络的多个中间层进行知识提取，而不是只从教师网络的最后一层输出中学习，该学习方法遵循以下两个策略：

1. PKD-Skip: 从每k层学习，假设教师网络的底层和高层中都包含重要信息，需要被学习到（如图2a所示）

2. PKD-Last: 从最后k层学习，假设教师网络越靠后的层包含越丰富的知识信息（如图2b所示）

​

图2a: PKD-Skip 学生网络学习教师网络每两层的输出 图2b: PKD-Last 学生网络从教师网络的最后六层学习

​

因为在BERT中仅使用最后一层的[CLS] token的输出来进行预测，且在其他BERT的变体模型中，如SDNet，是通过对每一层的[CLS] embedding的加权平均值进行处理并预测。由此可以推断，如果学生模型可以从任何教师网络中间层中的[CLS]表示中学习，那么它就有可能获得类似教师网络的泛化能力。

因此，Patient-KD中提出特殊的一种损失函数的计算方式：

其中，对于输入$x_i$，所有层的[CLS]的输出表示为： $h_i = [h{i,1}, h{i,2},..., h{i,k}] = BERT{k}(x_i) \in \mathbb{R}^{k\times d}$

$I{pt}$表示要从中提取知识的一组中间层，以从$BERT{12}$压缩到$BERT_6$为例，对于PKD-Skip策略，$I{pt} = {2,4,6,8,10}$；对于PKD-Last策略，$I{pt} = {7,8,9,10,11}$。M表示学生网络的层数，N是训练样本的数量，上标s和t分别代表学生网络和教师网络。

同时，Patient-KD中也使用了$L{DS}$和$L{CE}^S$两种损失函数用来衡量教师和学生网络的预测值的距离和学生网络在特定下游任务上的交叉熵损失。

最终的目标损失函数可以表示为：

实验结果

​

图3: results from the GLUE test server

作者将模型预测提交到GLUE并获得了在测试集上的结果，如图3所示。与fine-tuning和vanilla KD这两种方法相比，使用PKD训练的$BERT_3$和$BERT_6$在除MRPC外的几乎所有任务上都表现良好。其中，PKD代表Patient-KD-Skip方法。对于MNLI-m和MNLI-mm，六层模型比微调（FT）基线提高了1.1%和1.3%，

​

图4: PKD-Last 和 PKD-Skip 在GLUE基准上的对比

​

尽管这两种策略都比vanilla KD有所改进，但PKD-Skip的表现略好于PKD-Last。作者推测，这可能是由于每k层的信息提炼捕获了从低级到高级的语义，具备更丰富的内容和更多不同的表示，而只关注最后k层往往会捕获相对同质的语义信息。

​

图5: 参数量和推理时间对比

​

图5展示了$BERT_3$、$BERT_6$、$BERT_12$的推理时间即参数量, 实验表明Patient-KD方法实现了几乎线性的加速，$BERT_6$和$BERT_3$分别提速1.94倍和3.73倍。

DistilBERT

DistilBERT 算法综述

论文地址：DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

​

图1: 几个预训练模型的参数量统计

​

近年来，大规模预训练语言模型成为NLP任务的基本工具，虽然这些模型带来了显著的改进，但它们通常拥有数亿个参数（如图1所示），而这会引起两个问题。首先，大型预训练模型需要的计算成本很高。其次，预训练模型不断增长的计算和内存需求可能会阻碍语言处理应用的广泛落地。因此，作者提出DistilBERT，它表明小模型可以通过知识蒸馏从大模型中学习，并可以在许多下游任务中达到与大模型相似的性能，从而使其在推理时更轻、更快。

学生网络结构

学生网络DistilBERT具有与BERT相同的通用结构，但token-type embedding和pooler层被移除，层数减半。学生网络通过从教师网络中每两层抽取一层来进行初始化。

Training loss

L_{ce} 训练学生模仿教师模型的输出分布：

其中，t_i和s_i分别是教师网络和学生网络的预测概率。

同时使用了Hinton在2015年提出的softmax-temperature：

<