[Paper Review] Improving Multi-Task Deep Neural Networks via Knowledge Distillation for Natural Language Understanding

Xiaodong Liu, Pengcheng He+2 more|arXiv (Cornell University)|Apr 20, 2019

Topic Modeling20 references161 citations

TL;DR

The paper applies knowledge distillation to a multi-task DNN (MT-DNN) to transfer ensemble knowledge into a single model, achieving state-of-the-art GLUE scores.

ABSTRACT

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.

Motivation & Objective

• Motivate reducing deployment cost of ensemble MT-DNNs while retaining high NLU performance.
• Investigate whether knowledge distillation can transfer ensemble generalization to a single MT-DNN in a multi-task setting.
• Demonstrate improvements on GLUE by distilling multiple task-specific teachers into one student.
• Show robustness of the distilled model across tasks, including those without teachers.

Proposed method

• Train an ensemble of MT-DNNs (teachers) for selected tasks to produce soft targets.
• Compute soft targets by averaging ensemble predictions for each training sample.
• Train a single MT-DNN (student) with multi-task learning using both hard targets and soft targets from teachers.
• Optionally combine hard and soft targets with a weighted loss for tasks that have teachers.
• Fine-tune the distilled MT-DNN on each GLUE task after distillation.

Experimental results

Research questions

• RQ1 Can knowledge distillation from task-specific MT-DNN ensembles improve a single MT-DNN in a multi-task setting?
• RQ2 Does the distilled MT-DNN retain the gains of ensemble teachers and also benefit tasks without teachers?
• RQ3 What is the impact of distillation on GLUE performance compared to BERT-based and vanilla MT-DNN baselines?

Key findings

• MT-DNN KD outperforms vanilla MT-DNN on 7 of 9 GLUE tasks.
• MT-DNN KD achieves a GLUE score of 83.7% (single model), a 1.5% absolute improvement over the previous state of the art and 3.2% over BERT on the GLUE benchmark as of Apr 1, 2019.
• MT-DNN KD significantly improves CoLA and RTE tasks relative to MT-DNN.
• Distillation transfers the generalization ability of ensemble teachers to the student, with the student retaining most improvements of the ensemble.
• Even for tasks without teachers, MT-DNN KD shows notable gains over MT-DNN and approaches ensemble performance on some tasks.
• Ablation studies show that MT-DNN KD benefits both teacher-provided tasks and unassisted tasks, indicating effective knowledge transfer.