[Paper Review] MKD: a Multi-Task Knowledge Distillation Approach for Pretrained Language Models

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

Topic Modeling45 references19 citations

TL;DR

This paper proposes MKD, a multi-task knowledge distillation framework that jointly distills a lightweight student model from a multi-task BERT teacher across multiple NLP tasks. By leveraging shared representations from diverse tasks, MKD achieves state-of-the-art performance with faster inference—outperforming BERT-PKD and TinyBERT in speed while matching or exceeding their accuracy using only downstream task data, not external corpora.

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.

Motivation & Objective

• To address the high computational cost and inference latency of large pretrained language models.
• To overcome the limitation of task-specific knowledge distillation, which requires retraining for each new task.
• To develop a general, model-agnostic distillation framework that works across diverse student architectures.
• To improve student model performance and inference speed by leveraging multi-task learning during distillation.
• To reduce reliance on large-scale external text corpora during distillation, enhancing data availability and privacy compliance.

Proposed method

• The student model is jointly distilled from a multi-task BERT teacher across multiple NLP tasks, enabling shared representation learning.
• The distillation process uses knowledge from multiple tasks simultaneously, improving the student's generalization and robustness.
• The framework is model-agnostic, allowing application to both Transformer-based and LSTM-based student models.
• The method uses cross-task knowledge transfer without requiring external training data, relying only on downstream task data.
• A bi-attention mechanism is incorporated into the LSTM-based student to enhance sequence modeling capacity.
• WordPiece tokenization is used to reduce vocabulary size and improve rare-word handling.

Experimental results

Research questions

• RQ1 Can multi-task distillation improve the performance and generalization of distilled student models across diverse NLP tasks?
• RQ2 Does joint distillation from multiple tasks lead to faster inference and better efficiency compared to task-specific distillation?
• RQ3 Can a distillation framework avoid reliance on large external text corpora while maintaining competitive performance?
• RQ4 How does the inclusion of attention mechanisms affect the performance of lightweight LSTM-based student models?
• RQ5 What is the impact of different tokenization strategies on distillation effectiveness?

Key findings

• MKD-LSTM achieves faster inference (2.93s) than Distilled BiLSTM (1.36s) while maintaining higher accuracy, demonstrating superior efficiency and performance.
• Compared to BERT-PKD and TinyBERT, MKD achieves comparable or better performance without using external training data.
• The multi-task distilled LSTM model outperforms single-task distillation by 1.4 points on RTE when trained on all tasks, showing the benefit of cross-task knowledge sharing.
• The bi-attention mechanism in the LSTM student significantly improves performance over a simple BiLSTM, with gains of 2.2–6.1 points across six datasets.
• WordPiece tokenization reduces vocabulary size and improves rare-word handling, leading to better overall performance than word-level tokenization.
• MKD achieves state-of-the-art results on seven datasets across four tasks, outperforming existing distillation methods in both accuracy and inference speed.