当前位置：首页 > article >正文

DistilBERT 是 BERT 的精简版本，具有更小、更快、更经济、更轻便的特点。

article 2025/2/1 17:48:34

摘要

随着大规模预训练模型的迁移学习在自然语言处理（NLP）中变得越来越普遍，在边缘设备上或受限的计算训练/推理预算下运行这些大型模型仍然具有挑战性。在本研究中，我们提出了一种预训练较小通用语言表示模型的方法，称为 DistilBERT，该模型可以在广泛的任务上像其更大规模的对应模型一样进行微调并取得良好性能。尽管大多数先前的研究探讨了使用蒸馏技术构建特定任务模型的方法，但我们在预训练阶段利用知识蒸馏，并证明可以将 BERT 模型的规模减少 40%，同时保留其 97% 的语言理解能力，且速度提高 60%。为了利用更大模型在预训练期间学到的归纳偏差，我们引入了一种结合语言建模、蒸馏和余弦距离损失的三重损失函数。我们更小、更快、更轻的模型在预训练成本上更低，并通过概念验证实验和对比性设备端研究展示了其在设备端计算中的能力。

1 简介

过去两年，自然语言处理（NLP）领域见证了迁移学习方法的兴起，大规模预训练语言模型已成为许多 NLP 任务的基本工具 [Devlin 等人, 2018, Radford 等人, 2019, Liu 等人, 2019]。尽管这些模型带来了显著的性能提升，但它们通常具有数亿个参数，而当前关于预训练模型的研究表明，训练更大的模型仍然能够在下游任务中取得更好的性能。

在这里插入图片描述

随着模型规模不断增大的趋势，引发了一些担忧。首先，正如 Schwartz 等人 [2019] 和 Strubell 等人 [2019] 所提到的，模型计算需求的指数级增长带来了环境成本问题。其次，尽管在设备上实时运行这些模型有可能推动新颖且有趣的语言处理应用的发展，但这些模型日益增长的计算和内存需求可能会阻碍其广泛采用。

在本文中，我们展示了通过使用知识蒸馏预训练的较小语言模型，可以在许多下游任务上达到类似的性能。这些模型在推理时更轻量、更快速，同时所需的计算训练预算也更小。我们的通用预训练模型可以在多个下游任务上进行微调并取得良好性能，保持了较大模型的灵活性。我们还展示了我们的压缩模型足够小，可以在边缘设备（例如移动设备）上运行。

通过使用三重损失函数，我们展示了通过更大 Transformer 语言模型的监督进行蒸馏预训练的、规模缩小 40% 的 Transformer（Vaswani 等人 [2017]）可以在各种下游任务上实现类似的性能，同时在推理时速度提高 60%。进一步的消融研究表明，三重损失函数的所有组成部分对于实现最佳性能都很重要。

我们已将训练好的权重与训练代码一起发布在 HuggingFace 的 Transformers 库中 [Wolf 等人, 2019]。

2 知识蒸馏

知识蒸馏（Bucila 等人，2006；Hinton 等人，2015）是一种压缩技术，其中紧凑模型（学生模型）被训练以复制更大模型（教师模型）或模型集合的行为。

在监督学习中，分类模型通常通过最大化黄金标签的估计概率来预测实例类别。因此，标准的训练目标涉及最小化模型预测分布与训练标签的单热经验分布之间的交叉熵。在训练集上表现良好的模型将在正确类别上预测高概率的输出分布，而在其他类别上预测接近零的概率。但这些“接近零”的概率中，有些比其他更大，部分反映了模型的泛化能力及其在测试集上的表现。

训练损失：学生模型通过教师模型的软目标概率进行蒸馏损失训练： $L_{ce}=\sum_{i}t_{i}*\log(s_{i})$ ，其中 $t_{i}$ （分别为 $s_{i}$ ）是教师模型（分别为学生模型）估计的概率。该目标通过利用教师模型的完整分布提供了丰富的训练信号。根据 Hinton 等人 [2015]，我们使用了带温度的 softmax： $p_{i}= \frac {\exp ( z_{i}/ T) }{\sum _{j}\exp ( z_{j}/ T) }$ ，其中 $T$ 控制输出分布的平滑度， $z_i$ 是模型对类别 $i$ 的得分。训练时，相同的温度 $T$ 应用于学生模型和教师模型，而在推理时， $T$ 设置为 1 以恢复标准 softmax。

最终的训练目标是蒸馏损失 $L_{ce}$ 与监督训练损失（在我们的案例中是掩码语言建模损失 $L_{mlm}$ [Devlin 等人，2018]）的线性组合。我们发现添加余弦嵌入损失 $L_{cos})$ 是有益的，它有助于对齐学生模型和教师模型隐藏状态向量的方向。

3 DistilBERT：BERT 的蒸馏版本

学生模型架构：在本工作中，学生模型 DistilBERT 具有与 BERT 相同的通用架构。去除了 $t o k e n - t y p e$ 嵌入和池化器，同时将层数减少了一半。Transformer 架构中使用的大多数操作（线性层和层归一化）在现代线性代数框架中高度优化，我们的研究表明，张量的最后一个维度（隐藏大小维度）的变化对计算效率的影响（在固定参数预算下）比其他因素（如层数）的变化要小。因此，我们专注于减少层数。

学生模型初始化：除了上述优化和架构选择外，我们训练过程中的一个重要元素是找到子网络收敛的正确初始化。利用教师模型和学生模型之间的共同维度，我们通过从教师模型中每隔一层取一层来初始化学生模型。
在这里插入图片描述

在这里插入图片描述

蒸馏

我们应用了 Liu 等人 [2019] 最近提出的训练 BERT 模型的最佳实践。因此，DistilBERT 在非常大的批次上进行蒸馏，利用梯度累积（每批次最多 4K 个样本），使用动态掩码，并且没有下一句预测目标。

数据和计算能力

我们在与原始 BERT 模型相同的语料库上训练 DistilBERT：英文维基百科和多伦多图书语料库 [Zhu 等人，2015] 的拼接。DistilBERT 在 8 块 16GB V100 GPU 上训练了大约 90 小时。作为比较，RoBERTa 模型 [Liu 等人，2019] 需要在 1024 块 32GB V100 GPU 上训练 1 天。

4 实验

通用语言理解

我们在通用语言理解评估（GLUE）基准 [Wang 等人，2018] 上评估了 DistilBERT 的语言理解和泛化能力。GLUE 是一个包含 9 个数据集的集合，用于评估自然语言理解系统。我们通过微调 DistilBERT 报告了每个任务开发集上的分数，没有使用集成或多任务微调方案（这些方法与当前工作基本正交）。我们将结果与 GLUE 作者提供的基线进行比较：一个 ELMo (Peters 等人 [2018]) 编码器，后跟两个 BiLSTM。

表 II 显示了 9 个任务的结果以及宏观分数（各个分数的平均值）。在 9 个任务中，DistilBERT 始终与 ELMo 基线持平或有所改进（在 STS-B 上最多提高了 19 个准确率点）。DistilBERT 与 BERT 相比也表现出色，保留了 97% 的性能，同时参数减少了 40%。

4.1 下游任务基准

下游任务：我们进一步研究了 DistilBERT 在高效推理约束下的几个下游任务上的表现：一个分类任务（IMDb 情感分类 - Maas 等人 [2011]）和一个问答任务（SQuAD v1.1 - Rajpurkar 等人 [2016]）。

如表 2 所示，DistilBERT 在 IMDb 基准测试中的测试准确率仅比 BERT 低 0.6%，同时模型规模小了 40%。在 SQuAD 上，DistilBERT 与完整 BERT 的差距在 3.9 分以内。我们还研究了是否可以在适应阶段通过微调 DistilBERT 时添加另一个蒸馏步骤，使用之前在 SQuAD 上微调的 BERT 模型作为教师，在损失中添加一个额外的项（知识蒸馏）。

在这里插入图片描述

在这种情况下，我们能够在模型规模较小的情况下达到有趣的性能：79.8 F1 和 70.4 EM，即与完整模型的差距在 3 分以内。

规模和推理速度

为了进一步研究 DistilBERT 的速度提升与规模之间的权衡，我们在表 3 中比较了每个模型的参数数量以及在 CPU（Intel Xeon E5-2690 v3 Haswell @2.9GHz）上使用批量大小为 1 对 STSB 开发集进行完整推理所需的时间。DistilBERT 的参数比 BERT 少 40%，推理速度比 BERT 快 60%。

设备端计算

我们通过构建一个用于问答的移动应用程序，研究了 DistilBERT 是否可以用于边缘计算应用。我们将最新智能手机（iPhone 7 Plus）上的平均推理时间与之前基于 BERT-base 训练的问答模型进行了比较。排除分词步骤，DistilBERT 比 BERT 快 71%，整个模型的大小为 207 MB（可以通过量化进一步减小）。我们的代码已公开。

4.2 消融研究

在本节中，我们研究了三重损失和学生模型初始化对蒸馏模型性能的影响。我们报告了 GLUE 上的宏观分数。表 4 展示了与完整三重损失的差异：去除掩码语言建模损失影响较小，而两个蒸馏损失占据了性能的大部分。

5 相关工作

任务特定蒸馏

大多数先前的工作集中在构建任务特定的蒸馏设置。Tang 等人 [2019] 将微调的分类模型 BERT 转移到基于 LSTM 的分类器。Chatterjee [2019] 将在 SQuAD 上微调的 BERT 模型蒸馏到先前从 BERT 初始化的较小 Transformer 模型中。在本研究中，我们发现使用通用预训练蒸馏比任务特定蒸馏更有益。Turc 等人 [2019] 使用原始预训练目标训练较小的学生模型，然后通过蒸馏进行微调。如消融研究所示，我们发现利用教师模型的知识进行预训练并添加额外的蒸馏信号是有益的。

多蒸馏

Yang 等人 [2019] 使用多任务学习结合教师模型集合的知识来正则化蒸馏。作者应用多任务知识蒸馏从一组大型问答模型中学习紧凑的问答模型。多蒸馏的一个应用是多语言性：Tsai 等人 [2019] 采用与我们类似的方法，通过蒸馏从头开始预训练多语言模型。然而，如消融研究所示，利用教师模型的知识进行初始化和添加额外的损失可以带来显著的性能提升。

其他压缩技术

其他压缩技术也被研究用于压缩大型模型。最近在权重剪枝方面的发展表明，可以在测试时移除自注意力中的一些头部而不会显著降低性能（Michel 等人 [2019]）。某些层可以减少到一个头部。另一项研究利用量化来推导较小的模型（Gupta 等人 [2015]）。剪枝和量化与本研究是正交的。

6 结论与未来工作

我们介绍了 DistilBERT，这是一个通用预训练版本的 BERT，规模缩小了 40%，速度提高了 60%，同时保留了 97% 的语言理解能力。我们展示了通过蒸馏可以成功训练通用语言模型，并通过消融研究分析了各个组件的影响。我们进一步证明了 DistilBERT 是边缘应用中的一个引人注目的选择。

查看全文

http://www.kler.cn/a/521431.html