小版BERT也能出奇迹：最火的预训练语言库探索小巧之路

近日，HuggingFace 发布了 NLP transformer 模型——DistilBERT，该模型与 BERT 的架构类似，不过它仅使用了 6600 万参数（区别于 BERT_base 的 1.1 亿参数），却在 GLUE 基准上实现了后者 95% 的性能。

在过去 18 个月中，基于大规模语言模型的迁移学习显著提升了自然语言处理任务的 SOTA 成绩。这些预训练语言模型通常基于 Vaswani 等人提出的 Transformer 架构，这些模型的规模越来越大，训练数据集的规模也越来越大。最近英伟达提出的预训练语言模型拥有 83 亿参数：是 BERT-large 参数量的 24 倍、GPT-2 参数量的 5 倍。而 Facebook AI 最近提出的 RoBERTa 模型在 160GB 文本上训练得到。

社区中的一些人质疑训练越来越大 Transformer 的必要性，尤其是考虑到训练的资金成本和环境成本时。 该图展示了部分近期大模型及其参数量。

Hugging Face 直接体会到这些模型的流行度，因为其预训练语言库（包含这些模型中的大部分）在近几个月的安装量超过 40 万次。

• NLP 库地址：https://github.com/huggingface/pytorch-transformers

然而，尽管这些模型被更大的 NLP 社区接受，一个重要且有挑战性的问题出现了。如何将这些庞然大物投入到生产中？如何在低延迟约束下使用这些大模型？我们需要用（昂贵的）GPU 服务器执行大规模服务吗？

为了构建更尊重隐私的系统，Hugging Face 注意到在边缘设备上运行机器学习系统的需求在不断增长，而不是调用云 API，将隐私数据发往服务器。在智能手机等设备上运行的模型需要是轻量级、响应快和能源利用率高的！

最后但同样重要的一点，Hugging Face 越来越担忧这些大模型所需的指数级计算成本增长。

有很多技术可以解决前述问题。最常见的工具是量化（使用更小精度逼近全精度模型）和权重剪枝（移除网络中的部分连接）。想了解更多，可以查看这篇关于 BERT 量化的精彩博客：https://blog.rasa.com/compressing-bert-for-faster-prediction-2/。

Hugging Face 的研究者决定把重点放在知识蒸馏（distillation）上。蒸馏即，将较大模型（教师模型）压缩成较小模型（学生模型）的方法。

知识蒸馏： 迁移泛化能力

知识蒸馏是一种模型压缩方法，又叫师生学习（teacher-student learning）。它训练一个小模型，使之复制大模型（或模型集成）的行为。知识蒸馏由 Bucila 等人提出，几年后被 Hinton 等人推广（参见论文《Distilling the Knowledge in a Neural Network》）。Hugging Face 研究者使用的是 Hinton 等人的方法。

在监督学习中，分类模型通常用于预测类别，它利用对数似然信号最大化类别概率。在很多案例中，高性能模型预测的输出分布中，正确的类别具备高概率，而其他类别的概率则接近于零。

例如，desk chair（办公椅）可能会被误分类为 armchair（扶手椅），但通常不会被误认为是 mushroom（蘑菇）。这种不确定性被称为「暗知识」。

理解蒸馏的另一种方式是，它阻止模型对预测结果过于自信（类似于标签平滑）。

以下是一个实践示例。在语言建模过程中，我们通过观察词汇分布，可以轻松发现这种不确定性。下图展示了 Bert 对电影《卡萨布兰卡》中某句著名台词的 top 20 补全结果：

BERT_base 对被遮蔽 token 的 top 20 补全结果。 该语言模型确定了两个概率较高的 token（day 和 life）。

如何复制暗知识？

在师生训练中，我们训练学生网络来模拟教师网络的完整输出分布（它的知识）。

我们使学生网络和教师网络具备同样的输出分布，从而使学生网络实现同样的泛化。

我们不对硬目标类别（正确类别的 one-hot 编码）使用交叉熵来进行训练，而是对软目标类别（教师网络的概率）执行交叉熵，从而将教师网络的知识迁移到学生网络。这样训练损失函数就变成了：

其中 t 表示教师网络的 logit 值，s 表示学生网络的 logit 值。该损失函数具备更丰富的训练信号，因为软目标类别比单个硬目标类别提供更多约束。

为了进一步揭示类别分布的多样性，Hinton 等人提出了 softmax-temperature：

其中T 表示温度参数，当 T → 0 时，分布接近于 one-hot 目标向量，当 T →+∞ 时，则得到均匀分布。在训练过程中对教师网络和学生网络使用同样的温度参数，进而为每一个训练样本提供更多信号。在推断时，T 被设置为 1，恢复标准的 Softmax 函数。

PyTorch 动手实践： 压缩 BERT

Hugging Face 研究者想利用知识蒸馏压缩大型语言模型。对于蒸馏，研究者使用 KL 散度作为损失函数，因为最优化过程与交叉熵是等价的：

在计算 q（学生网络的分布）的梯度时，得到了同样的梯度。这允许研究者利用 PyTorch 实现执行更快速的计算：

使用教师网络 BERT 的监督信号，研究者训练得到较小的语言模型——DistilBERT。（研究者使用的是 Bert 的英语 bert-base-uncased 版本）。

按照 Hinton 等人的方法，训练损失是蒸馏损失和遮蔽语言建模损失的线性组合。学生模型是 BERT 的较小版本，研究者移除了 token 类型的嵌入和 pooler（用于下一句分类任务），保留了 BERT 的其余架构，不过网络层数只有原版的 1/2。

备注 1：为什么不减少隐藏层大小呢？将它从 768 减少到 512 即可将参数总量减少约 1/2。但是，在现代框架中，大部分运算是经过高度优化的，张量最后一维（隐藏维度）的变化对 Transformer 架构中使用的大部分运算影响较小。在研究者的实验中，相比隐藏层大小，层数才是推断阶段的决定性因素。

研究者的早期实验表明，在该案例中，交叉熵损失会带来更好的性能。因此，他们假设在语言建模设置中，输出空间（词汇）要比下游任务输出空间的维度大得多。而在 L2 损失中，logit 可能会相互抵消。

训练子网络不仅仅关乎架构，它还需要找出子网络收敛的合适初始化（例如彩票假设论文《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》）。因此，研究者基于教师网络 Bert 对学生网络 DistilBERT 进行初始化，将层数减半，对学生网络使用与教师网络一样的隐藏层大小。

研究者还使用 RoBERTa 论文中的一些训练 trick，这篇论文证明 Bert 的训练方式对最终性能有着重要影响。遵循 RoBERTa 的训练方式，研究者利用梯度累积以非常大的批次（每个批次多达 4000 个样本）训练 DistilBERT，训练使用了动态遮挡（dynamic masking），并移除了下一句预测目标。

该训练设置主动接受资源方面的限制。研究者在 8 块 16GB V100 GPU 上训练 DistilBERT，训练时长接近三天半，训练数据为 Toronto Book Corpus 和英文维基百科（与原版 BERT 的训练数据相同）。

DistilBert 的部分代码来自于 Facebook XLM 的代码，部分代码来自 Google AI BERT 的 Hugging Face PyTorch 版本。这些代码可在 Hugging Face 的 NLP 库中获取，该库还包含多个 DistilBert 训练版本和微调版本，及其复现代码。

模型性能： DistilBERT 测试

研究者在 GLUE 基准的开发集上对比了 DistilBERT 和两个基线模型的性能，基线模型分别是 BERT base（DistilBERT 的教师模型）和来自纽约大学的强大非 transformer 基线模型：ELMo + BiLSTMs。研究者使用纽约大学发布的 ELMo 基线 jiant 库和 BERT 基线模型的 PyTorch-Transformers 版本。

如下表所示，DistilBERT 在参数量分别是基线模型的 1/2 和 1/3 的情况下，性能可与后者媲美。在 9 项任务中，DistilBERT 的性能通常等同于或优于 ELMo 基线（在 QNLI 任务上的准确率超出后者 14 个百分点）。令人惊讶的是，DistilBERT 的性能堪比 BERT：在参数量比 BERT 少 40% 的情况下，准确率达到了后者的 95%。

在 GLUE 基准开发集上的模型对比结果。 ELMo 的性能结果来自原论文，BERT 和 DistilBERT 的性能结果是使用不同种子进行 5 次运行后的中位数 。

至于推断性能，DistilBERT 的推断速度比 BERT 快 60%，规模也比后者小；DistilBERT 的推断速度比 ELMo+BiLSTM 快 120%，规模也比后者小很多。

为了进一步调查 DistilBERT 的加速/规模权衡（speed-up/size trade-off），研究者对比了每个模型的参数量和在 STS-B 开发集上使用一块 CPU、批大小为 1 的情况下完成一个完整 pass 的推断时间，如上表所示。

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

研究者进一步研究了在高效推断约束下，DistilBERT 在下游任务上的应用。研究者对紧凑的预训练语言模型DistilBERT 进行微调，用于分类任务。这是结合蒸馏预训练和迁移学习的绝妙方式！

研究者使用 IMDB 评论情感分类数据集，该数据集包含 5 万条英文评论（被标注为积极或消极）：其中 2.5 万条作为训练数据，另外一半作为测试数据（均类别均衡）。研究者使用一块 12GB K80 GPU 进行训练。

首先，在数据集上训练 bert-base-uncased。该模型达到了 99.98% 的准确率（3 次运行的平均值），几乎完美！

然后使用同样的超参数训练 DistilBERT。该模型达到了 99.53% 的准确率（3 次运行的平均值），在延迟降低 60%、规模减少 40% 的情况下，DistilBERT 的性能仅比原版 BERT 低 0.5%！

少即是多： 小模型也能出奇迹

Hugging Face 对 DistilBERT 的潜力非常看好。DistilBERT 只是个开始，它也提出了许多问题：使用知识蒸馏技术，我们可以把大模型压缩到什么程度？这些技术可用于进一步探索和洞察大模型中存储的知识吗？在压缩过程中损失了语言学/语义学的哪些方面？……

目前，HuggingFace 的这项研究已经开源，并进行知识共享。他们认为这是每个人参与 NLP 进展，并收获最新进展果实的最快和最公平的路径。

• GitHub 地址：https://github.com/huggingface

• Medium 页面：http://www.medium.com/huggingface