Pea-KD: Parameter-efficient and accurate Knowledge Distillation

1 INTRODUCTION . How can we improve the accuracy of knowledge distillation ( KD ) with smaller number of parameters ? KD uses a well-trained large teacher model to train a smaller student model . Conventional KD method ( Hinton et al . ( 2006 ) ) trains the student model using the teacher model ’ s predictions as targets . That is , the student model uses not only the true labels ( hard distribution ) but also the teacher model ’ s predictions ( soft distribution ) as targets . Since better KD accuracy is directly linked to better model compression , improving KD accuracy is valuable and crucial . Naturally , there have been many studies and attempts to improve the accuracy of KD . Sun et al . ( 2019 ) introduced Patient KD , which utilizes not only the teacher model ’ s final output but also the intermediate outputs generated from the teacher ’ s layers . Jiao et al . ( 2019 ) applied additional KD in the pretraining step of the student model . However , existing KD methods share the limitation of students having lower model capacity compared to their teacher models , due to their smaller size . In addition , there is no proper initialization guide established for the student models , which especially becomes important when the student model is small . These limitations lower than desired levels of student model accuracy . In this paper , we propose Pea-KD ( Parameter-efficient and accurate Knowledge Distillation ) , a novel KD method designed especially for Transformer-based models ( Vaswani et al . ( 2017 ) ) that significantly improves the student model ’ s accuracy . Pea-KD is composed of two modules , Shuffled Parameter Sharing ( SPS ) and Pretraining with Teacher ’ s Predictions ( PTP ) . Pea-KD is based on the following two main ideas . 1 . We apply SPS in order to increase the effective model capacity of the student model without increasing the number of parameters . SPS has two steps : 1 ) stacking layers that share parameters and 2 ) shuffling the parameters between shared pairs of layers . Doing so increases the model ’ s effective capacity which enables the student to better replicate the teacher model ( details in Section 3.2 ) . 2 . We apply a pretraining task called PTP for the student . Through PTP , the student model learns general knowledge about the teacher and the task . With this additional pretraining , the student more efficiently acquires and utilizes the teacher ’ s knowledge during the actual KD process ( details in Section 3.3 ) . Throughout the paper we use Pea-KD applied on BERT model ( PeaBERT ) as an example to investigate our proposed approach . We summarize our main contributions as follows : • Novel framework for KD . We propose SPS and PTP , a parameter sharing method and a KD-specialized initialization method . These methods serve as a new framework for KD to significantly improve accuracy . • Performance . When tested on four widely used GLUE tasks , PeaBERT improves student ’ s accuracy by 4.4 % on average and up to 14.8 % maximum when compared to the original BERT model . PeaBERT also outperforms the existing state-of-the-art KD baselines by 3.5 % on average . • Generality . Our proposed method Pea-KD can be applied to any transformer-based models and classification tasks with small modifications . Then we conclude in Section 5 . 2 RELATED WORK . Pretrained Language Models . The framework of first pretraining language models and then finetuning for downstream tasks has now become the industry standard for Natural Language Processing ( NLP ) models . Pretrained language models , such as BERT ( Devlin et al . ( 2018 ) ) , XLNet ( Yang et al . ( 2019 ) ) , RoBERTa ( Liu et al . ( 2019 ) ) and ELMo ( Peters et al . ( 2018 ) ) prove how powerful pretrained language models can be . Specifically , BERT is a language model consisting of multiple transformer encoder layers . Transformers ( Vaswani et al . ( 2017 ) ) can capture long-term dependencies between input tokens by using a self-attention mechanism . Self-attention calculates an attention function using three components , query , key , and value , each denoted as matrices Q , K , and V. The attention function is defined as follows : A = QKT√ dk ( 1 ) Attention ( Q , K , V ) = softmax ( A ) V ( 2 ) It is known that through pretraining using Masked Language Modeling ( MLM ) and Next Sentence Prediction ( NSP ) , the attention matrices in BERT can capture substantial linguistic knowledge . BERT has achieved the state-of-the-art performance on a wide range of NLP tasks , such as the GLUE benchmark ( Wang et al . ( 2018 ) ) and SQuAD ( Rajpurkar et al . ( 2016 ) ) . However , these modern pretrained models are very large in size and contain millions of parameters , making them nearly impossible to apply on edge devices with limited amount of resources . Model Compression . As deep learning algorithms started getting adopted , implemented , and researched in diverse fields , high computation costs and memory shortage have started to become challenging factors . Especially in NLP , pretrained language models typically require a large set of parameters . This results in extensive cost of computation and memory . As such , Model Compression has now become an important task for deep learning . There have already been many attempts to tackle this problem , including quantization ( Gong et al . ( 2014 ) ) and weight pruning ( Han et al . ( 2015 ) ) . Two promising approaches are KD ( Hinton et al . ( 2015 ) ) and Parameter Sharing , which we focus on in this paper . Knowledge Distillation ( KD ) . As briefly covered in Section 1 , KD transfers knowledge from a well-trained and large teacher model to a smaller student model . KD uses the teacher models predictions on top of the true labels to train the student model . It is proven through many experiments that the student model learns to imitate the soft distribution of the teacher model ’ s predictions and ultimately performs better than learning solely from the original data . There have already been many attempts to compress BERT using KD . Patient Knowledge Distillation ( Sun et al . ( 2019 ) ) extracts knowledge not only from the final prediction of the teacher , but also from the intermediate layers . TinyBERT ( Jiao et al . ( 2019 ) ) uses a two-stage learning framework and applies KD in both pretraining and task-specific finetuning . DistilBERT ( Sanh et al . ( 2019 ) ) uses half of the layers of BERT-base model and applies KD during pretraining and finetuning of BERT . Insufficient capacity and the absence of a clear initialization guide are some of the existing KD method ’ s area of improvement . Parameter Sharing . Sharing parameters across different layers is a widely used idea for model compression . There have been several attempts to apply parameter sharing in transformer architecture and BERT model . However , the existing parameter sharing methods exhibit a large tradeoff between model performance and model size . They reduce the model ’ s size significantly but also suffers from a great loss in performance as a result . 3 PROPOSED METHODS . In the following , we provide an overview of the main challenges faced in KD and our methods to address them in Section 3.1 . We then discuss the precise procedures of SPS and PTP in Sections 3.2 and 3.3 . Lastly , we explain our final method , PeaBERT and the training details in Section 3.4 . 3.1 OVERVIEW . BERT-base model contains over 109 million parameters . Its extensive size makes model deployment often infeasible and computationally expensive in many cases , such as on mobile devices . As a result , industry practitioners commonly use a smaller version of BERT and apply KD . However , the existing KD methods face the following challenges : 1 . Insufficient model capacity of the student model . Since the student model contains fewer number of parameters than the teacher model , its model capacity is also lower . The smaller and simpler the student model gets , the gap between the student and the teacher grows , making it increasingly difficult for the student to replicate the teacher model ’ s accuracy . The limited capacity hinders the student model ’ s performance . How can we enlarge the student model ’ s capacity while maintaining the same number of parameters ? 2 . Absence of proper initial guide for the student model . There is no widely accepted and vetted guide to selecting the student ’ s initial state of the KD process . In most cases , a truncated version of pretrained BERT-base model is used . In reality , this hinders the student from reproducing the teacher ’ s results . Can we find a better method for the student ’ s KD initialization ? We propose the following main ideas to address these challenges : 1 . Shuffled Parameter Sharing ( SPS ) : increasing ‘ effective ’ model capacity of the student . SPS consists of two steps . Step 1 increases the student ’ s effective model capacity by stacking parameter-shared layers . Step 2 further improves the model capacity by shuffling the shared parameters . Shuffling enables more efficient use of the parameters by learning diverse set of information . As a result , the SPS-applied student model achieves much higher accuracy while still using the same number of parameters ( details in Section 3.2 ) . 2 . Pretraining with Teacher ’ s Predictions ( PTP ) : a novel pretraining task utilizing teacher ’ s predictions for student initialization . To address the limitation of the initial guide , we propose PTP , a novel pretraining method for the student by utilizing teacher model ’ s predictions . Through PTP , the student model pre-learns knowledge latent in the teacher ’ s softmax output which is hard to learn during the conventional KD process . This helps the student better acquire and utilize the teacher ’ s knowledge during the KD process ( details in Section 3.3 ) . The following subsections describe the procedures of SPS , PTP , and PeaBERT in detail . 3.2 SHUFFLED PARAMETER SHARING ( SPS ) . SPS improves student ’ s model capacity while using the same number of parameters , addressing the capacity limitations of a typical KD . SPS is composed of the following two steps . Step1 . Paired Parameter Sharing . We start with doubling the number of layers in the student model . We then share the parameters between the bottom half and the upper half of the model , as graphically represented in Figure 1a . Step2 . Shuffling . We shuffle the Query and Key parameters between the shared pairs . That is , for the shared upper half of layers we use the original Query parameters as Key parameters , and the original Key parameters as Query parameters . We call this architecture SPS , which is depicted in Figure 1b . For the 6-layer student case we slightly modify the architecture as shown in Figure 1c . We apply SPS on the top 3 layers only . Motivation behind step 1 : Increased model capacity due to a more complex structure . By applying step 1 , the student model now has more layers while maintaining the same number of parameters used . This leads to a higher model complexity , which we expect will increase model capacity and performance . Motivation behind step 2 : Increased parameter efficiency and additional regularization effect . The intuition behind shuffling becomes clear when we compare the model with and without step 2 . Let us consider the case depicted in Figure1b . In step 1 , the first layer ’ s Query parameter is used as a Query parameter in both the first layer and in the shared fourth layer . Since this parameter is used only as a Query parameter throughout the model , it will only learn the information relevant to Query . In step 2 , on the other hand , this parameter ’ s function changes due to shuffling . The first layer ’ s Query parameter is used as a Key parameter in the shared fourth layer . Thus it has the opportunity to learn the important features of both the Query and the Key , gaining a more diverse and wider breadth of knowledge . The shuffling mechanism allows the parameters to learn diverse information , thereby maximizing learning efficiency . We can also view this as a type of regularization or multi-task learning . Since these parameters need to learn the features of both the Query and the Key , this prevents overfitting to either one . Therefore , we expect the shuffling mechanism to improve the generalization ability of the model .

Summary: This paper proposes a distillation method for BERT. The work is based on two-fold main ideas. First, as the student model is usually smaller in the number of parameters, the model capacity is limited. The authors propose to stack the layers that share parameters to counter this limitation. Second, the authors argue that the initialization of the student model is crucial, so they propose an pre-training strategy for boosting the student's performance.