Tutoring Helps Students Learn Better: Improving Knowledge Distillation for BERT with Tutor Network (2022)

Abstract: Pre-trained language models have achieved remarkable successes in natural language processing tasks, coming at the cost of increasing model size. To address this issue, knowledge distillation (KD) has been widely applied to compress language models. However, typical KD approaches for language models have overlooked the difficulty of training examples, suffering from incorrect teacher prediction transfer and sub-efficient training. In this paper, we propose a novel KD framework, Tutor-KD, which improves the distillation effectiveness by controlling the difficulty of training examples during pre-training. We introduce a tutor network that generates samples that are easy for the teacher but difficult for the student, with training on a carefully designed policy gradient method. Experimental results show that Tutor-KD significantly and consistently outperforms the state-of-the-art KD methods with variously sized student models on the GLUE benchmark, demonstrating that the tutor can effectively generate training examples for the student.

## introduction pre-trained language models (plms) have achieved great success in extensive natural language processing (nlp) tasks by learning generalized language representations from large text corpora @xcite @xcite @xcite . bert @xcite and its variants @xcite @xcite , have shown significant performance improvement on natural language understanding tasks through learning bidirectional contextualized representations. however, due to the high memory footprints and computational costs, these models are challenging to be used in resource-constrained situations, such as mobile devices. kd @xcite is a promising approach for compressing pre-trained transformers, by transferring knowledge from a large source network (i.e., a teacher) to a small target network (i.e., a student). previous studies have shown that the performance of the student can be significantly improved by learning informative learning signals, such as output probabilities @xcite , intermediate feature representations, and attention probabilities @xcite @xcite from the pre-trained teacher. despite these compelling results, typical kd approaches used in language modeling have two key limitations arising from ignoring the difficulty of training examples. first, since the student mimics learning signals from the teacher regardless of whether it is right or wrong, incorrect teacher predictions from kd for overly difficult examples can be transferred to the student. second, with the ignorance of the difficulty of training examples, the performance of the language models on downstream tasks can be significantly affected @xcite , which may lead to sub-efficient kd. in this paper, we propose tutor-kd, a novel kd framework for plms to address the aforementioned limitations. our key idea is to introduce a tutor network that controls the difficulty of training examples during kd. specifically, the tutor generates training examples, which are easy for the teacher but difficult for the student, by replacing the masked tokens with corrupted tokens via masked language modeling (mlm). to accurately identify the difficulty of the tutor-generated training samples, we propose a novel method to train the tutor based on a policy gradient with carefully designed rewards. the tutor network is therefore optimized to generate tokens with relatively lower teacher losses (i.e., more accurate predictions) and consequently prevent overly difficult samples from being generated. simultaneously, the student trained with generated examples can be benefited to learn more effectively due to the increased difficulty. we conduct extensive experiments on downstream nlp tasks using various sizes of student models. our experimental results show that the proposed approach significantly improves the language model distillation performance. specifically, the 6-layer model with 768 hidden dimensions distilled from bert-base outperforms the teacher model on the general language understanding evaluation (glue) benchmark @xcite . moreover, our framework shows notable effectiveness for extremely small-sized student models that are designed 7.5× smaller than bert-base. finally, we demonstrate that our designed rewards can generate effective samples for both the teacher and the student. as a summary of our main contributions: ## methodology in this section, we introduce our proposed kd framework with its implementation details. figure 2 illustrates the overall architecture of the proposed framework. the core idea is to define the respective difficulty level of training examples for the teacher and the student while maximizing the rewards for training the tutor. ## analysis to better understand the main advantages of our tutor-kd, we conduct several analysis experiments. we perform all experiments on a 6-layer student model with 9m parameters using the same corpora. ## conclusion we have presented tutor-kd, a novel kd framework that controls the difficulty of training examples. with the carefully designed rewards on the policy gradient method, our tutor network is trained to generate training examples that are easy for the teacher but difficult for the student. through extensive experiments, we have verified that tutor-kd significantly improves kd effectiveness. specifically, our student models outperform the state-ofthe-art kd baselines with various sizes of models on the glue benchmark. furthermore, we have demonstrated that our tutor network can generate effective samples for training student models, resulting in consistent performance improvements.