LightPAFF: A Two-Stage Distillation Framework for Pre-training and Fine-tuning

1 INTRODUCTION . Recently , the pre-training and fine-tuning frameworks ( Devlin et al. , 2018 ; Radford et al. , 2018 ; 2019 ; Yang et al. , 2019b ; Liu et al. , 2019 ; Lample & Conneau , 2019 ; Song et al. , 2019 ) have achieved great success in NLP by learning universal language representations on large-scale language corpus and transferring knowledge from the pre-trained model to downstream tasks . Based on the carefully designed unsupervised pre-training strategy , e.g. , masked language modeling ( Devlin et al. , 2018 ) , causal language modeling ( Radford et al. , 2018 ; 2019 ) , permutation language modeling ( Yang et al. , 2019b ) and masked sequence to sequence modeling ( Song et al. , 2019 ) , the pre-trained model can extract semantic information , understand meanings , generate sentences , even capture common-sense knowledge from large scale corpus . Some of the most representative pretrained models are BERT ( Devlin et al. , 2018 ) , XLNet ( Yang et al. , 2019b ) for language understanding tasks , GPT/GPT-2 ( Radford et al. , 2018 ; 2019 ) for language modeling tasks , and XLM ( Lample & Conneau , 2019 ) , MASS ( Song et al. , 2019 ) for sequence to sequence based language generation tasks . However , pre-trained models are usually of huge model parameter size , e.g. , BERT has more than 300 million parameters while GPT-2 has more than 1.5 billion parameters , which brings challenges in memory cost and inference latency when deploying the pre-trained models online or on mobile devices . One straightforward solution is to use smaller model for pre-training and fine-tuning , which will cause much accuracy degradation that can not be acceptable for online usage . Previous works ( Yang et al. , 2019a ; Tang et al. , 2019 ; Sun et al. , 2019 ) have tried knowledge distillation to reduce the number of parameters for BERT . However , they just distill the knowledge of fine-tuned BERT to the task-specific small model , still with big accuracy drop compared with the original finetuned BERT model . Meanwhile , most works focus on compressing BERT while few works study on reducing the size of other models like GPT-2 , XLNet , XLM , MASS . Knowledge distillation ( Bucilu et al. , 2006 ; Hinton et al. , 2015 ; Yu et al. , 2013 ) is an effective and practical way to transfer the knowledge from a cumbersome teacher model to a lightweight student model . However , unlike the previous works ( Yang et al. , 2019a ; Tang et al. , 2019 ; Sun et al. , 2019 ) that only leverage the knowledge in the fine-tuned task-specified model , the knowledge in the pretrained model is also helpful and even more general for language understanding and generation . In this work , we propose LightPAFF , a Lightweight Pre-training And Fine-tuning Framework that incorporates knowledge distillation in both pre-training and fine-tuning stages and transfers knowledge from a pre-trained big model into a smaller model . As a result , LightPAFF can greatly reduce the number of model parameters and thus memory cost and inference latency , without losing much of accuracy . As a general framework , LightPAFF can be applied on a variety of pre-trained models ( e.g. , BERT , GPT/GPT-2 , XLNet , RoBERTa , XLM , MASS ) and downstream tasks ( e.g. , language understanding/modeling/generation ) . In this paper , we carefully choose BERT , GPT-2 and MASS as three study cases , considering the model diversity that BERT leverages masked language modeling for pre-training and is for downstream language understanding tasks , GPT-2 leverages causal language modeling for pre-training and is for language modeling and generation tasks , while MASS leverages masked sequence to sequence modeling for pre-training and is for sequence to sequence based language generation tasks . We formulate LightPAFF on the pre-training and fine-tuning tasks of BERT , GPT-2 and MASS uniformly , and conduct experiments on three language understanding tasks ( SST2 , QQP , polyphone disambiguation ) for BERT , three language modeling tasks ( WikiText-2 , PTB , WikiText-103 ) for GPT-2 , and three sequence to sequence based language generation tasks ( WMT Zh-En , En-De , En-Fr translation ) for MASS . Experimental results demonstrate that LightPAFF reduces the model parameter of BERT/GPT-2/MASS by nearly 5× and improves the inference speed by 5× ∼7× , while achieves similar accuracy with the original BERT/GPT-2/MASS model 1 . 2 TWO-STAGE DISTILLATION FRAMEWORK . Different from existing knowledge distillation methods for pre-trained models ( Yang et al. , 2019a ; Tang et al. , 2019 ; Sun et al. , 2019 ) , which only perform distillation in the fine-tuning stage , we perform distillation not only in the fine-tuning stage but also in the pre-training stage to ensure the accuracy of the final student model . The basic idea of our two-stage distillation is shown in Figure 1 , and we call this Light-weight Pre-training And Fine-tuning Framework as LightPAFF for short . LightPAFF runs in four steps : 1 . Pre-train the big teacher model using the pre-training dataset . 2 . Fine-tune the big teacher model using the dataset of a downstream task . 3 . Distill the pre-trained big teacher model and obtain a light-weight pre-trained model using the pre-training dataset . 4 . Fine-tune the light-weight pre-trained model for the downstream task by using the dataset of the task and distilling the fine-tuned big teacher model . As can be seen , the first and second steps are exactly the same as previous pre-training methods such as BERT ( Devlin et al. , 2018 ) and GPT-2 ( Radford et al. , 2019 ) . The third step performs distillation in the pre-training stage , and the fourth step performs distillation in the fine-tuning stage . We adopt a general formulation for knowledge distillation in LightPAFF : L ( θ ) = ∑ ( x , y ) ( 1− λ ) ·MLE ( x , y ; θ ) + λ ·KL ( P ( x ; θT ) , P ( x ; θ ) ) , ( 1 ) 1We will release the codes and models once accepted . where MLE ( x , y ; θ ) is the maximal likelihood loss of the student model θ over a training data pair ( x , y ) , KL ( P ( x ; θT ) , P ( x ; θ ) ) is the KL divergence between the predicted probability distribution on the data x of the teacher model θT and the distribution of the student model θ , and λ is a hyperparameter to trade off the two loss terms . LightPAFF is a general framework and can be applied to many pre-training fine-tuning tasks . In the following sections , we choose three representative models considering their diversity and code availability : BERT ( Devlin et al. , 2018 ) , which is popular and effective for language understanding tasks such as GLUE benchmark , GPT-2 ( Radford et al. , 2019 ) , which is popular and effective for language modeling tasks , and MASS ( Song et al. , 2019 ) , which is popular and effective for sequence to sequence generation tasks such as machine translation . Detailed descriptions of two-stage knowledge distillation for BERT , GPT-2 and MASS can be found in Appendix ( Section 1 ) . We conduct experimental study of LightPAFF on BERT in Section 3 , GPT-2 in Section 4 and MASS in Section 5 , and then conduct some general analyses of LightPAFF in Section 6 . 3 LIGHTPAFF FOR BERT . We follow Equation 1 to conduct knowledge distillation for BERT in both pre-training and finetuning stages . The BERT teacher model provides the probability distribution of the masked tokens to the student model during pre-training distillation2 , while provides the probability distribution of the class label to the student model during fine-tuning distillation . In the following sections , we describe the experimental setting and results of LightPAFF for BERT . 3.1 EXPERIMENTAL SETTING . Datasets For BERT pre-training , we use BooksCorpus ( 800M words ) and Wikipedia ( 2500M words ) for English BERT as used in Devlin et al . ( 2018 ) , and NewsCorpus ( 2000M words ) and Wikipedia ( 300M words ) for Chinese BERT . For downstream language understanding tasks , we choose SST-2 ( Socher et al. , 2013 ) and QQP ( Iyer et al. , 2017 ) on English and polyphone disambiguation ( PolyDis ) ( Zhang et al. , 2002 ) on Chinese . Additionally , we add 1M unlabeled data for SST-2 , 2M for QQP and 1M for PolyDis in the fine-tuning stage , where we simply use the teacher model to generate probability distribution to train the student model . More descriptions about datasets can be found in Appendix ( Section 2 ) . Model Configuration We denote L , H , A as the number of layers , hidden size and the number of attention heads in Transformer . For BERT teacher model , we choose the pre-trained BERT-base model ( L=12 , H=768 , A=12 ) with 110M parameters released by Devlin et al . ( 2018 ) . For BERT student model , we use the configuration ( L=3 , H=512 , A=8 ) with 25M parameters for English and 20M parameters for Chinese , where the parameter difference is due to the different vocabulary sizes . Training and Inference The hyperparameter λ in Equation 1 is set as 1.0 in pre-training according to the performance of pre-training on a held out validation set , and the λ in fine-tuning is chosen according to the validation performance on each downstream task . We pre-train BERT on 8 NVIDIA V100 GPUs , each with 16GB memory , and with batch size of 400 sentences and max sequence length of 128 tokens . We fine-tune on 4 NVIDIA V100 GPUs with batch size of 32 sentences . 3.2 RESULTS . Accuracy of LightPAFF The results of LightPAFF on BERT tasks ( SST-2 , QQP and PolyDis ) are shown in Table 1 . It can be seen that LightPAFF achieves similar accuracy ( maintains nearly 99.5 % of accuracy ) with the BERT teacher model , but reduces the model size by nearly 5× ( from 110M to 25M or 20M ) . Compared with the model with the same size of parameters ( Transformer small and BERT small ) and the previous works ( Patient KD and Distilled BiLSTM ) that only use 2BERT leverages masked language modeling and next sentence prediction for pre-training . Considering some works ( Yang et al. , 2019b ; Liu et al. , 2019 ; Joshi et al. , 2019 ) have achieved good results without next sentence prediction , we only distill BERT on the masked language modeling task in pre-training . knowledge distillation in the fine-tuning stage , LightPAFF achieves much better accuracy3 . The trade-off between accuracy and size of the student model can be found in Appendix ( Section 3 ) . Since BERT leverages additional unlabeled data for distillation , we also study the effect of unlabeled data in Appendix ( Section 4 ) . Inference Speedup of LightPAFF We further measure the inference latency of LightPAFF on both GPU and CPU . As shown in Table 2 , LightPAFF achieves about 6× speedup over the big model on GPU and about 7× speedup on CPU in BERT task , demonstrating the advantages of LightPAFF for inference speedup . 4 LIGHTPAFF FOR GPT-2 . Since GPT-2 ( Radford et al. , 2019 ) uses language modeling for both pre-training and fine-tuning , the knowledge distillation is similar in the two stages , where the GPT-2 teacher model provides the probability distribution of each token in a sentence to the student model during distillation . In the following sections , we describe the experimental setting and results of LightPAFF for GPT-2 .

Summary: The paper sets to solve the problem that many language models such as GPT-2 despite having achieved great success in their language tasks are too big to deploy in applications. They propose LightPAFF, a framework that allows to transfer knowledge from a big teacher model to a lightweight student model, thus solving the deployability issue at hand. They conduct experiments showing LightPAFF achieves a 5x reduction in the number of parameters and a 5x-7x improvement on the inference speed, while roughly preserving performance.