感兴趣的朋友可以点击此处获取重现DistilBERT训练以及DistilBERT预训练权重的代码。

过去18个月以来，大规模语言模型的迁移学习可谓遍地开花，在几乎所有自然语言处理任务当中都实现了显著的性能改进。

作为以Vaswani等人的Transformer架构为基础的解决方案，这些经过预训练的语言模型正变得日益庞大，并仍在立足更大的数据集进行训练。英伟达公司的最新模型拥有多达83亿个参数：24倍于BERT，5倍于GPT-2；而Facebook AI拿出的RoBERTa则利用160 GB文本训练而成。

社区中的不少从业者开始怀疑，到底有没有必要训练这些越来越臃肿的Transformer，毕竟其在训练的经济与环境成本方面已经呈现出失控的状态。通过上图，我们一起来看部分最新大型模型及其参数数量（单位为百万）。

在Hugging Face，我们亲身体验了这些模型高涨的人气，因为我们的NLP库（打包了其中大部分模型）在短短几个月之内就得到超过40万次安装。

然而，随着这些模型进入更大的社区，一个重要甚至说极具挑战性的问题开始出现——我们该如何把这些庞然大物投入生产？我们如何在低延迟约束条件下使用这类大型模型？我们是否需要昂贵的GPU服务器才能实现大规模服务？

对于许多研究人员及开发人员而言，这可能是个最现实的问题。

为了构建更尊重隐私的系统，我们意识到有必要在边缘位置运行机器学习系统，从而尽可能避免以调用云API的方式将个人数据发送至服务器端。这就意味着，我们需要能够在智能手机等小型设备上运行轻、反应灵敏且资源需求量较低的模型版本！

最后但同样重要的是，我们也越来越关注模型扩展过程当中，严苛计算资源需求所带来的环境成本。

那么，我们该如何帮助这些庞然大物成功瘦身？

不少现有技术都有望解决问题。最常见的工具包括量化（对准确率影响较小的网络权重进行近似化）以及权重修剪（删除网络中的某些连接）。对于此类技术，推荐大家参阅Rasa发布的BERT量化博文。

但我们最终决定专注于模型蒸馏：这是一种能够将大型模型（被称为「老师」）压缩为较小模型（即「学生」）的技术。

知识蒸馏：迁移泛化能力

知识蒸馏（有时也称为师生学习）是一种压缩技术，要求对小型模型进行训练，以使其拥有类似于大型模型（或者模型集合）的行为特征。这项技术由Bucila等人提出，并得到了Hinton等人的推广。我们这里采用的，正是Hinton采取的方法。

在监督学习领域，我们在训练分类模型时往往会利用对数似然信号实现概率最大化（logits的softmax），进而预测出正确类。在大多数情况下，性能良好的模型能够利用具有高概率的正确类预测输出分布，同时其它类的发生概率则接近于零。

但是，某些“接近于零”的概率要比其它概率更大，这在一定程度上反映出模型的泛化能力。

例如，把普通椅子误认为扶手椅虽然属于错误，但这种错误远比将其误认为蘑菇来得轻微。这种不确定性，有时被称为“暗知识”。

我们也可以从另一个角度来理解蒸馏——用于防止模型对预测结果太过确定（类似于标签平滑）。

以下为具体实例。在语言建模当中，我们可以通过查看词汇表中的分布轻松观察到这种不确定性。下图为BERT对《卡萨布兰卡》电影当中经典台词下一句用词的猜测：

BERT提出的20大高概率用词猜测结果。语言模型确定了两个可能性最高的选项（day与life），接下来的词汇相比之下概率要低得多。

我们如何复制这些“暗知识”？

在师生训练当中，我们训练学生网络，用于模拟老师网络的全部输出分布（也就是知识）。

我们通过匹配输出分布的方式训练学生网络，从而实现与老师网络相同的泛化方式。

我们并没有在硬目标上使用交叉熵训练（正确类的独热编码），而是通过软目标（老师概率）将交叉熵从老师处传递给学生。我们的训练损失因此变为：

其中t为来自老师的logit，s为学生的logit。

这个损失函数属于更丰富的训练信号，因为单一示例要比单一硬目标拥有更高的强制约束效果。

为了进一步揭示分类结果的质量，Hinton等人提出了softmax温度的概念：

T为该温度参数。

当T → 0时，分布变为Kronecker（相当于独热目标矢量）；当 T →+∞时，则变为均匀分布。在训练过程中，将相同的温度参数应用于学生与老师网络，即可进一步为每个训练示例揭示更多信号。在推论当中，T被设置为1以恢复标准Softmax。

PyTorch编码——压缩BERT

我们希望利用蒸馏方法对大型语言模型加以压缩。在蒸馏方面，我们使用Kullback-Leibler损失函数，因为其拥有相同的优化效果：

在计算关于q（学生网络分布）的梯度时，我们获得了相同的梯度结果。我们可以利用PyTorch实现加快计算速度：

PyTorch中的知识蒸馏训练步骤。点击此处复制gist。

利用老师信号，我们能够训练出一套较小的语言模型，我们称之为DistilBERT，属于BERT的监督产物（我们使用BERT的英文bert-base-uncased版本）。

根据Hinton等人的发现，训练损失函数属于蒸馏损失与masked语言建模损失的线性组合。我们的学生网络属于BERT的一套小型版本，其中删除了token-type嵌入与pooler（用于下一句分类任务），但其余部分架构保持不变，而层数也减少至原本的二分之一。

总体而言，我们的蒸馏模型DistilBERT在总体参数数量上约为BERT的一半，但在GLUE语言理解基准测试中能够保留95%的BERT性能表现。

注1 — 为什么不降低隐藏层的大小？
将768层减至512层，意味着总参数量约下降至原本的二分之一。希，在现代框架当中，大多数运算都经过高度优化，而且张量的最终维度（隐藏维度）的变化会对Transformer架构（线性分层与层规范化）中的大部分运算产生小幅影响。在我们的实验中，层数对于推理时间的影响要远高于隐藏层的大小。
因此，更小并不代表着一定更快……

注2 — Tang等人在蒸馏工作当中，直接在下游任务内使用L2距离作为蒸馏损失
我们的早期实验结果表明，在本案例中，交叉熵损失会明显提高性能水平。我们假定在语言建模设置当中，输出空间（词汇表）要明显大于下游任务输出空间的维度。因此，logits可以在L2损失中相互补偿。

训练子网络的核心不只是建立架构，还要求我们为子网络找到正确的初始化方式以实现收敛。因此，我们以作为老师的Bert为基础对学生DistilBERT进行初始化，将层数削减一半，并采用相同的隐藏大小。

我们还用到了最近RoBERTa论文当中提到的一些训练技巧，这也再次证明BERT模型的训练方式对其最终表现有着至关重要的影响。与RoBERTa类似，我们对DIstilBERT进行大批次训练，使用梯度累积（每批最多4000个例子）、配合动态遮挡并删除了下一句预测目标。

我们在训练设置中对资源进行了主动限制。我们利用多伦多图书语料库与英语维基百科的串联数据集（与原始BERT相同），并配合八块16 GB V100 GPU进行了约三天半的训练。

DistilBert的代码部分来自Facebook XLM，也有一部分来自我们PyTorch版本的Google AI Bert（可点此获取），以及针对DistilBert进行的精心调优。这一切，都是为了更好地重现BERT的预测性能。

DistilBERT模型性能测试

我们将DistilBERT在GLUE基准测试开发集上的性能与两项基准进行了比较：其一为BERT基础（DistilBERT的老师），其二为来自纽约大学的强大非transformer基准——ELMo上的两个BilSTM。我们利用纽约大学的jiant库获取ELMo基准，并使用pytorch-transformers获取BERT基准。

如下表所示，DistilBERT的性能与基准相比更好一些，而参数数量只分别相当于二者的一半以及三分之一。在9项任务当中，DistilBERT的ELMo基准成绩一直等同或者领先（在QNLI上的准确率高出14%）。DistilBERT的表现确实远超BERT：我们保留了95%以上的性能，同时将参数减少了40%。

在GLUE基准测试开发集中的比较结果以及由作者上报的ELMo结果。BERT与DistilBERT结果来自5次单独运行后的中位数取值。

在推理时间方面，DistilBERT比BERT快60%，体积比BERT小60%，比ELMo + BiLSTM快120%且模型体积更小。

为了进一步研究DistilBERT的加速/大小平衡点，我们在上表中比较了各个模型的参数数量，以及在CPU上完全处理STS-B开发集（批量大小为1）所需要的推理时间。

下游任务：蒸馏与迁移学习

我们进一步研究了DistilBERT在有效推理约束下的下游应用效果。我们通过分类任务调优，实现对这套紧凑预训练语言模型的迁移。事实证明，这是种实现蒸馏预训练与迁移学习的好方法！

从IMDB Review数据集中提取到的电影评论。

我们选择了IMDB影评情感区中的素材，该分区共包含5万条英文评论，且标记为正面或负面：我们使用其中2万5千条进行训练，另外2万5千条进行测试（同时配合平衡类）。整个训练过程在单一12 GB K80上进行。

首先，我们在自己的数据集上训练bert-base-uncased。我们亲爱的BERT老师达到了99.98%的准确率（3次运行取平均值）。相当完美！

接下来，我们训练DistilBERT，使用同样的超参数。压缩模型的准确率达到99.53%（3次运行取平均值）。性能的绝对差为0.5%，延迟降低60%，大小减少40%。

NLP技术的另一种常见应用是问题解答。我们在SQuAD 1.1数据集上比较了BERT bert-base-uncased版本与DistilBERT的结果。在开发集上，BERT的F1得分为88.5，EM（完全匹配）得分为81.2。我们利用同样的超参数进行DistilBERT训练，F1分数与EM分数分别为85.1与76.5，同BERT成绩的差距分别为3分与5分。

我们还研究了能否在适应阶段利用经过调优的BERT作为老师，配合知识蒸馏损失对DistilBERT实现SQuAD数据集上的调优。

在新案例中，我们将问题回答模型蒸馏为以往通过知识蒸馏预训练完成的语言模型，从而实现调优！这样，老师与学生将能够相互转换。

如此一来，考虑到网络规模，我们能够获得非常有趣的结果：F1得分为86.2，EM得分为78.1。与完整模型相比，差距保持在3分以内！

少即是多：小型模型也能带来理想性能

我们对DistilBERT的潜力感到非常兴奋。目前的成果只是刚刚起步，也给我们提出了很多新的问题：我们能够利用知识蒸馏技术将这些模型压缩到怎样的程度？这些技术能否用于进一步理解大型模型中存储的知识？在这类压缩当中，我们损失掉的是哪些语言/语义元素？……

在HuggingFadce，我们一直将开源与知识共享视为自己的使命。所以，大家可以点击此处访问我们的GitHub库，这是我们每个人参与NLP深度学习项目并获取卓越成果的最简单、也最公平的方式。

因此，配合本篇博文，我们在pytorch-transformer库当中发布了实验代码（主要是重现训练与DistilBERT调优代码）以及一套经过训练的DistilBERT版本，感兴趣的朋友可以随意取用。

原文链接：

https://medium.com/huggingface/distilbert-8cf3380435b5

创作场景

更小、更快、更便宜、更轻量：开源 DistilBERT，BERT 的精简版本