T iny BERT : Distilling BERT for Natural Language Understanding (2020)

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” TinyBERT. Then, we introduce a new two-stage learning framework for TinyBERT, which performs Transformer distillation at both the pre-training and task-specific learning stages. This framework ensures that TinyBERT can capture the general-domain as well as the task-specific knowledge in BERT. TinyBERT4 with 4 layers is empirically effective and achieves more than 96.8% the performance of its teacher BERT-Base on GLUE benchmark, while being 7.5x smaller and 9.4x faster on inference. TinyBERT4 is also significantly better than 4-layer state-of-the-art baselines on BERT distillation, with only ~28% parameters and ~31% inference time of them. Moreover, TinyBERT6 with 6 layers performs on-par with its teacher BERT-Base.

## introduction pre-training language models then fine-tuning on downstream tasks has become a new paradigm for natural language processing (nlp). pre-trained language models (plms), such as bert @xcite , xlnet @xcite , roberta @xcite , albert @xcite , t5 @xcite and electra @xcite , have achieved great success in many nlp tasks (e.g., the glue benchmark @xcite and the challenging multi-hop reasoning task @xcite ). however, plms usually have a large number of parameters and take long inference time, which are difficult to be deployed on edge devices such as mobile phones. recent studies @xcite @xcite demonstrate that there is redundancy in plms. therefore, it is crucial and feasible to reduce the computational overhead and model storage of plms while retaining their performances. there have been many model compression techniques @xcite proposed to accelerate deep model inference and reduce model size while maintaining accuracy. the most commonly used techniques include quantization @xcite , weights pruning @xcite , and knowledge distillation (kd) @xcite . in this paper, we focus on knowledge distillation, an idea originated from @xcite , in a teacher-student framework. kd aims to transfer the knowledge embedded in a large teacher network to a small student network where the student network is trained to reproduce the behaviors of the teacher network. based on the framework, we propose a novel distillation method specifically for the transformer-based models @xcite , and use bert as an example to investigate the method for large-scale plms. kd has been extensively studied in nlp @xcite as well as for pre-trained language models @xcite @xcite . the pre-training-then-fine-tuning paradigm firstly pretrains bert on a large-scale unsupervised text corpus, then fine-tunes it on task-specific dataset, which greatly increases the difficulty of bert distillation. therefore, it is required to design an effective kd strategy for both training stages. to build a competitive tinybert, we firstly propose a new transformer distillation method to distill the knowledge embedded in teacher bert. specifically, we design three types of loss functions to fit different representations from bert layers: 1) the output of the embedding layer; 2) the hidden states and attention matrices derived from the transformer layer; 3) the logits output by the prediction layer. the attention based fitting is inspired by the recent findings @xcite that the attention weights learned by bert can capture substantial linguistic knowledge, and it thus encourages the linguistic knowledge can be well transferred from teacher bert to student tiny-bert. then, we propose a novel two-stage learning framework including the general distillation and the task-specific distillation, as illustrated in figure 1 . at general distillation stage, the original bert without fine-tuning acts as the teacher model. the student tinybert mimics the teacher's behavior through the proposed transformer distillation on general-domain corpus. after that, we obtain a general tinybert that is used as the initialization of student model for the further distillation. at the task-specific distillation stage, we first do the data augmentation, then perform the distillation on the augmented dataset using the fine-tuned bert as the teacher model. it should be pointed out that both the two stages are essential to improve the performance and generalization capability of tinybert. the main contributions of this work are as follows: 1) we propose a new transformer distillation method to encourage that the linguistic knowledge encoded in teacher bert can be adequately transferred to tinybert; 2) we propose a novel two-stage learning framework with performing the proposed transformer distillation at both the pretraining and fine-tuning stages, which ensures that tinybert can absorb both the general-domain and task-specific knowledge of the teacher bert. 3) we show in the experiments that our tinybert 4 can achieve more than 96.8% the performance of teacher bert base on glue tasks, while having much fewer parameters (∼13.3%) and less inference time (∼10.6%), and significantly outperforms other state-of-the-art baselines with 4 layers on bert distillation; 4) we also show that a 6-layer tinybert 6 can perform on-par with the teacher bert base on glue. ## preliminaries in this section, we describe the formulation of transformer @xcite and knowledge distillation @xcite . our proposed transformer distillation is a specially designed kd method for transformer-based models. ## method in this section, we propose a novel distillation method for transformer-based models, and present a two-stage learning framework for our model distilled from bert, which is called tinybert. ## experiments in this section, we evaluate the effectiveness and efficiency of tinybert on a variety of tasks with different model settings.