摘要

基于预训练的语言模型（如BERT）显著提升了众多自然语言处理任务的性能。然而，预训练语言模型通常计算开销较大，难以在资源受限的设备上高效运行。为在保持模型准确率的同时加速推理并减小模型规模，本文提出一种专为基于Transformer的模型知识蒸馏（Knowledge Distillation, KD）设计的新型Transformer蒸馏方法。通过该方法，大型教师模型BERT所蕴含的丰富知识可被高效迁移至小型学生模型Tiny-BERT中。随后，本文进一步提出一种面向Tiny-BERT的新颖两阶段学习框架，在预训练阶段与特定任务学习阶段均实施Transformer蒸馏。该框架确保Tiny-BERT能够同时捕捉BERT中的通用领域知识与特定任务知识。实验结果表明，仅含4层的Tiny-BERT在GLUE基准测试中性能达到其教师模型BERTBASE的96.8%以上，模型体积缩小7.5倍，推理速度提升9.4倍。相较于当前最先进的4层BERT蒸馏方法，Tiny-BERT在参数量仅为其约28%、推理时间约为其31%的情况下，仍表现出显著更优的性能。此外，含6层的Tiny-BERT在性能上已与教师模型BERTBASE相当。