[Paper Review] TinyBERT: Distilling BERT for Natural Language Understanding

Xiaoqi Jiao, Yichun Yin+6 more|arXiv (Cornell University)|Sep 23, 2019

Topic Modeling53 references136 citations

TL;DR

TinyBERT uses a novel Transformer distillation method and a two-stage learning framework to compress BERT into smaller, faster models with competitive GLUE performance.

ABSTRACT

Language model pre-training, such as BERT, has significantly improved the performances of many natural language processing tasks. However, pre-trained language models are usually computationally expensive, so it is difficult to efficiently execute them on resource-restricted devices. To accelerate inference and reduce model size while maintaining accuracy, we first propose a novel Transformer distillation method that is specially designed for knowledge distillation (KD) of the Transformer-based models. By leveraging this new KD method, the plenty of knowledge encoded in a large teacher BERT can be effectively transferred to a small student Tiny-BERT. Then, we introduce a new two-stage learning framework for TinyBERT, which performs Transformer distillation at both the pretraining and task-specific learning stages. This framework ensures that TinyBERT can capture he general-domain as well as the task-specific knowledge in BERT. TinyBERT with 4 layers is empirically effective and achieves more than 96.8% the performance of its teacher BERTBASE on GLUE benchmark, while being 7.5x smaller and 9.4x faster on inference. TinyBERT with 4 layers is also significantly better than 4-layer state-of-the-art baselines on BERT distillation, with only about 28% parameters and about 31% inference time of them. Moreover, TinyBERT with 6 layers performs on-par with its teacher BERTBASE.

Motivation & Objective

• Motivate reducing computational overhead of pre-trained language models for edge devices while preserving accuracy.
• Introduce a Transformer-specific knowledge distillation method to transfer teacher BERT knowledge to a smaller student.
• Propose a two-stage learning framework (general distillation and task-specific distillation) to capture general-domain and task-specific knowledge.
• Demonstrate that TinyBERT achieves substantial speedups and parameter reductions with competitive performance on GLUE.

Proposed method

• Propose a Transformer distillation loss with three components: embedding-layer distillation, attention-based distillation, and hidden-state distillation, plus prediction-layer distillation.
• Use a layer-mapping function g(m) to align student and teacher layers for distillation.
• Train in two stages: general distillation on a large general corpus using un-finetuned BERT as teacher, then task-specific distillation with augmented data using a fine-tuned BERT as teacher.
• In task-specific distillation, perform data augmentation combining BERT predictions and GloVe similarities to expand training data.
• Evaluate on GLUE benchmarks comparing TinyBERT (4 and 6 layers) to prior KD baselines and BERT BASE as teacher.

Experimental results

Research questions

• RQ1 Can Transformer-specific knowledge distillation effectively transfer knowledge from BERT to a smaller student?
• RQ2 Does a two-stage distillation framework (pre-training distillation and task-specific distillation) improve TinyBERT performance over single-stage approaches?
• RQ3 How do embedding-, attention-, and hidden-state-level distillation contribute to final performance?
• RQ4 What are the trade-offs in parameters, FLOPs, and inference speed when compressing BERT to TinyBERT?
• RQ5 How close can a 4- or 6-layer TinyBERT get to BERT BASE on GLUE tasks?

Key findings

• TinyBERT 4 reaches over 96.8% of BERT BASE performance on GLUE while being ~7.5x smaller and ~9.4x faster at inference.
• TinyBERT 6 matches BERT BASE performance on GLUE.
• TinyBERT 4 outperforms 4-layer KD baselines (BERT-PKD, DistilBERT 4) by at least 4.4% on average.
• TinyBERT 4 achieves strong results despite only ~13.3% of BERT BASE parameters and ~10.6% of inference time.
• Two-stage learning (general distillation plus task-specific distillation with data augmentation) is crucial for performance gains.
• Attention-based distillation yields substantial gains, and combining it with hidden-state distillation is complementary.