Improved Knowledge Distillation for Pre-trained Language Models via Knowledge Selection (2022)

Abstract: Knowledge distillation addresses the problem of transferring knowledge from a teacher model to a student model.In this process, we typically have multiple types of knowledge extracted from the teacher model.The problem is to make full use of them to train the student model.Our preliminary study shows that: (1) not all of the knowledge is necessary for learning a good student model, and (2) knowledge distillation can benefit from certain knowledge at different training steps.In response to these, we propose an actor-critic approach to selecting appropriate knowledge to transfer during the process of knowledge distillation.In addition, we offer a refinement of the training algorithm to ease the computational burden.Experimental results on the GLUE datasets show that our method outperforms several strong knowledge distillation baselines significantly.

## introduction pre-trained language models (plms) have significantly advanced state-of-the-art on various natural language processing tasks, such as sentiment analysis @xcite ), text classification @xcite , and question answering @xcite . despite the remarkable results, plms have a large number of parameters which make them expensive for deployment @xcite . to solve this problem, recent works resort to knowledge distillation (kd) @xcite to compress and accelerate the plms @xcite @xcite . its key idea is to transfer the knowledge from a large plm (i.e., teacher model) into a lightweight model (i.e., student model) without a significant performance loss. in this process, we typically have multiple types of knowledge extracted from the teacher model, such as response knowledge, feature knowledge, and relation knowledge @xcite . recent works mainly focus on how the student model learns the transferred knowledge more efficiently, such as designing training schemes @xcite @xcite or enriching task-specific data @xcite . however, few works consider how to make full use of the multiple types of knowledge into the training of the student model. our preliminary study (see section 4.1) shows that not all of the knowledge is necessary for learning a good student model when conducting distillation with diverse knowledge. furthermore, inspired by dynamic kd @xcite , we assume that learning from certain knowledge at different training steps is beneficial to kd. we conduct probing experiments to verify this assumption in section 4.2. specifically, we repeat kd 200 times and randomly select knowledge at each training step. we find that different distilled student models have a distinct performance. for example, the best student model is nearly 10% higher than the worst one in terms of accuracy on the rte dataset. in addition, it notably exceeds the student model that learns from fixed knowledge. based on our preliminary study, we have the following suggestion: the distilled student model could achieve superior performance if it learns appropriate knowledge at each training step. this suggestion motivates us to investigate how to select the appropriate knowledge for the student model during the process of kd. based on the above findings, we propose an actor-critic approach to selecting appropriate knowledge to transfer at different training steps @xcite . this approach first uses an actor-critic algorithm to implement a knowledge selection module via a long-term reward optimization. this optimization can consider the influence of knowledge selection on future training steps. furthermore, we develop a multi-phase training approach that divides the distillation process into many phases and provides a particular reward at the phase end. compared to the usual actor-critic algorithm, it can ease the burden of computing rewards. after that, we perform kd by employing the trained knowledge selection module to select knowledge at each training step. we test the proposed approach on six glue datasets @xcite which involve question answering, sentiment analysis, and textual entailment. experimental results show that our method significantly outperforms the vanilla kd method @xcite and other competitive kd methods. notably, results also show our bert 6 (6layer bert) student model can achieve more than 98.5% of the performance of teacher bert base while keeping much fewer parameters (∼61%). in addition, we prove that when armed with data augmentation, our approach can yield further improvements and be superior to tinybert @xcite , leading to a strong baseline with the data augmentation. ## related work knowledge distillation (kd) @xcite is widely used to compress and accelerate the pretrained language models (plms) @xcite @xcite . its core idea is to transfer the knowledge from a large plm (i.e., teacher model) into a lightweight model (i.e., student model). recent works on kd could be classified into three groups. the first group focused on utilizing various knowledge from the teacher model to distill the student model. for example, multiple teacher models are leveraged to provide a more diverse and accurate knowledge for the student model @xcite . @xcite exploited knowledge from intermediate layers of the teacher model during distillation. moreover, weight distillation @xcite transfers the knowledge in parameters of the teacher model to the student model. the second group tended to design effective strategies to facilitate the student model to learn the knowledge from different types of knowledge, such as enriching the task-specific data @xcite and the two-stage learning framework @xcite . the third group that has attracted less attention generally explored appropriate data and teacher models for student models in the kd. for example, @xcite proposed a data selection strategy that only selects some vital data for the student model according to its competency. @xcite attempted to adjust the weights to the multiple teacher models during distillation. different from these methods, in this paper, we are concerned with how to select appropriate knowledge for student models during the process of kd. to this end, we train an effective knowledge selection model via an actor-critic algorithm and employ it to select appropriate knowledge to transfer at each training step. ## preliminary study due to its unique information and character, we assume that each type of knowledge has a different impact on kd results. based on this, we conduct preliminary studies to probe the relationship between knowledge and kd. ## conclusion in this paper, we focus on making full use of the multiple types of knowledge into the distilling of the student model. we have proposed an actorcritic approach to selecting appropriate knowledge to transfer during the process of knowledge distillation. to ease the burden of computing rewards during training, we propose a multi-phase training approach and an exploration reward. our experiments on glue datasets show that our method significantly outperforms several strong knowledge distillation baselines.