MixKD: Towards Efficient Distillation of Large-scale Language Models

1 INTRODUCTION . Recent language models ( LM ) pre-trained on large-scale unlabeled text corpora in a self-supervised manner have significantly advanced the state of the art across a wide variety of natural language processing ( NLP ) tasks ( Devlin et al. , 2018 ; Liu et al. , 2019c ; Yang et al. , 2019 ; Joshi et al. , 2020 ; Sun et al. , 2019b ; Clark et al. , 2020 ; Lewis et al. , 2019 ; Bao et al. , 2020 ) . After the LM pretraining stage , the resulting parameters can be fine-tuned to different downstream tasks . While these models have yielded impressive results , they typically have millions , if not billions , of parameters , and thus can be very expensive from storage and computational standpoints . Additionally , during deployment , such large models can require a lot of time to process even a single sample . In settings where computation may be limited ( e.g . mobile , edge devices ) , such characteristics may preclude such powerful models from deployment entirely . One promising strategy to compress and accelerate large-scale language models is knowledge distillation ( Zhao et al. , 2019 ; Tang et al. , 2019 ; Sun et al. , 2020 ) . The key idea is to train a smaller model ( a “ student ” ) to mimic the behavior of the larger , stronger-performing , but perhaps less practical model ( the “ teacher ” ) , thus achieving similar performance with a faster , lighter-weight model . A simple but powerful method of achieving this is to use the output probability logits produced by the teacher model as soft labels for training the student ( Hinton et al. , 2015 ) . With higher entropy than one-hot labels , these soft labels contain more information for the student model to learn from . ∗Equal contribution Previous efforts on distilling large-scale LMs mainly focus on designing better training objectives , such as matching intermediate representations ( Sun et al. , 2019a ; Mukherjee & Awadallah , 2019 ) , learning multiple tasks together ( Liu et al. , 2019a ) , or leveraging the distillation objective during the pre-training stage ( Jiao et al. , 2019 ; Sanh et al. , 2019 ) . However , much less effort has been made to enrich task-specific data , a potentially vital component of the knowledge distillation procedure . In particular , tasks with fewer data samples provide less opportunity for the student model to learn from the teacher . Even with a well-designed training objective , the student model is still prone to overfitting , despite effectively mimicking the teacher network on the available data . In response to these limitations , we propose improving the value of knowledge distillation by using data augmentation to generate additional samples from the available task-specific data . These augmented samples are further processed by the teacher network to produce additional soft labels , providing the student model more data to learn from a large-scale LM . Intuitively , this is akin to a student learning more from a teacher by asking more questions to further probe the teacher ’ s answers and thoughts . In particular , we demonstrate that mixup ( Zhang et al. , 2018 ) can significantly improve knowledge distillation ’ s effectiveness , and we show with a theoretical framework why this is the case . We call our framework MixKD . We conduct experiments on 6 GLUE datasets ( Wang et al. , 2019 ) across a variety of task types , demonstrating that MixKD significantly outperforms knowledge distillation ( Hinton et al. , 2015 ) and other previous methods that compress large-scale language models . In particular , we show that our method is especially effective when the number of available task data samples is small , substantially improving the potency of knowledge distillation . We also visualize representations learned with and without MixKD to show the value of interpolated distillation samples , perform a series of ablation and hyperparameter sensitivity studies , and demonstrate the superiority of MixKD over other BERT data augmentation strategies . 2 RELATED WORK . 2.1 MODEL COMPRESSION . Compressing large-scale language models , such as BERT , has attracted significant attention recently . Knowledge distillation has been demonstrated as an effective approach , which can be leveraged during both the pre-training and task-specific fine-tuning stages . Prior research efforts mainly focus on improving the training objectives to benefit the distillation process . Specifically , Turc et al . ( 2019 ) advocate that task-specific knowledge distillation can be improved by first pre-training the student model . It is shown by Clark et al . ( 2019 ) that a multi-task BERT model can be learned by distilling from multiple single-task teachers . Liu et al . ( 2019b ) propose learning a stronger student model by distilling knowledge from an ensemble of BERT models . Patient knowledge distillation ( PKD ) , introduced by Sun et al . ( 2019a ) , encourages the student model to mimic the teacher ’ s intermediate layers in addition to output logits . DistilBERT ( Sanh et al. , 2019 ) reduces the depth of BERT model by a factor of 2 via knowledge distillation during the pre-training stage . In this work , we evaluate MixKD on the case of task-specific knowledge distillation . Notably , it can be extended to the pre-training stage as well , which we leave for future work . Moreover , our method can be flexibly integrated with different KD training objectives ( described above ) to obtain even better results . However , we utilize the BERT-base model as the testbed in this paper without loss of generality . 2.2 DATA AUGMENTATION IN NLP . Data augmentation ( DA ) has been studied extensively in computer vision as a powerful technique to incorporate prior knowledge of invariances and improve the robustness of learned models ( Simard et al. , 1998 ; 2003 ; Krizhevsky et al. , 2012 ) . Recently , it has also been applied and shown effective on natural language data . Many approaches can be categorized as label-preserving transformations , which essentially produce neighbors around a training example that maintain its original label . For example , EDA ( Wei & Zou , 2019 ) propose using various rule-based operations such as synonym replacement , word insertion , swap or deletion to obtain augmented samples . Back-translation ( Yu et al. , 2018 ; Xie et al. , 2019 ) is another popular approach belonging to this type , which relies on pre-trained translation models . Additionally , methods based on paraphrase generation have also been leveraged from the data augmentation perspective ( Kumar et al. , 2019 ) . On the other hand , label-altering techniques like mixup ( Zhang et al. , 2018 ) have also been proposed for language ( Guo et al. , 2019 ; Chen et al. , 2020 ) , producing interpolated inputs and labels for the models predict . The proposed MixKD framework leverages the ability of mixup to facilitate the student learning more information from the teacher . It is worth noting that MixKD can be combined with arbitrary labelpreserving DA modules . Back-translation is employed as a special case here , and we believe other advanced label-preserving transformations developed in the future can benefit the MixKD approach as well . 2.3 MIXUP . Mixup ( Zhang et al. , 2018 ) is a popular data augmentation strategy to increase model generalizability and robustness by training on convex combinations of pairs of inputs and labels ( xi , yi ) and ( xj , yj ) : x′ = λxi + ( 1− λ ) xj ( 1 ) y′ = λyi + ( 1− λ ) yj ( 2 ) with λ ∈ [ 0 , 1 ] and ( x′ , y′ ) being the resulting virtual training example . This concept of interpolating samples was later generalized with Manifold mixup ( Verma et al. , 2019a ) and also found to be effective in semi-supervised learning settings ( Verma et al. , 2019b ; c ; Berthelot et al. , 2019b ; a ) . Other strategies include mixing together samples resulting from chaining together other augmentation techniques ( Hendrycks et al. , 2020 ) , or replacing linear interpolation with the cutting and pasting of patches ( Yun et al. , 2019 ) . 3 METHODOLOGY . 3.1 PRELIMINARIES . In NLP , an input sample i is often represented as a vector of tokens wi = { wi,1 , wi,2 , ... , wi , T } , with each token wi , t ∈ RV a one-hot vector often representing words ( but also possibly subwords , punctuation , or special tokens ) and V being the vocabulary size . These discrete tokens are then mapped to word embeddings xi = { xi,1 , xi,2 , ... , xi , T } , which serve as input to the machine learning model f . For supervised classification problems , a one-hot label yi ∈ RC indicates the ground-truth class of xi out of C possible classes . The parameters θ of f are optimized with some form of stochastic gradient descent so that the output of the model f ( xi ) ∈ RC is as close to yi as possible , with cross-entropy as the most common loss function : LMLE = − 1 n n∑ i yi · log ( f ( xi ) ) ( 3 ) where n is the number of samples , and · is the dot product . 3.2 KNOWLEDGE DISTILLATION FOR BERT . Consider two models f and g parameterized by θT and θS , respectively , with |θT | |θS | . Given enough training data and sufficient optimization , f is likely to yield better accuracy than g , due to higher modeling capacity , but may be too bulky or slow for certain applications . Being smaller in size , g is more likely to satisfy operational constraints , but its weaker performance can be seen as a disadvantage . To improve g , we can use the output prediction f ( xi ) on input xi as extra supervision for g to learn from , seeking to match g ( xi ) with f ( xi ) . Given these roles , we refer to g as the student model and f as the teacher model . While there are a number of recent large-scale language models driving the state of the art , we focus here on BERT ( Devlin et al. , 2018 ) models . Following Sun et al . ( 2019a ) , we use the notation BERTk to indicate a BERT model with k Transformer ( Vaswani et al. , 2017 ) layers . While powerful , BERT models also tend to be quite large ; for example , the default bert-base-uncased ( BERT12 ) has∼110M parameters . Reducing the number of layers ( e.g . using BERT3 ) makes such models significantly more portable and efficient , but at the expense of accuracy . With a knowledge distillation set-up , however , we aim to reduce this loss in performance . 3.3 MIXUP DATA AUGMENTATION FOR KNOWLEDGE DISTILLATION . While knowledge distillation can be a powerful technique , if the size of the available data is small , then the student has only limited opportunities to learn from the teacher . This may make it much harder for knowledge distillation to close the gap between student and teacher model performance . To correct this , we propose using data augmentation for knowledge distillation . While data augmentation ( Yu et al. , 2018 ; Xie et al. , 2019 ; Yun et al. , 2019 ; Kumar et al. , 2019 ; Hendrycks et al. , 2020 ; Shen et al. , 2020 ; Qu et al. , 2020 ) is a commonly used technique across machine learning for increasing training samples , robustness , and overall performance , a limited modeling capacity constrains the representations the student is capable of learning on its own . Instead , we propose using the augmented samples to further query the teacher model , whose large size often allows it to learn more powerful features . While many different data augmentation strategies have been proposed for NLP , we focus on mixup ( Zhang et al. , 2018 ) for generating additional samples to learn from the teacher . Mixup ’ s vicinal risk minimization tends to result in smoother decision boundaries and better generalization , while also being cheaper to compute than methods such as backtranslation ( Yu et al. , 2018 ; Xie et al. , 2019 ) . Mixup was initially proposed for continuous data , where interpolations between data points remain in-domain ; its efficacy was demonstrated primarily on image data , but examples in speech recognition and tabular data were also shown to demonstrate generality . Directly applying mixup to NLP is not quite as straightforward as it is for images , as language commonly consists of sentences of variable length , each comprised of discrete word tokens . Since performing mixup directly on the word tokens doesn ’ t result in valid language inputs , we instead perform mixup on the word embeddings at each time step xi , t ( Guo et al. , 2019 ) . This can be interpreted as a special case of Manifold mixup Verma et al . ( 2019a ) , where the mixing layer is set to the embedding layer . In other words , mixup samples are generated as : x′i , t = λxi , t + ( 1− λ ) xj , t ∀t ( 4 ) y′i = λyi + ( 1− λ ) yj ( 5 ) with λ ∈ [ 0 , 1 ] ; random sampling of λ from a Uniform or Beta distribution are common choices . Note that we index the augmented sample with i regardless of the value of λ . Sentence length variability can be mitigated by grouping mixup pairs by length . Alternatively , padding is a common technique for setting a consistent input length across samples ; thus , if x ( i ) contains more word tokens than x ( j ) , then the extra word embeddings are mixed up with zero paddings . We find this approach to be effective , while also being much simpler to implement . We query the teacher model with the generated mixup sample x′i , producing output prediction f ( x ′ i ) . The student is encouraged to imitate this prediction on the same input , by minimizing the objective : LTMKD = d ( f ( x′i ) , g ( x′i ) ) ( 6 ) where d ( · , · ) is a distance metric for distillation , with temperature-adjusted cross-entropy and mean square error ( MSE ) being common choices . Since we have the mixup samples already generated ( with an easy-to-generate interpolated pseudolabel y′i ) , we can also train the student model on these augmented data samples in the usual way , with a cross-entropy objective : LSM = − 1 n n∑ i y′i · log ( g ( x′i ) ) ( 7 ) Our final objective for MixKD is a sum of the original data cross-entropy loss , student cross-entropy loss on the mixup samples , and knowledge distillation from the teacher on the mixup samples : L = LMLE + αSMLSM + αTMKDLTMKD ( 8 ) where αSM and αTMKD are hyperparameters weighting the loss terms .

Summary: This paper applies mixup (Zhang et al., 2018) to augment training data to improve knowledge distillation in NLP tasks . Mixup was originally proposed to augment data for continuous data. To apply mixup to textual data, this paper applies mixup to the word/token embeddings instead of the tokens themselves. Some theoretical analysis has been done, and the experimental results show improved metrics over baseline methods such as DistillBERT.