DistilBERT:用蒸馏把 BERT 做小做快

快速答案

DistilBERT 在预训练阶段(而非微调之后)蒸馏 BERT,得到一个 6 层、参数少 40%(6600 万对 1.1 亿)的模型,推理约快 60%,同时保留 BERT-base 约 97% 的 GLUE 成绩。它的贡献在于证明:一个通用的蒸馏模型,只在和教师相同的语料上训练一次,迁移到下游任务时几乎不输完整模型。蒸馏一次,到处复用同一份权重。

做蒸馏的三重损失

DistilBERT 用三个目标同时训练这个小学生模型。蒸馏损失对齐教师完整的软标签分布,并用温度让学生也学到教师对”哪些错误答案更可信”的判断;标准的掩码语言建模损失让它扎根于真实文本;余弦嵌入损失则让学生隐状态的方向与教师对齐。把软标签和 MLM 结合起来,才让一个深度减半的模型贴近 BERT,而不是只在微调后用任务标签蒸馏。

从 BERT 砍掉了什么

学生保留了 BERT 的隐藏维度和注意力设计(它继承了 Transformer 编码器),但把层数从 12 砍到 6,并去掉了 token 类型嵌入和 pooler。关键的一点是,这 6 层从训练好的教师里每隔一层取一层来初始化,所以蒸馏从一个好起点出发,而非随机权重。训练用和 BERT 相同的数据(英文维基百科和 Toronto Book Corpus)、大批量和动态掩码,并丢掉了下一句预测。

关键结果

• 小 40%: 6600 万参数,对比 BERT-base 的 1.1 亿。
• 推理约快 60%(CPU),端上约快 71%(在 iPhone 上实测),模型约 207MB。
• 保留 BERT-base 约 97% 的 GLUE 成绩,这是九个 GLUE 任务上的平均值。
• 预训练阶段只蒸馏一次: 同一份 DistilBERT 权重可在各任务上微调,不像按任务蒸馏那样每个任务都要重做。

“保留 97%“掩盖了什么

97% 是 GLUE 上的平均值,各任务有差异,在更难、更重推理的数据集上差距会变大。而且这是对 BERT-base 的保留率,不是对 RoBERTa 或更大教师,所以”教师的 97%“只和你选的教师一样强。提速也只在推理侧:你仍要付一次 BERT 规模的成本去训练那个被蒸馏的教师。

局限与存疑

DistilBERT 需要一个已经训练好的教师,所以它没有消除预训练 BERT 的成本,只是把它摊薄。后续工作(TinyBERT、MobileBERT)用逐层蒸馏和注意力蒸馏把压缩推得更远,所以 DistilBERT 是易读的基线,而非模型压缩的前沿。原始版本只支持英文,且准确率差距意味着:对质量敏感、对延迟不敏感的应用,仍应选用完整模型。

今天还该用 DistilBERT 吗(工程判断)

对延迟或成本敏感、能接受掉几个点准确率的 NLP 场景,DistilBERT 仍是合理的默认选择,也是最易部署的压缩编码器之一。需要最高准确率时,用 BERT-large 或 RoBERTa;需要比 40% 更紧的显存或延迟预算时,评估后续那些更激进的蒸馏方法。它的长期价值是方法论上的:它把通用的、预训练期的蒸馏变成了标准工具。

常见问题

DistilBERT 的三重损失怎么工作?

它结合三个目标:对齐教师温度软化后 logits 的蒸馏损失、在真实文本上的标准掩码语言建模损失,以及对齐学生与教师隐状态方向的余弦嵌入损失。软标签携带的信息比硬标签多,这正是让 6 层模型贴近 BERT 的关键。

“保留 BERT 的 97%“是说每个任务都 97% 吗?

不是。97% 是九个 GLUE 任务上的平均值,更难的任务掉得更多,且对比对象是 BERT-base 而非更强的教师。把它当作总体平均,而非逐任务保证。

DistilBERT 从 BERT 砍掉了什么,有多大?

它把层数从 12 砍到 6,去掉 token 类型嵌入和 pooler,保留隐藏维度,得到 6600 万参数(对比 1.1 亿)。学生在蒸馏前从教师每隔一层初始化。

压缩时该选 DistilBERT 还是 TinyBERT、MobileBERT?

DistilBERT 是最简单、最易读的基线,只在输出和嵌入层蒸馏。TinyBERT 和 MobileBERT 用逐层和注意力蒸馏压得更狠,需要比 DistilBERT 的 40% 缩减更紧的预算时选它们。

阅读 arXiv 原文。