TinyBERT: Distilling BERT for Natural Language Understanding翻译

摘要

语言模型预训练（例如BERT）已大大改善了许多自然语言处理任务的性能。但是，预训练的语言模型通常在计算上昂贵，因此很难在资源受限的设备上有效执行它们。为了加快推理速度并减小模型大小且同时保持精度，我们首先提出了一种新的Transformer蒸馏方法，该方法专门为基于Transformer模型的知识蒸馏（KD）设计。通过利用这种新的KD方法，可以将大型“teacher” BERT中编码的大量知识有效地转移给小型“student” TinyBERT。然后，我们为TinyBERT引入了一个新的两阶段学习框架，该框架在预训练阶段和特定于任务的学习阶段都执行Transformer蒸馏。该框架确保TinyBERT可以捕获BERT中的通用领域以及特定于任务的知识。

1.介绍

预训练语言模型，然后对下游任务进行微调，已成为自然语言处理（NLP）的新范例。预训练的语言模型（PLM），例如BERT，XLNet，RoBERTa，ALBERT ，T5和ELECTRA在许多NLP任务（例如GLUE基准测试和具有挑战性的多步推理任务）中都取得了巨大的成功。但是，PLM通常具有大量参数，并且推理时间长，难以部署在移动电话等边缘设备上。最近的研究证明PLM中存在冗余。因此，在保持性能的同时减少PLM的计算开销和模型存储至关重要且可行。
　　许多模型压缩技术已经被提出，以加速深度模型推理并减小模型大小，同时保持准确性。最常用的技术包括量化，权重修剪和知识蒸馏（KD）。在本文中，我们专注于知识蒸馏，这是一个由Hinton et al. (2015)提出的想法，在teacher-student框架中。KD旨在将嵌入在大型teacher网络中的知识转移到小型student网络，在该网络中训练student网络以重现teacher网络的行为。在此框架的基础上，我们针对基于transformer的模型提出了一种新的蒸馏方法，并以BERT为例研究了大型PLM的KD方法。
　　KD已经在NLP和预训练的语言模型中进行了广泛的研究。预训练然后微调范例首先在大型无监督文本语料库上对BERT进行预训练，然后在特定于任务的数据集上对其进行微调，这极大地增加了BERT蒸馏的难度。因此，需要为两个训练阶段设计有效的KD策略。

　　为了构建具有竞争力的TinyBERT，我们首先提出一种新的Transformer蒸馏方法，以提取teacher BERT中嵌入的知识。具体来说，我们设计三种类型的损失函数以适应BERT不同层的表示形式：1）嵌入层的输出；2）源自Transformer层的隐藏状态和注意力矩阵；3）预测层输出的logits。基于注意力的拟合是受最近的发现的启发而来的，该发现表明BERT学习的注意力权重可以捕获大量的语言知识，从而鼓励语言知识可以很好地从teacher BERT转移到student TinyBERT。然后，我们提出了一个新的两阶段蒸馏框架，包括通用蒸馏和特定任务的蒸馏，如图1所示。在通用蒸馏阶段，没有微调的原始BERT充当teacher模型。student TinyBERT通过使用通用域语料库中的Transformer蒸馏来模仿老师的行为。之后，我们获得了一个通用的TinyBERT，该TinyBERT用作student模型的初始化以进行进一步的蒸馏。在特定任务的蒸馏阶段，我们首先进行数据增强，然后使用微调的BERT作为teacher模型在增强的数据集上进行蒸馏。应该指出的是，这两个阶段对于提高TinyBERT的性能和泛化能力都是必不可少的。
　　这项工作的主要贡献如下：1）我们提出了一种新的Transformer蒸馏方法，以鼓励将以teacher BERT编码的语言知识适当地转移给TinyBERT；2）我们提出了一种新的两阶段学习框架，该框架可以在预训练和微调阶段执行的transformer蒸馏，从而确保TinyBERT可以吸收teacher BERT的一般领域知识和特定任务知识。3）我们在实验中证明，我们的4层TinyBERT可以在GLUE任务上实现达到teacher BERT-BASE96.8％的性能，而参数（〜13.3％）和推理时间（〜10.6％）却少得多，并且明显优于其他最先进的基线；4）我们还展示了6层TinyBERT可以与GLUE上的teacher BERT-BASE相当。

2.预备知识

在本节中，我们描述了Transformer和Knowledge Distillation的方式。我们提出的transformer蒸馏是针对基于Transformer模型的一种特殊设计的KD方法。

2.1 Transformer层

最近的大多数预训练语言模型（例如BERT，XLNet和RoBERTa）都是使用Transformer层构建的，该层可以通过自注意力机制捕获输入字符之间的长距离依赖关系。具体来说，标准的Transformer层包括两个主要子层：多头注意力（MHA）和全连接的前馈网络（FFN）。
　　（1）多头注意力（MHA）
　　注意力函数的计算取决于queries，keys和values三个组成部分，分别表示为矩阵$\text{Q}$，$\text{K}$和$\text{V}$。注意力函数可以表述为：
$$\text{A}=\frac{\text{Q}{\text{K}}^T}{\sqrt{d_k}},\text{\hspace{1em}(1)}$$
$$Attention(\text{Q},\text{K},\text{V})=softmax(\text{A})\text{V},\text{\hspace{1em}(2)}$$
其中$d_k$是keys的维数并用作缩放因子，$\text{A}$是根据点积运算从$\text{Q}$和$\text{K}$的计算出的注意力矩阵。最终函数输出被计算为值$\text{V}$的加权和，权重是通过对矩阵$\text{A}$的每一列应用$softmax()$运算来计算的。BERT中的注意力矩阵可以捕获大量的语言知识，因此在我们提出的蒸馏方法中起着至关重要的作用。
　　通过将来自不同表示子空间的注意力头进行串联来定义多头注意力，如下所示：
$$MHA(\text{Q},\text{K},\text{V})=Concat(h_1,...,h_k)\text{W},\text{\hspace{1em}(3)}$$
其中$k$是注意力头的数量，而$h_i$表示第$i$个注意力头，它由$Attention()$函数使用来自不同表示子空间的输入计算得出。矩阵$\text{W}$充当线性变换。
　　（2）全连接的前馈网络（FFN）
　　Transformer层还包含一个全连接的前馈网络，其公式如下：
$$FFN(x)=max(0,x{\text{W}}_1+b_1){\text{W}}_2+b_2.\text{\hspace{1em}(4)}$$
我们可以看到，FFN包含两个线性变换和一个ReLU激活。

2.2 知识蒸馏

KD旨在通过训练student网络以模仿teacher网络的行为，从而将大型teacher网络$T$的知识转移到小型student网络$S$。令$f^T$和$f^S$分别代表teacher和student网络的行为函数。行为函数的目标是将网络输入转换为一些有用的表示，可以将其定义为网络中任何层的输出。在Transformer蒸馏的背景下，MHA层或FFN层的输出或某些中间表示（例如注意力矩阵$\text{A}$）可用作行为函数。正式地，可以将KD建模为最小化以下目标函数：
$${\mathcal{L}}_{KD}=\sum _{x\in \mathcal{X}}L(f^S(x),f^T(x)),\text{\hspace{1em}(5)}$$
其中$L(\cdot )$是一个损失函数，用于评估teache和studen网络之间的差异，$x$是文本输入，$\mathcal{X}$表示训练数据集。因此，研究问题的关键成为如何定义有效的行为函数和损失函数。与以前的KD方法不同，除了特定于任务的训练阶段之外，我们还需要考虑如何在BERT的预训练阶段执行KD。

3.方法

在本节中，我们为基于Transformer的模型提出一种新的蒸馏方法，并为从BERT提取的模型提出一个两阶段学习框架，这称为TinyBERT。

3.1 Transformer蒸馏

　　提出的Transformer蒸馏是一种专门设计用于Transformer网络的KD方法。在这项工作中，student和teacher网络都是使用Transformer层构建的。为了清楚起见，我们在介绍我们的方法之前先阐明问题。
　　（1）问题定义
　　假设student模型具有$M$个Transformer层，teacher模型具有$N$个Transformer层，我们首先从teacher模型的$N$个层中选择$M$个，以进行Transformer层蒸馏。接下来将函数$n=g(m)$定义为从student层到teacher层的索引之间的映射函数，这意味着第$m$层srudent模型从第$g(m)$层teacher模型中学习信息。准确地说，我们将$0$设为嵌入层的索引，将$M+1$设为预测层的索引，并将相应的层映射分别定义为$0=g(0)$和$N+1=g(M+1)$。在实验部分研究了选择不同映射函数对性能的影响。正式地，student可以通过最小化以下目标从teacher那里获得知识：
$${\mathcal{L}}_{model}=\sum _{x\in \mathcal{X}}\sum _{m=0}^{M+1}{\lambda }_m{\mathcal{L}}_{layer}(f_m^S(x),f_{g(m)}^T(x)),\text{\hspace{1em}(6)}$$
其中${\mathcal{L}}_{layer}$表示指定模型层（例如，Transformer层或嵌入层）的损失函数，$f_m(x)$表示从第$m$层得到的行为函数，而${\lambda }_m$是表示对第$m$层进行蒸馏的重要性超参数。
　　（2）Transformer-layer蒸馏
　　提出的Transformer层蒸馏包括基于注意力的蒸馏和基于隐藏状态的蒸馏，如图2所示。基于注意力的蒸馏是受最近发现的启发而来的，该发现是BERT学习的注意力权重可以捕获丰富的语言知识。这种语言知识包括语法和共指信息，这对于自然语言理解至关重要。因此，我们提出了基于注意力的蒸馏，以鼓励语言知识可以从teacher（BERT）转移到student（TinyBERT）。具体来说，student学习在teacher网络中拟合多头注意力矩阵，目标定义为：
$${\mathcal{L}}_{attn}=\frac{1}{h}\sum _{i=1}^{h}MES({\text{A}}_i^S,{\text{A}}_i^T),\text{\hspace{1em}(7)}$$
其中$h$是注意力头的数量，${\text{A}}_i\in {\mathbb{R}}^{l\times l}$表示对应于第$i$个teacher或student头的注意力矩阵，$l$是输入文本的长度，$MSE()$表示均方误差损失函数。在这项工作中，（为归一化的）注意力矩阵${\text{A}}_i$代替了它的softmax输出$softmax({\text{A}}_i)$被用作拟合目标，因为我们的实验表明，前一种设置具有更快的收敛速度和更好的性能。
　　除了基于注意力的蒸馏外，我们还从Transformer层的输出中提取知识，目标函数如下：
$${\mathcal{L}}_{hidn}=MSE({\text{H}}^S{\text{W}}_h,{\text{H}}^T),\text{\hspace{1em}(8)}$$
其中矩阵${\text{H}}^S\in {\mathbb{R}}^{l\times d^{\prime }}$和${\text{H}}^T\in {\mathcal{R}}^{l\times d}$分别指代student和teacher网络的隐藏状态，这些状态由等式4计算得出。标量值$d$和$d^{\prime }$表示teacher和student模型的隐藏维度，并且$d^{\prime }$通常小于$d$，以获得较小的student网络。矩阵${\text{W}}_h\in {\mathbb{R}}^{d^{\prime }\times d}$是可学习的线性变换，它将student网络的隐藏状态转换到与teacher网络隐藏状态维度相同的空间。
　　（3）Embedding-layer蒸馏
　　类似于基于隐藏状态的蒸馏，我们还执行嵌入层蒸馏，目标函数是：
$${\mathcal{L}}_{embd}=MSE({\text{E}}^S{\text{W}}_e,{\text{E}}^T),\text{\hspace{1em}(9)}$$
其中矩阵${\text{E}}^S$和${\text{E}}^T$分别指的是student和teacher网络的嵌入。在本文中，它们具有与隐藏状态矩阵相同的维度。矩阵${\text{W}}_e$是线性变换，起着与${\text{W}}_h$相似的作用。
　　（4）Prediction-layer蒸馏
　　除了模仿中间层的行为外，我们还使用知识蒸馏来拟合teacher模型的预测。具体来说，我们对student网络的logits与教师的logits之间的软交叉熵损失进行了惩罚：
$${\mathcal{L}}_{pred}=CE(z^T/t,z^S/t),\text{\hspace{1em}(10)}$$
其中$z^S$和$z^T$分别是student和teacher预测的logits矢量，$CE$表示交叉熵损失，$t$表示温度值。在我们的实验中，我们发现$t=1$表现最好。
　　使用以上蒸馏目标函数（即方程式7、8、9和10），我们可以统一student和teacher网络之间相应层的蒸馏损失：
$${\mathcal{L}}_{layer}=\{\begin{array}{ll}{\mathcal{L}}_{embed},& m=0\\ {\mathcal{L}}_{hidn}+{\mathcal{L}}_{attn},& M\ge m>0\\ {\mathcal{L}}_{pred},& m=M+1\end{array}\text{\hspace{1em}(11)}$$

3.2 TinyBERT学习

　　BERT的应用通常包括两个学习阶段：预训练和微调。BERT在预训练阶段学到的大量知识非常重要，也应将其转移到压缩模型中。因此，我们提出了一个新的两阶段学习框架，包括通用蒸馏和特定于任务的蒸馏，如图1所示。通用蒸馏有助于TinyBERT学习预训练的BERT中嵌入的丰富知识，这对于改进TinyBERT的泛化能力具有重要作用。特定任务的蒸馏进一步向TinyBERT传授了来自微调BERT的知识。通过两步蒸馏，我们可以大大减少teacher和student模型之间的差距。
　　（1）General Distillation
　　我们使用原始的BERT（无需微调）作为teacher，并使用大型文本语料库作为训练数据。通过对来自常规领域的文本执行Transformer蒸馏（此时不进行预测蒸馏），我们获得了可以针对下游任务进行微调的常规TinyBERT。但是，由于隐藏/嵌入大小和层数的显着减少，一般TinyBERT的性能通常比BERT差。
　　（2）Task-specific Distillation
　　先前的研究表明，复杂的模型（例如微调的BERT）会因特定领域任务的过度参数化而性能变差。因此，较小的模型有可能达到与BERT相当的性能。为此，我们提出通过特定任务的蒸馏来生产具有竞争力的微调TinyBERT。在任务特定的蒸馏中，我们在增强的任务特定的数据集上重新执行提出的Transformer蒸馏。具体来说，将微调的BERT用作teacher，并提出一种数据增强方法来扩展特定于任务的训练集。通过训练更多与任务相关的示例，可以进一步提高学生模型的泛化能力。
　　（3）Data Augmentation
　　我们结合了预训练的语言模型BERT和GloVe的单词嵌入来进行单词级替换，以进行数据增强。具体来说，我们使用语言模型来预测单片单词的单词替换，并使用单词嵌入来检索最相似的单词作为多片单词的单词替换。我们还定义了一些超参数来控制句子的替换率和扩充数据集的数量。数据增强过程的更多详细信息显示在算法1中。我们为所有实验都设置了$p_t=0.4,N_a=20,K=15$。
　　以上两个学习阶段是相辅相成的：通用蒸馏为特定任务的蒸馏提供了良好的初始化，而扩展数据上的特定任务的蒸馏通过专注于特定于任务的知识进一步改善了TinyBERT。尽管模型规模显着减小，但通过数据增强和通过在预训练和微调阶段执行提出的Transformer蒸馏方法，TinyBERT可以在各种NLP任务中取得有竞争力的性能。