TinyBERT: Distilling BERT for Natural Language Understanding——用于自然语言理解的BERT压缩模型

这篇文章的核心内容是关于如何高效地压缩预训练语言模型BERT，以使其更适合在资源受限的设备上运行，同时保持较高的性能。文章提出了一个名为TinyBERT的模型，通过知识蒸馏（Knowledge Distillation, KD）技术将BERT的知识转移到一个更小的模型中。以下是文章的主要研究内容总结：

1. 研究背景

BERT等预训练语言模型在自然语言处理（NLP）任务中表现出色，但其庞大的模型参数和高计算成本限制了其在边缘设备（如移动电话）上的应用。
为了在保持性能的同时减少模型大小和计算成本，文章提出了一种基于知识蒸馏的方法，将BERT的知识转移到一个更小的模型TinyBERT中。

2. TinyBERT的两阶段学习框架

通用蒸馏（General Distillation）：使用未经微调的BERT作为教师模型，对大规模文本语料库进行蒸馏，以提取BERT中的通用知识。这一阶段的目标是为TinyBERT提供一个良好的初始化。
特定任务蒸馏（Task-specific Distillation）：在微调后的BERT的指导下，对特定任务的数据集进行蒸馏，以进一步提升TinyBERT在特定任务上的性能。这一阶段还包括数据增强，通过生成更多与任务相关的示例来提高模型的泛化能力。

3. Transformer蒸馏方法

提出了一种新的Transformer蒸馏方法，通过设计三种类型的损失函数（嵌入层、Transformer层和预测层）来确保知识从教师模型有效地转移到学生模型。
Transformer层蒸馏包括基于注意力的蒸馏和基于隐藏状态的蒸馏，这两种方法相互补充，能够更好地捕获BERT中的语言知识。

4. 实验结果

在GLUE基准测试的多个任务上，TinyBERT4（4层模型）和TinyBERT6（6层模型）均表现出色，显著优于其他基线模型，如BERTTINY、BERT-PKD、DistilBERT等。
TinyBERT4在模型大小和推理速度上分别比BERTBASE小7.5倍和快9.4倍，同时在性能上保持了96.8%的教师模型水平。
TinyBERT6在性能上与BERTBASE相当，进一步证明了该方法的有效性。

5. 消融研究

通过移除或修改学习框架中的关键步骤（如通用蒸馏、特定任务蒸馏和数据增强），验证了每个步骤对TinyBERT性能的贡献。
实验表明，所有步骤都对TinyBERT的性能至关重要，尤其是特定任务蒸馏和数据增强对模型的泛化能力提升更为显著。

6. 映射函数的影响

研究了不同层映射策略对TinyBERT性能的影响，发现均匀策略（从BERT的不同层均匀提取知识）在多个任务上表现最佳。

7. 未来工作

探索如何将更宽和更深的教师模型（如BERTLARGE）的知识更有效地转移到TinyBERT中。
结合量化和剪枝等技术，进一步压缩预训练语言模型。

文章通过提出一种新的知识蒸馏方法和两阶段学习框架，成功地将BERT的知识转移到TinyBERT中，显著减少了模型大小和推理时间，同时保持了较高的性能。这种方法为在资源受限的设备上部署BERT模型提供了一种有效的解决方案。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

摘要

语言模型预训练（如BERT）显著提升了众多自然语言处理任务的性能。然而，预训练语言模型通常计算成本高昂，难以高效地在资源受限的设备上执行。为了在保持准确率的同时加速推理并减小模型大小，我们首先提出了一种专为Transformer基础模型的知识蒸馏（KD）设计的新型Transformer蒸馏方法。通过利用这种新的KD方法，大型"教师"BERT中编码的丰富知识可以有效地转移到小型"学生"TinyBERT中。接着，我们为TinyBERT引入了一种新的两阶段学习框架，在预训练和特定任务学习阶段均执行Transformer蒸馏。该框架确保TinyBERT能够捕获BERT中的通用领域知识以及特定任务知识。

具有4层的TinyBERT4在经验上表现出色，在GLUE基准测试中达到了其教师BERTBASE性能的96.8%以上，同时模型大小缩小了7.5倍，推理速度加快了9.4倍。TinyBERT4还显著优于具有4层的BERT蒸馏最新基线模型，仅包含其约28%的参数和约31%的推理时间。此外，具有6层的TinyBERT6在性能上与教师BERTBASE相当。

1 引言

预训练语言模型然后在下游任务上进行微调，已成为自然语言处理（NLP）的新范式。预训练语言模型（PLMs），如BERT（Devlin等人，2019年）、XLNet（Yang等人，2019年）、RoBERTa（Liu等人，2019年）、ALBERT（Lan等人，2020年）、T5（Raffel等人，2019年）和ELECTRA（Clark等人，2020年），在许多NLP任务（例如GLUE基准测试（Wang等人，2018年）和具有挑战性的多跳推理任务（Ding等人，2019年））中取得了巨大成功。然而，PLMs通常参数众多，推理时间长，难以部署在移动电话等边缘设备上。最近的研究（Kovaleva等人，2019年；Michel等人，2019年；Voita等人，2019年）表明PLMs存在冗余。因此，减少PLMs的计算开销和模型存储，同时保留其性能是至关重要且可行的。

为了加速深度模型推理并减小模型大小，同时保持准确率，已经提出了许多模型压缩技术（Han等人，2016年）。最常用的技术包括量化（Gong等人，2014年）、权重剪枝（Han等人，2015年）和知识蒸馏（KD）（Romero等人，2014年）。在本文中，我们专注于知识蒸馏，这一思想源自Hinton等人（2015年），在教师-学生框架中实现。KD旨在将嵌入在大型教师网络中的知识转移到小型学生网络中，学生网络被训练以重现教师网络的行为。基于该框架，我们提出了一种专为基于Transformer的模型设计的新型蒸馏方法，并以BERT为例研究了针对大规模PLMs的方法。

KD已在NLP（Kim和Rush，2016年；Hu等人，2018年）以及预训练语言模型（Sanh等人，2019年；Sun等人，2019年，2020年；Wang等人，2020年）中得到了广泛研究。预训练-然后-微调范式首先在大规模无监督文本语料库上预训练BERT，然后在特定任务数据集上进行微调，这大大增加了BERT蒸馏的难度。因此，需要设计一种有效的KD策略，用于预训练和微调阶段。为了构建具有竞争力的TinyBERT，我们首先提出了一种新的Transformer蒸馏方法，以提取教师BERT中嵌入的知识。具体而言，我们设计了三种类型的损失函数，以拟合BERT层的不同表示：1）嵌入层的输出；2）来自Transformer层的隐藏状态和注意力矩阵；3）预测层输出的logits。基于注意力的拟合受到最近发现的启发（Clark等人，2019年），即BERT学习到的注意力权重可以捕获大量的语言知识，从而鼓励将语言知识从教师BERT转移到学生TinyBERT。然后，我们提出了一个包括通用蒸馏和特定任务蒸馏的新型两阶段学习框架，如图1所示。在通用蒸馏阶段，未经微调的原始BERT作为教师模型。学生TinyBERT通过在通用领域语料库上进行提出的Transformer蒸馏来模仿教师的行为。之后，我们获得了一个通用TinyBERT，用作进一步蒸馏的学生模型的初始化。在特定任务蒸馏阶段，我们首先进行数据增强，然后在增强的数据集上使用微调后的BERT作为教师模型进行蒸馏。需要指出的是，这两个阶段对于提升TinyBERT的性能和泛化能力都是至关重要的。

本文的主要贡献如下：1）我们提出了一种新的Transformer蒸馏方法，以充分将教师BERT中编码的语言知识转移到TinyBERT中；2）我们提出了一个在预训练和微调阶段均执行提出的Transformer蒸馏的新型两阶段学习框架，确保TinyBERT能够吸收教师BERT的通用领域和特定任务知识；3）我们在实验中表明，我们的TinyBERT4在GLUE任务上达到了教师BERTBASE性能的96.8%以上，同时参数数量（约13.3%）和推理时间（约10.6%）都大幅减少，并且显著优于具有4层的BERT蒸馏最新基线模型；4）我们还表明，具有6层的TinyBERT6在GLUE上与教师BERTBASE的性能相当。

2 预备知识

在本节中，我们描述了Transformer（Vaswani等人，2017年）和知识蒸馏（Hinton等人，2015年）的公式化。我们提出的Transformer蒸馏是一种专为基于Transformer的模型设计的KD方法。

2.1 Transformer层

最近的大多数预训练语言模型（例如BERT、XLNet和RoBERTa）都是基于Transformer层构建的，这些层可以通过自注意力机制捕获输入标记之间的长期依赖关系。具体而言，标准的Transformer层包括两个主要子层：多头注意力（MHA）和全连接前馈网络（FFN）。

多头注意力（MHA）

注意力函数的计算依赖于查询、键和值三个组成部分，分别表示为矩阵Q、K和V。注意力函数可以表示为：

逐位置前馈网络（FFN）

Transformer层还包含一个全连接的前馈网络，其公式为：

可以看出，FFN包含两个线性变换和一个ReLU激活函数。

2.2 知识蒸馏

3 方法

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

3.1 Transformer蒸馏

我们提出的Transformer蒸馏是一种专为Transformer网络设计的KD方法。在本工作中，学生和教师网络均基于Transformer层构建。为清晰起见，我们在介绍方法之前先对问题进行公式化。

问题公式化

假设学生模型有M个Transformer层，教师模型有N个Transformer层，我们首先从教师模型的N层中选择M层进行Transformer层蒸馏。然后定义一个函数n=g(m)，作为学生层索引与教师层索引之间的映射函数，这意味着学生模型的第m层从教师模型的第g(m)层学习信息。具体而言，我们将嵌入层的索引设为0，预测层的索引设为M+1，相应的层映射分别定义为0=g(0)和N+1=g(M+1)。实验部分研究了不同映射函数选择对性能的影响。正式地，学生可以通过最小化以下目标来从教师那里获取知识：

Transformer层蒸馏

我们提出的Transformer层蒸馏包括基于注意力的蒸馏和基于隐藏状态的蒸馏，如图2所示。基于注意力的蒸馏是受到最近发现的启发，即BERT学习到的注意力权重可以捕获丰富的语言知识（Clark等人，2019年）。这种语言知识包括句法和共指信息，对于自然语言理解至关重要。因此，我们提出基于注意力的蒸馏，以鼓励将语言知识从教师（BERT）转移到学生（TinyBERT）。具体而言，学生学习拟合教师网络中的多头注意力矩阵，目标定义为：

嵌入层蒸馏

与基于隐藏状态的蒸馏类似，我们还执行嵌入层蒸馏，目标为：

预测层蒸馏

除了模仿中间层的行为外，我们还使用知识蒸馏来拟合教师模型的预测结果，如Hinton等人（2015年）所述。具体而言，我们惩罚学生网络的logits与教师的logits之间的软交叉熵损失：

3.2 TinyBERT学习

BERT的应用通常包括两个学习阶段：预训练和微调。BERT在预训练阶段学到的丰富知识对于压缩模型来说至关重要，应该转移到TinyBERT中。因此，我们提出了一个包括通用蒸馏和特定任务蒸馏的新型两阶段学习框架，如图1所示。通用蒸馏有助于TinyBERT学习预训练BERT中嵌入的丰富知识，这对于提高TinyBERT的泛化能力起到了重要作用。特定任务蒸馏进一步教导TinyBERT从微调后的BERT中获取知识。通过两步蒸馏，我们可以显著缩小教师和学生模型之间的差距。

通用蒸馏

我们使用未经微调的原始BERT作为教师模型，并使用大规模文本语料库作为训练数据。通过对通用领域文本执行Transformer蒸馏，我们获得了一个可以用于下游任务微调的通用TinyBERT。然而，由于隐藏层/嵌入层大小和层数的显著减少，通用TinyBERT通常比BERT表现更差。

特定任务蒸馏

先前的研究表明，复杂的模型（如微调后的BERTs）对于特定领域的任务存在过度参数化的问题（Kovaleva等人，2019年）。因此，较小的模型有可能达到与BERT相当的性能。为此，我们提出通过特定任务蒸馏来产生具有竞争力的微调TinyBERTs。在特定任务蒸馏中，我们在增强的特定任务数据集上重新执行提出的Transformer蒸馏。具体而言，使用微调后的BERT作为教师模型，并提出一种数据增强方法来扩展特定任务的训练集。通过使用更多与任务相关的示例进行训练，可以进一步提高学生模型的泛化能力。

数据增强

我们结合预训练语言模型BERT和GloVe（Pennington等人，2014年）词嵌入进行单词级别的替换以实现数据增强。具体而言，我们使用语言模型预测单片词的单词替换（Wu等人，2019年），并使用词嵌入检索多片词的最相似单词作为单词替换。定义了一些超参数来控制句子的替换比例和增强数据集的数量。数据增强程序的更多细节如算法1所示。我们在所有实验中将pt设为0.4，Na设为20，K设为15。

上述两个学习阶段是相互补充的：通用蒸馏为特定任务蒸馏提供了良好的初始化，而特定任务蒸馏在增强数据上进一步改进TinyBERT，专注于学习特定任务的知识。尽管模型大小显著减少，但通过在预训练和微调阶段执行提出的Transformer蒸馏方法，并结合数据增强，TinyBERT可以在各种自然语言处理任务中实现具有竞争力的性能。

4 实验

在本节中，我们在多种任务上评估TinyBERT的有效性和效率，并使用不同的模型设置。

4.1 数据集

我们在通用语言理解评估（GLUE）基准测试（Wang等人，2018年）上评估TinyBERT，该基准测试包括2个单句任务：CoLA（Warstadt等人，2019年）和SST-2（Socher等人，2013年）；3个句子相似性任务：MRPC（Dolan和Brockett，2005年）、STS-B（Cer等人，2017年）和QQP（Chen等人，2018年）；以及4个自然语言推理任务：MNLI（Williams等人，2018年）、QNLI（Rajpurkar等人，2016年）、RTE（Bentivogli等人，2009年）和WNLI（Levesque等人，2012年）。这些任务的评估指标可以在GLUE论文（Wang等人，2018年）中找到。

4.2 TinyBERT设置

TinyBERT学习包括通用蒸馏和特定任务蒸馏。对于通用蒸馏，我们将最大序列长度设置为128，并使用英文维基百科（25亿词）作为文本语料库，在预训练的BERTBASE的监督下执行中间层蒸馏，持续3个epoch，同时保持其他超参数与BERT预训练（Devlin等人，2019年）相同。对于特定任务蒸馏，在微调BERT的监督下，我们首先在增强数据上执行中间层蒸馏，持续20个epoch4，批量大小为32，学习率为5e-5，然后在增强数据上执行预测层蒸馏，持续3个epoch5，批量大小从{16, 32}中选择，学习率从{1e-5, 2e-5, 3e-5}中选择，基于验证集进行调整。在特定任务蒸馏中，单句任务的最大序列长度设置为64，序列对任务的最大序列长度设置为128。

4.3 基线

我们将TinyBERT与BERTTINY、BERTSMALL6（Turc等人，2019年）以及几种最新的知识蒸馏（KD）基线进行比较，包括BERT-PKD（Sun等人，2019年）、PD（Turc等人，2019年）、DistilBERT（Sanh等人，2019年）和MobileBERT（Sun等人，2020年）。BERTTINY表示直接预训练一个小BERT，其模型架构与TinyBERT4相同。在训练BERTTINY时，我们遵循原始BERT（Devlin等人，2019年）中描述的相同学习策略。为了进行公平比较，我们使用发布的代码训练了一个4层的BERT4-PKD7和一个4层的DistilBERT48，并使用建议的超参数对这些4层基线进行微调。对于6层基线，我们使用报告的数字或使用发布的模型在GLUE测试集上评估结果。

4.4 GLUE上的实验结果

我们将模型预测提交给官方GLUE评估服务器，以获得测试集9上的结果，如表1所示。4层学生模型的实验结果表明：1）由于模型大小的显著减少，BERTTINY（或BERTSMALL）与BERTBASE之间存在较大的性能差距；2）TinyBERT4在所有GLUE任务上均优于BERTTINY，平均性能提升了6.8%。这表明提出的KD学习框架可以有效地提高小型模型在各种下游任务上的性能；3）TinyBERT4显著优于4层的最新KD基线（即BERT4-PKD和DistilBERT4），性能提升至少4.4%，同时参数数量减少了约28%，推理速度加快了3.1倍；4）与教师BERTBASE相比，TinyBERT4模型大小缩小了7.5倍，推理速度加快了9.4倍，同时保持了具有竞争力的性能；5）对于具有挑战性的CoLA数据集（预测语言可接受性判断的任务），所有4层蒸馏模型与教师模型相比都存在较大的性能差距，而TinyBERT4在4层基线上实现了显著的性能提升；6）我们还将TinyBERT与从24层IB-BERTLARGE蒸馏而来的24层MobileBERTTINY进行了比较。结果表明，TinyBERT4仅使用38.7%的浮点运算次数（FLOPs）就达到了与24层模型相同的平均分数；7）当我们增加模型容量到TinyBERT6时，其性能可以进一步提升，并且在平均性能上比相同架构的基线高出2.6%，与教师模型的性能相当；8）与另一种两阶段基线PD相比，PD首先预训练一个小BERT，然后在特定任务上使用这个小模型进行蒸馏，而TinyBERT通过通用蒸馏在特定任务阶段初始化学生模型。我们在附录C中分析了这两种初始化方法。此外，BERT-PKD和DistilBERT通过使用预训练BERT的一些层来初始化学生模型，这使得学生模型必须保持与教师模型相同的Transformer层（或嵌入层）大小设置。在我们的两阶段蒸馏框架中，TinyBERT通过通用蒸馏进行初始化，使其在选择模型配置时更加灵活。

更多比较

我们通过引入更多基线（如Poor Man's BERT（Sajjad等人，2020年）、BERT-of-Theseus（Xu等人，2020年）和MiniLM（Wang等人，2020年））来展示TinyBERT的有效性，其中一些基线仅在GLUE开发集上报告了结果。此外，我们还在SQuAD v1.1和v2.0上评估了TinyBERT。由于篇幅限制，我们在附录A和B中呈现了结果。

4.5 消融研究

在本节中，我们进行消融研究，以调查：a）图1中提出的两阶段TinyBERT学习框架的不同步骤的贡献；b）公式11中不同的蒸馏目标的贡献。

4.5.1 学习步骤的影响

提出的两阶段TinyBERT学习框架包括三个关键步骤：GD（通用蒸馏）、TD（特定任务蒸馏）和DA（数据增强）。移除每个单独学习步骤的性能分析结果如表2所示。结果表明，所有三个步骤对于提出的方法都至关重要。TD和DA在所有四个任务中具有相当的效果。我们注意到，特定任务步骤（TD和DA）在所有任务中比预训练步骤（GD）更有帮助。另一个有趣的观察是，GD在CoLA任务中的贡献大于MNLI和MRPC任务。我们推测，GD学到的语言泛化能力（Warstadt等人，2019年）在语言可接受性判断任务中发挥了重要作用。

4.5.2 蒸馏目标的影响

我们研究了蒸馏目标对TinyBERT学习的影响。提出了几个基线，包括不进行Transformer层蒸馏的学习（w/o Trm）、不进行嵌入层蒸馏（w/o Emb）或不进行预测层蒸馏（w/o Pred）10。结果如表3所示，表明所有提出的蒸馏目标都有用。不进行Transformer层蒸馏的性能（w/o Trm）显著下降，从75.6降至56.3。这种显著下降的原因在于学生模型的初始化。在预训练阶段，为基于Transformer的模型蒸馏获得良好的初始化至关重要，而在不进行Transformer层蒸馏的情况下，这一阶段没有来自上层的监督信号来更新Transformer层的参数。此外，我们还研究了Transformer层蒸馏中注意力（Attn）和隐藏状态（Hidn）的贡献。我们发现，基于注意力的蒸馏比基于隐藏状态的蒸馏影响更大。同时，这两种知识蒸馏方法相互补充，使它们成为我们在实验中针对基于Transformer模型的最重要的蒸馏技术。

4.6 映射函数的影响

我们还研究了不同的映射函数n=g(m)对TinyBERT学习的影响。我们在第4.2节中描述的原始TinyBERT使用了均匀策略，并与两个典型基线进行了比较，包括顶部策略（g(m)=m+N−M;0<m≤M）和底部策略（g(m)=m;0<m≤M）。比较结果如表4所示。我们发现，顶部策略在MNLI任务上的表现优于底部策略，但在MRPC和CoLA任务上表现更差，这证实了不同任务依赖于BERT不同层的知识的观察结果。均匀策略涵盖了BERTBASE从底层到顶层的知识，并在所有任务中都取得了比其他两个基线更好的性能。为特定任务自适应选择层是一个具有挑战性的问题，我们将其作为未来的研究方向。

5 相关工作

预训练语言模型压缩

一般来说，预训练语言模型（PLMs）可以通过低秩近似（Ma等人，2019年；Lan等人，2020年）、权重共享（Dehghani等人，2019年；Lan等人，2020年）、知识蒸馏（Tang等人，2019年；Sanh等人，2019年；Turc等人，2019年；Sun等人，2020年；Liu等人，2020年；Wang等人，2020年）、剪枝（Cui等人，2019年；McCarley，2019年；F.等人，2020年；Elbayad等人，2020年；Gordon等人，2020年；Hou等人，2020年）或量化（Shen等人，2019年；Zafrir等人，2019年）来进行压缩。在本文中，我们的关注点是知识蒸馏。

PLMs的知识蒸馏

已经有一些工作尝试将预训练语言模型（PLMs）蒸馏到更小的模型中。BiLSTMSOFT（Tang等人，2019年）将BERT中的特定任务知识蒸馏到单层BiLSTM中。BERT-PKD（Sun等人，2019年）不仅从教师的最后一层提取知识，还在微调阶段从中间层提取知识。DistilBERT（Sanh等人，2019年）在大规模语料库上进行预训练阶段的蒸馏。同期工作MobileBERT（Sun等人，2020年）通过在预训练阶段进行逐步知识转移，将带有瓶颈结构的BERTLARGE蒸馏成一个24层的简化版本。MiniLM（Wang等人，2020年）也在预训练阶段进行了深度自注意力蒸馏。相比之下，我们提出了一个新的两阶段学习框架，通过一种新的Transformer蒸馏方法在预训练和微调阶段从BERT中提取知识。

预训练轻量级PLMs

其他相关工作旨在直接预训练轻量级PLMs。Turc等人（2019年）预训练了24个小BERT模型，并表明在较小架构的背景下，预训练仍然很重要，微调预训练的紧凑模型可以具有竞争力。ALBERT（Lan等人，2020年）通过嵌入因子化和跨层参数共享来减少模型参数。由于ALBERT没有减少Transformer块的隐藏大小或层数，因此它仍然有大量的计算。另一个同期工作，ELECTRA（Clark等人，2020年）提出了一个名为"替换标记检测"的样本高效任务，以加速预训练，并展示了一个12层的ELECTRAsmall，其性能与TinyBERT4相当。与这些小PLMs不同，TinyBERT4是一个4层模型，可以实现更高的加速。

6 结论与未来工作

在本文中，我们介绍了一种新的基于Transformer的蒸馏方法，并进一步提出了一个两阶段的TinyBERT框架。广泛的实验表明，TinyBERT在显著减少BERTBASE模型大小和推理时间的同时，实现了具有竞争力的性能，为在边缘设备上部署基于BERT的自然语言处理模型提供了一种有效的方法。在未来的研究中，我们将研究如何有效地将更宽和更深的教师模型（例如BERTLARGE）的知识转移到学生TinyBERT中。将蒸馏与量化/剪枝相结合，将是进一步压缩预训练语言模型的一个有希望的方向。