Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MKD: a Multi-Task Knowledge Distillation Approach for Pretrained Language Models

Linqing Liu, Huan Wang, Jimmy Lin, Richard Socher, Caiming Xiong

TL;DR

MKD introduces a general, multi-task knowledge distillation framework that distills across multiple NLP tasks from a multi-task refined teacher into a single, shared student. It is model-agnostic, supporting both Transformer and LSTM student architectures, and relies on logits-based distillation without requiring external large corpora. Empirical results on GLUE show MKD-LSTM achieving strong efficiency and competitive accuracy, while MKD-Transformer provides additional gains over prior KD methods. The work demonstrates the value of cross-task representation sharing for data-efficient, scalable distillation across diverse NLP tasks.

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.

MKD: a Multi-Task Knowledge Distillation Approach for Pretrained Language Models

TL;DR

MKD introduces a general, multi-task knowledge distillation framework that distills across multiple NLP tasks from a multi-task refined teacher into a single, shared student. It is model-agnostic, supporting both Transformer and LSTM student architectures, and relies on logits-based distillation without requiring external large corpora. Empirical results on GLUE show MKD-LSTM achieving strong efficiency and competitive accuracy, while MKD-Transformer provides additional gains over prior KD methods. The work demonstrates the value of cross-task representation sharing for data-efficient, scalable distillation across diverse NLP tasks.

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.

Paper Structure

This paper contains 17 sections, 8 equations, 2 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Model Architecture
  4. Multi-Task Refined Teacher Model
  5. LSTM-based Student Model
  6. Transformer-based Student Model
  7. Multi-Task Distillation
  8. An Intuitive Explanation
  9. Experimental Setup
  10. Datasets
  11. Implementation Details
  12. Methods and Baselines
  13. Result and Discussions
  14. Model Quality Analysis
  15. Ablation Study
  16. ...and 2 more sections

Figures (2)

  • Figure 1: The left figure represents task-specific KD. The distillation process needs to be performed for each different task. The right figure represents our proposed multi-task KD. The student model consists of shared layers and task-specific layers.
  • Figure 2: Architecture for the bi-attentive student neural network.