TinyBERT: Distilling BERT for Natural Language Understanding

1 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 ( Devlin et al. , 2018 ) , XLNet ( Yang et al. , 2019 ) , RoBERTa ( Liu et al. , 2019 ) and SpanBERT ( Joshi et al. , 2019 ) , have achieved great success in many NLP tasks ( e.g. , the GLUE benchmark ( Wang et al. , 2018 ) and the challenging multi-hop reasoning task ( Ding et al. , 2019 ) ) . However , PLMs usually have an extremely large number of parameters and need long inference time , which are difficult to be deployed on edge devices such as mobile phones . Moreover , recent studies ( Kovaleva et al. , 2019 ) also demonstrate that there is redundancy in PLMs . Therefore , it is crucial and possible to reduce the computational overhead and model storage of PLMs while keeping their performances . There has been many model compression techniques ( Han et al. , 2015a ) proposed to accelerate deep model inference and reduce model size while maintaining accuracy . The most commonly used techniques include quantization ( Gong et al. , 2014 ) , weights pruning ( Han et al. , 2015b ) , and knowledge distillation ( KD ) ( Romero et al. , 2014 ) . In this paper we focus on knowledge distillation , an idea proposed by Hinton et al . ( 2015 ) in a teacher-student framework . KD aims to transfer the knowledge embedded in a large teacher network to a small student network . 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 Transformer-based models ( Vaswani et al. , 2017 ) , and use BERT as an example to investigate the KD methods for large scale PLMs . KD has been extensively studied in NLP ( Kim & Rush , 2016 ; Hu et al. , 2018 ) , while designing KD methods for BERT has been less explored . The pre-training-then-fine-tuning paradigm firstly pre-trains BERT on a large scale unsupervised text corpus , then fine-tunes it on task-specific dataset , which greatly increases the difficulty of BERT distillation . Thus we are required to design an ef- 1Our code and models will be made publicly available . fective KD strategy for both 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 several 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 ( Clark et al. , 2019 ) that the attention weights learned by BERT can capture substantial linguistic knowledge , which encourages that the linguistic knowledge can be well transferred from teacher BERT to student TinyBERT . However , it is ignored in existing KD methods of BERT , such as Distilled BiLSTMSOFT ( Tang et al. , 2019 ) , BERT-PKD ( Sun et al. , 2019 ) and DistilBERT2 . Then , we propose a novel two-stage learning framework including the general distillation and the task-specific distillation . At the general distillation stage , the original BERT without fine-tuning acts as the teacher model . The student TinyBERT learns to mimic the teacher ’ s behavior by executing the proposed Transformer distillation on the large scale corpus from general domain . We obtain a general TinyBERT that can be fine-tuned for various downstream tasks . At the task-specific distillation stage , we perform the data augmentation to provide more task-specific data for teacherstudent learning , and then re-execute the Transformer distillation on the augmented data . Both the two stages are essential to improve the performance and generalization capability of TinyBERT . A detailed comparison between the proposed method and other existing methods is summarized in Table 1 . The Transformer distillation and two-stage learning framework are two key ideas of the proposed method . 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 well transferred to TinyBERT . 2 ) We propose a novel two-stage learning framework with performing the proposed Transformer distillation at both the pre-training and fine-tuning stages , which ensures that TinyBERT can capture both the general-domain and task-specific knowledge of the teacher BERT . 3 ) We show experimentally that our TinyBERT can achieve more than 96 % the performance of teacher BERTBASE 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 on BERT distillation . 2 PRELIMINARIES . We firstly describe the formulation of Transformer ( Vaswani et al. , 2017 ) and Knowledge Distillation ( Hinton et al. , 2015 ) . Our proposed Transformer distillation is a specially designed KD method for Transformer-based models . 2.1 TRANSFORMER LAYER . Most of the recent pre-trained language models ( e.g. , BERT , XLNet and RoBERTa ) are built with Transformer layers , which can capture long-term dependencies between input tokens by selfattention mechanism . Specifically , a standard Transformer layer includes two main sub-layers : multi-head attention ( MHA ) and fully connected feed-forward network ( FFN ) . 2https : //medium.com/huggingface/distilbert-8cf3380435b5 3The student learns from the [ CLS ] ( a special classification token of BERT ) hidden states of the teacher . 4The output of pre-training tasks ( such as dynamic masking ) is used as the supervision signal . Multi-Head Attention ( MHA ) . The calculation of attention function depends on the three components of queries , keys and values , which are denoted as matrices Q , K and V respectively . The attention function can be formulated as follows : A = QKT√ dk , ( 1 ) Attention ( Q , K , V ) = softmax ( A ) V , ( 2 ) where dk is the dimension of keys and acts as a scaling factor , A is the attention matrix calculated from the compatibility of Q and K by dot-product operation . The final function output is calculated as a weighted sum of values V , and the weight is computed by applying softmax ( ) operation on the each column of matrix A . According to Clark et al . ( 2019 ) , the attention matrices in BERT can capture substantial linguistic knowledge , and thus play an essential role in our proposed distillation method . Multi-head attention is defined by concatenating the attention heads from different representation subspaces as follows : MultiHead ( Q , K , V ) = Concat ( head1 , . . . , headh ) W , ( 3 ) where h is the number of attention heads , and headi denotes the i-th attention head , which is calculated by the Attention ( ) function with inputs from different representation subspaces , the matrix W acts as a linear transformation . Position-wise Feed-Forward Network ( FFN ) . Transformer layer also contains a fully connected feed-forward network , which is formulated as follows : FNN ( x ) = max ( 0 , xW1 + b1 ) W2 + b2 . ( 4 ) We can see that the FFN contains two linear transformations and one ReLU activation . 2.2 KNOWLEDGE DISTILLATION . KD aims to transfer the knowledge of a large teacher network T to a small student network S. The student network is trained to mimic the behaviors of teacher networks . Let fT and fS represent the behavior functions of teacher and student networks , respectively . The behavior function targets at transforming network inputs to some informative representations , and it can be defined as the output of any layer in the network . In the context of Transformer distillation , the output of MHA layer or FFN layer , or some intermediate representations ( such as the attention matrix A ) can be used as behavior function . Formally , KD can be modeled as minimizing the following objective function : LKD = ∑ x∈X L ( fS ( x ) , fT ( x ) ) , ( 5 ) where L ( · ) is a loss function that evaluates the difference between teacher and student networks , x is the text input and X denotes the training dataset . Thus the key research problem becomes how to define effective behavior functions and loss functions . Different from previous KD methods , we also need to consider how to perform KD at the pre-training stage of BERT in addition to the task-specific training stage . 3 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 . 3.1 TRANSFORMER DISTILLATION . The proposed Transformer distillation is a specially designed KD method for Transformer networks . Figure 1 displays an overview of the proposed KD method . In this work , both the student and teacher networks are built with Transformer layers . For a clear illustration , we firstly formulate the problem before introducing our method . Problem Formulation . Assuming that the student model has M Transformer layers and teacher model has N Transformer layers , we choose M layers from the teacher model for the Transformerlayer distillation . The function n = g ( m ) is used as a mapping function from student layers to teacher layers , which means that the m-th layer of student model learns the information from the nth layer of teacher model . The embedding-layer distillation and the prediction-layer distillation are also considered . We set 0 to be the index of embedding layer andM+1 to be the index of prediction layer , and the corresponding layer mappings are defined as 0 = g ( 0 ) and N + 1 = g ( M + 1 ) respectively . The effect of the choice of different mapping functions on the performances will be studied in the experiment section . Formally , the student can acquire knowledge from the teacher by minimizing the following objective : Lmodel = M+1∑ m=0 λmLlayer ( Sm , Tg ( m ) ) , ( 6 ) where Llayer refers to the loss function of a given model layer ( e.g. , Transformer layer or embedding layer ) and λm is the hyper-parameter that represents the importance of the m-th layer ’ s distillation . Transformer-layer Distillation . The proposed Transformer-layer distillation includes the attention based distillation and hidden states based distillation , which is shown in Figure 1 ( b ) . The attention based distillation is motivated by the recent findings that attention weights learned by BERT can capture rich linguistic knowledge ( Clark et al. , 2019 ) . This kind of linguistic knowledge includes the syntax and coreference information , which is essential for natural language understanding . Thus we propose the attention based distillation to encourage that the linguistic knowledge can be transferred from teacher BERT to student TinyBERT . Specifically , the student learns to fit the matrices of multihead attention in the teacher network , and the objective is defined as : Lattn = 1 h h∑ i=1 MSE ( ASi , A T i ) , ( 7 ) where h is the number of attention heads , Ai ∈ Rl×l refers to the attention matrix corresponding to the i-th head of teacher or student , l is the input text length , and MSE ( ) means the mean squared error loss function . In this work , the ( unnormalized ) attention matrix Ai is used as the fitting target instead of its softmax output softmax ( Ai ) , since our experiments show that the former setting has a faster convergence rate and better performances . In addition to the attention based distillation , we also distill the knowledge from the output of Transformer layer ( as shown in Figure 1 ( b ) ) , and the objective is as follows : Lhidn = MSE ( HSWh , HT ) , ( 8 ) where the matrices HS ∈ Rl×d′ and HT ∈ Rl×d refer to the hidden states of student and teacher networks respectively , which are calculated by Equation 4 . The scalar values d and d′ denote the hidden sizes of teacher and student models , and d′ is often smaller than d to obtain a smaller student network . The matrix Wh ∈ Rd ′×d is a learnable linear transformation , which transforms the hidden states of student network into the same space as the teacher network ’ s states . Embedding-layer Distillation . We also perform embedding-layer distillation , which is similar to the hidden states based distillation and formulated as : Lembd = MSE ( ESWe , ET ) , ( 9 ) where the matrices ES and HT refer to the embeddings of student and teacher networks , respectively . In this paper , they have the same shape as the hidden state matrices . The matrix We is a linear transformation playing a similar role as Wh . Prediction-Layer Distillation . In addition to imitating the behaviors of intermediate layers , we also use the knowledge distillation to fit the predictions of teacher model ( Hinton et al. , 2015 ) . Specifically , we penalize the soft cross-entropy loss between the student network ’ s logits against the teacher ’ s logits : Lpred = −softmax ( zT ) · log softmax ( zS/ t ) , ( 10 ) where zS and zT are the logits vectors predicted by the student and teacher respectively , log softmax ( ) means the log likelihood , t means the temperature value . In our experiment , we find that t = 1 performs well . Using the above distillation objectives ( i.e . Equations 7 , 8 , 9 and 10 ) , we can unify the distillation loss of the corresponding layers between the teacher and the student network : Llayer ( Sm , Tg ( m ) ) = Lembd ( S0 , T0 ) , m = 0 Lhidn ( Sm , Tg ( m ) ) + Lattn ( Sm , Tg ( m ) ) , M ≥ m > 0 Lpred ( SM+1 , TN+1 ) , m =M + 1 ( 11 ) In our experiments , we firstly perform intermediate layer distillation ( M ≥ m ≥ 0 ) , then perform the prediction-layer distillation ( m =M + 1 ) .

Summary: This paper proposes a new knowledge distillation method for BERT models. A number of modifications to the vanilla knowledge distillation method of Hinton et al (2015) are proposed. First, authors suggest adding L2 loss functions between alignment matrices, embedding layer values and prediction layer values. Second, authors propose run knowledge-distillation twice, once with the original pre-trained BERT model as teacher, and then again with task specific fine-tuned BERT as a new teacher. Third, authors emphasize the use of data augmentation for successful knowledge distillation. In Table 2, authors claim a significant lift across GLUE benchmarks with respect to other baseline methods with comparable model size.