[论文解读] MKD: a Multi-Task Knowledge Distillation Approach for Pretrained Language Models

Linqing Liu, Huan Wang+3 more|arXiv (Cornell University)|Nov 9, 2019

Topic Modeling参考文献 45被引用 19

一句话总结

本文提出MKD，一种多任务知识蒸馏框架，通过在多个NLP任务上联合蒸馏来自多任务BERT教师模型的轻量级学生模型。通过利用不同任务间的共享表征，MKD在实现更快推理速度的同时达到最先进性能——在速度上优于BERT-PKD和TinyBERT，且仅使用下游任务数据（无需外部语料）即可匹配或超越其准确率。

ABSTRACT

Pretrained language models have led to significant performance gains in many NLP tasks. However, the intensive computing resources to train such models remain an issue. Knowledge distillation alleviates this problem by learning a light-weight student model. So far the distillation approaches are all task-specific. In this paper, we explore knowledge distillation under the multi-task learning setting. The student is jointly distilled across different tasks. It acquires more general representation capacity through multi-tasking distillation and can be further fine-tuned to improve the model in the target domain. Unlike other BERT distillation methods which specifically designed for Transformer-based architectures, we provide a general learning framework. Our approach is model agnostic and can be easily applied on different future teacher model architectures. We evaluate our approach on a Transformer-based and LSTM based student model. Compared to a strong, similarly LSTM-based approach, we achieve better quality under the same computational constraints. Compared to the present state of the art, we reach comparable results with much faster inference speed.

研究动机与目标

• 为解决大规模预训练语言模型带来的高计算成本和推理延迟问题。
• 克服特定任务知识蒸馏的局限性，即每新增一个任务都需要重新训练。
• 开发一种通用的、与模型无关的蒸馏框架，适用于多种学生模型架构。
• 通过在蒸馏过程中引入多任务学习，提升学生模型的性能和推理速度。
• 减少蒸馏过程中对外部大规模文本语料的依赖，提升数据可用性并符合隐私合规要求。

提出的方法

• 学生模型在多个NLP任务上联合从多任务BERT教师模型中进行蒸馏，实现共享表征学习。
• 蒸馏过程同时利用多个任务的知识，提升学生模型的泛化能力和鲁棒性。
• 该框架与模型无关，可应用于基于Transformer和基于LSTM的学生模型。
• 该方法通过跨任务知识迁移实现，无需外部训练数据，仅依赖下游任务数据。
• 在基于LSTM的学生模型中引入双注意力机制，以增强序列建模能力。
• 采用WordPiece分词方式，以减小词汇表规模并改善罕见词的处理能力。

实验结果

研究问题

• RQ1 多任务蒸馏能否提升学生模型在多样化NLP任务上的性能与泛化能力？
• RQ2 与特定任务蒸馏相比，从多个任务中联合蒸馏是否能带来更快的推理速度和更高的效率？
• RQ3 是否可以设计一种蒸馏框架，在不依赖大规模外部语料的情况下仍保持竞争力的性能？
• RQ4 注意力机制的引入对轻量级LSTM学生模型的性能有何影响？
• RQ5 不同分词策略对蒸馏效果有何影响？

主要发现

• MKD-LSTM的推理速度（2.93秒）快于Distilled BiLSTM（1.36秒），同时保持更高准确率，证明其在效率与性能上的优越性。
• 与BERT-PKD和TinyBERT相比，MKD在不使用外部训练数据的情况下，实现了相当或更优的性能。
• 在所有任务上联合训练的多任务蒸馏LSTM模型在RTE任务上比单任务蒸馏高出1.4分，表明跨任务知识共享具有显著优势。
• LSTM学生模型中引入的双注意力机制显著提升了性能，相较于简单BiLSTM，在六个数据集上准确率提升2.2至6.1分。
• WordPiece分词方式减小了词汇表规模并改善了罕见词的处理，整体性能优于词级别分词。
• MKD在四个任务的七个数据集上达到最先进结果，在准确率和推理速度上均优于现有蒸馏方法。