TinyBERT: Distilling BERT for Natural Language Understanding

EMNLP 2020

论文链接：[1909.10351] TinyBERT: Distilling BERT for Natural Language Understanding

代码链接：GitHub - huawei-noah/Pretrained-Language-Model: Pretrained language model and its related optimization techniques developed by Huawei Noah's Ark Lab.

摘要

• 传统的预训练语言模型，如 BERT，虽然在自然语言处理任务中取得了显著进展，但由于其计算开销大，难以在资源有限的设备上高效执行
• 为了解决这个问题，TinyBERT 提出了一个新颖的 Transformer 蒸馏方法，专门针对基于 Transformer 的模型进行知识蒸馏。通过这种方法，大型 BERT 模型中蕴含的丰富知识可以有效地转移到较小的学生模型 TinyBERT 上。
• 实验结果表明，TinyBERT4（4 层模型）能够达到 BERTBASE 在 GLUE 基准测试中的 96.8% 的性能，同时在模型大小和推理速度上有显著优化：模型体积缩小了 7.5 倍，推理速度提高了 9.4 倍。TinyBERT4 相比于现有的 4 层 BERT 蒸馏模型，参数量和推理时间分别减少了约 28% 和 31%。另外，TinyBERT6（6 层模型）在性能上与 BERTBASE 相当

需要解决的问题

• 预训练语言模型的挑战
  • 高计算和内存需求：BERT 模型通常非常庞大，BERTBASE 有 1.1 亿个参数，BERTLARGE 甚至有 3.4 亿个参数。这使得 BERT 模型在计算和内存上非常消耗资源
  • 推理速度慢
• 资源受限设备的应用需求：移动设备通常具有较低的计算能力和内存，无法高效运行像 BERT 这样庞大的模型。
• 解决目标：减小模型体积、加速推理速度、保持准确性

解决方法与贡献

• 提出了一种新的 Transformer 蒸馏方法：设计了三种不同的损失函数，分别对应 BERT 层中的不同表示
• 提出了一个新的两阶段学习框架，包括通用蒸馏和任务特定蒸馏
  • ​​​​​​​通用蒸馏阶段：在该阶段，原始的未调优 BERT 作为教师模型，TinyBERT 模仿教师模型的行为，通过在通用领域语料上进行 Transformer 蒸馏。通过这一阶段，获得了一个通用的 TinyBERT，该模型将作为学生模型进一步蒸馏的初始化。
  • 任务特定蒸馏阶段：首先进行数据增强，然后在增强的数据集上进行蒸馏，教师模型为经过调优的 BERT。在此阶段，TinyBERT 能够更好地吸收与特定任务相关的知识。

具体方法

Transformer Distillation

（Transformer-layer、Embedding-layer、Prediction-layer Distillation）

TinyBERT Learning（两阶段学习框架）