[论文解读] Distilling BERT into Simple Neural Networks with Unlabeled Transfer Data

Subhabrata Mukherjee, Ahmed Hassan Awadallah|arXiv (Cornell University)|Oct 4, 2019

Natural Language Processing Techniques参考文献 28被引用 9

一句话总结

本文提出使用领域内未标注的迁移数据，将大型预训练 BERT 模型蒸馏到基于 RNN 的轻量级学生网络中，实现与教师模型相当或更优的性能。该方法实现了高达 26 倍的参数压缩和 51 倍的延迟加速，在仅需有限标注数据的情况下，跨 41 种语言实现了 BERT F1 分数的 95%。

ABSTRACT

Recent advances in pre-training huge models on large amounts of text through self supervision have obtained state-of-the-art results in various natural language processing tasks. However, these huge and expensive models are difficult to use in practise for downstream tasks. Some recent efforts use knowledge distillation to compress these models. However, we see a gap between the performance of the smaller student models as compared to that of the large teacher. In this work, we leverage large amounts of in-domain unlabeled transfer data in addition to a limited amount of labeled training instances to bridge this gap for distilling BERT. We show that simple RNN based student models even with hard distillation can perform at par with the huge teachers given the transfer set. The student performance can be further improved with soft distillation and leveraging teacher intermediate representations. We show that our student models can compress the huge teacher by up to 26x while still matching or even marginally exceeding the teacher performance in low-resource settings with small amount of labeled data. Additionally, for the multilingual extension of this work with XtremeDistil (Mukherjee and Hassan Awadallah, 2020), we demonstrate massive distillation of multilingual BERT-like teacher models by upto 35x in terms of parameter compression and 51x in terms of latency speedup for batch inference while retaining 95% of its F1-score for NER over 41 languages.

研究动机与目标

• 解决在低资源和生产环境中部署大型计算成本高昂的预训练语言模型（如 BERT）的挑战。
• 通过利用大量领域内未标注的迁移数据，提升小规模学生模型在低资源环境下的性能。
• 探究简单的基于 RNN 的学生架构是否可通过知识蒸馏匹配或超越大型 BERT 教师模型的性能。
• 研究训练策略、蒸馏目标以及教师模型容量对学生模型性能的影响。

提出的方法

• 采用两阶段训练策略：首先在未标注数据上优化表示损失，然后在标注数据上联合优化 logits 损失和交叉熵损失，并逐步解冻网络层。
• 使用硬蒸馏方法，将教师模型预测的类别标签作为学生模型的目标。
• 通过使用教师模型输出的温度缩放软标签来增强性能，以改善泛化能力。
• 利用教师网络的中间隐藏表示，引导学生模型的学习，超越最终的 logits 输出。
• 在未对未标注数据进行预训练的情况下训练学生模型（基于 BiLSTM），而是依赖教师模型的知识蒸馏。
• 结合标注数据的交叉熵损失与教师输出和隐藏状态的蒸馏损失，对学生模型进行优化。

实验结果

研究问题

• RQ1 当使用知识蒸馏和领域内未标注数据训练时，简单的基于 RNN 的学生模型能否达到与大型 BERT 教师模型相当的性能？
• RQ2 在低资源环境中，大规模未标注迁移数据的引入如何影响学生模型的性能？
• RQ3 哪种训练策略与损失组合能为小型学生模型带来最佳的蒸馏性能？
• RQ4 与硬蒸馏相比，软蒸馏或中间表示蒸馏是否能进一步提升学生模型的性能？
• RQ5 当学生模型大小固定时，教师模型的大小如何影响蒸馏性能？

主要发现

• 简单的基于 BiLSTM 的学生模型在高达 26 倍参数压缩的情况下，性能与 BERT Base 和 BERT Large 相当。
• 在标注数据有限的低资源环境中，学生模型的性能略高于教师模型。
• 软蒸馏和中间表示蒸馏进一步提升了学生模型的性能，超越了硬蒸馏的效果。
• 该方法在多语言 BERT 类似模型上实现了高达 35 倍的参数压缩和 51 倍的延迟加速，同时在 41 种语言上保持了 95% 的 F1 分数。
• 与联合优化相比，分阶段训练并逐步解冻层的策略在蒸馏性能上表现最佳。
• 当学生模型大小达到饱和后，更大的教师模型仅带来边际性能提升。