Source-linked AI summary
Improving Multi-Task Deep Neural Networks via Knowledge Distillation for Natural Language Understanding
Xiaodong Liu, Pengcheng He, Weizhu Chen, Jianfeng Gao
TL;DR
The paper tackles the high serving cost of MT-DNN ensembles while seeking their generalization benefits. It extends knowledge distillation to multi-task learning by training a single MT-DNN student from task-specific ensemble teachers and their soft targets. The distilled model reaches 83.7% on GLUE and improves over the vanilla MT-DNN on 7 of 9 tasks.
Problem
Ensembles improve NLU generalization but are prohibitively expensive to deploy, especially when built from large pre-trained MT-DNN models.
Method
Task-specific ensembles of MT-DNN teachers generate soft targets, which are combined with correct targets to train a single MT-DNN student through multi-task learning.
Results
7 of 9 GLUE tasks improve over vanilla MT-DNN, while the single-model GLUE score reaches 83.7%.
Takeaways & Limitations
The distilled MT-DNN retains nearly all ensemble improvements while keeping the same model size as vanilla MT-DNN.
Takeaways & Limitations
Future work is needed to improve the combination of soft and hard targets and to explore distillation using unlabeled data or equal-complexity teacher and student models.
Abstract
from arXiv · showhide
This paper explores the use of knowledge distillation to improve a Multi-Task Deep Neural Network (MT-DNN) (Liu et al., 2019) for learning text representations across multiple natural language understanding tasks. Although ensemble learning can improve model performance, serving an ensemble of large DNNs such as MT-DNN can be prohibitively expensive. Here we apply the knowledge distillation method (Hinton et al., 2015) in the multi-task learning setting. For each task, we train an ensemble of different MT-DNNs (teacher) that outperforms any single model, and then train a single MT-DNN (student) via multi-task learning to \emph{distill} knowledge from these ensemble teachers. We show that the distilled MT-DNN significantly outperforms the original MT-DNN on 7 out of 9 GLUE tasks, pushing the GLUE benchmark (single model) to 83.7\% (1.5\% absolute improvement\footnote{ Based on the GLUE leaderboard at https://gluebenchmark.com/leaderboard as of April 1, 2019.}). The code and pre-trained models will be made publicly available at https://github.com/namisan/mt-dnn.
1 Introduction
The paper addresses the deployment cost of ensemble NLU models by extending knowledge distillation to MT-DNN multi-task learning. A single distilled MT-DNN improves performance across GLUE tasks while retaining deployability.
- Ensemble learning improves generalization and state-of-the-art NLU performance but is prohibitively expensive to deploy.Large pre-trained models make ensemble serving especially costly.
- Knowledge distillation transfers knowledge from large models to a lighter single model without significant performance loss.
- The method trains task-specific MT-DNN ensembles as teachers, then uses their soft targets and correct targets to train one MT-DNN student across tasks.The student is trained through multi-task learning on the teachers’ offline outputs.
- 83.7% on GLUE represents a 3.2% absolute improvement over BERT and a 1.5% improvement over the previous state-of-the-art model.The comparison is based on the GLUE leaderboard as of April 1, 2019.
- The distilled MT-DNN outperforms the vanilla MT-DNN on 7 of 9 GLUE tasks, including tasks without teachers.
2 MT-DNN
MT-DNN learns shared contextual representations with task-specific output layers. Its training combines BERT initialization, multi-task optimization, and task-specific fine-tuning for varied NLU tasks.
- MT-DNN shares lower representation layers across tasks while using task-specific top layers for prediction.
- The lexicon encoder maps tokens to embeddings by combining word, segment, and positional information.
- The Transformer encoder converts input embeddings into shared contextual representations using a multilayer bidirectional architecture.These representations are trained by multi-task objectives.
- Task-specific output layers support classification, text generation, and relevance ranking through suitable task-specific decoders or predictors.
- For text classification, the [CLS] contextual embedding is passed to a task-specific logistic regression with softmax.
- MT-DNN initializes shared layers from pre-trained BERT and learns shared and task-specific parameters with mini-batch stochastic gradient descent.
3 Knowledge Distillation
The proposed distillation procedure trains ensembles as task-specific teachers, averages their predictions into soft targets, and trains one MT-DNN student using multi-task objectives. Soft targets provide richer supervision for transferring teacher generalization.
- For selected tasks with labeled data, ensembles of MT-DNN models are trained as teachers before student training begins.
- The distillation procedure follows MT-DNN multi-task training with BERT-initialized shared layers, mini-batches, gradient computation, and parameter updates.
- Each teacher produces soft targets by averaging class probabilities from its component networks.
- Soft targets encode relative class probabilities, can provide more information per example than hard targets, and reduce gradient variance between samples.
- When correct targets are available, combining soft and hard targets can significantly improve distilled-model performance.The experiments found no significant difference from changing the relative weights of the two loss terms.
- The student is trained across tasks with a loss combining cross-entropy on correct targets and cross-entropy on teacher soft targets.
4 Experiments
The experiments evaluate MT-DNNKD against established baselines and ensemble teachers on GLUE, showing that a single distilled model achieves strong benchmark performance and transfers gains beyond teacher-covered tasks.
- Experimental Setup: The evaluation covers nine GLUE NLU tasks and compares MT-DNNKD with BERT, STILT, Snorkel MeTaL, and MT-DNN.The study also includes an ablation examining the relative contribution of knowledge distillation for multi-task learning.
- GLUE Main Results: 83.7% establishes MT-DNNKD as the best performer on the GLUE leaderboard, ahead of the second-best model ALICE by 0.5%.The comparison uses official GLUE test datasets and results recorded April 1, 2019.
- GLUE Main Results: 65.4% vs. 61.5% on CoLA and 85.1% vs. 75.5% on RTE are the largest reported per-task improvements over MT-DNN.The paper relates these larger gains to the relatively small amounts of in-domain data for both tasks.
- Ablation Study: MT-DNNKD retains nearly all of the improvement achieved by the ensemble teachers while remaining simpler than those ensembles.The ablation compares the distilled student with MT-DNN and task-specific MT-DNN-ensemble models on development sets.
- Ablation Study: MT-DNNKD improves MRPC, CoLA, and STS-B even though their ensemble models are not used as teachers.On these tasks without teachers, its performance is close to that of the corresponding ensemble models.
5 Conclusion
The distilled MT-DNN transfers knowledge from ensemble teachers into a single model, achieving new state-of-the-art results on 7 of 9 GLUE tasks while retaining nearly all ensemble improvements at vanilla model size. Future work targets improved target combination, unlabeled-data distillation, and distillation beyond model compression.
- 7 out of 9 GLUE tasks reached new state-of-the-art results with the distilled MT-DNN, including tasks without teachers.The single-model GLUE benchmark reached 83.7%.
- The distilled MT-DNN retains nearly all ensemble improvements while keeping the vanilla MT-DNN model size.
- Future work will investigate combining soft and hard targets, distilling teachers over unlabeled data, and improving models without reducing complexity.