BERT Learns to Teach: Knowledge Distillation with Meta Learning (2022)

Abstract: We present Knowledge Distillation with Meta Learning (MetaDistil), a simple yet effective alternative to traditional knowledge distillation (KD) methods where the teacher model is fixed during training. We show the teacher network can learn to better transfer knowledge to the student network (i.e., learning to teach ) with the feedback from the performance of the distilled student network in a meta learning framework. Moreover, we introduce a pilot update mechanism to improve the alignment between the inner-learner and meta-learner in meta learning algorithms that focus on an improved inner-learner. Experiments on various benchmarks show that MetaDistil can yield significant improvements compared with traditional KD algorithms and is less sensitive to the choice of different student capacity and hyperparameters, facilitating the use of KD on different tasks and models.

## introduction with the prevalence of large neural networks with millions or billions of parameters, model compression is gaining prominence for facilitating efficient, eco-friendly deployment for machine learning applications. among techniques for compression, knowledge distillation (kd) @xcite has shown effectiveness in both computer vision and natural language processing tasks @xcite @xcite @xcite @xcite @xcite @xcite @xcite . previous works often train a large model as the "teacher"; then they fix the teacher and train a "student" model to mimic the behavior of the teacher, in order to transfer the knowledge from the teacher to the student. however, this paradigm has the following drawbacks: (1) the teacher is unaware of the student's capacity. recent studies in pedagogy suggest student-centered learning, which considers students' characteristics and learning capability, has shown effectiveness improving students' performance @xcite . however, in conventional knowledge distillation, the student passively accepts knowledge from the teacher, without regard for the student model's learning capability and performance. recent works @xcite introduce student-aware distillation by jointly training the teacher and the student with task-specific objectives. however, there is still space for improvement since: (2) the teacher is not optimized for distillation. in previous works, the teacher is often trained to optimize its own inference performance. however, the teacher is not aware of the need to transfer its knowledge to a student and thus usually does so suboptimally. a real-world analogy is that a phd student may have enough knowledge to solve problems themselves, but requires additional teaching training to qualify as a professor. to address these two drawbacks, we propose knowledge distillation with meta learning (metadistil), a new teacher-student distillation framework using meta learning @xcite to exploit feedback about the student's learning progress to improve the teacher's knowledge transfer ability throughout the distillation process. on the basis of previous formulations of bi-level optimization based meta learning @xcite , we propose a new mechanism called pilot update that aligns the learning of the bi-level learners (i.e., the teacher and the student). we illustrate the workflow of metadistil in figure 1 . the teacher in metadistil is trainable, which enables the teacher to adjust to its student network and also improves its training batches quiz samples l @xmath0." motivated by the idea of studentcentered learning, we allow the teacher to adjust its output based on the performance of the student model on a "quiz set," which is a separate reserved data split from the original training set. for each training step, we first copy the student s to s and update s by a common knowledge distillation loss. we call this process a "teaching experiment." in this way, we can obtain an experimental student s that can be quizzed. then, we sample from the quiz set, and calculate the loss of s on these samples. we use this loss as a feedback signal to meta-update the teacher by calculating second derivatives and performing gradient descent @xcite . finally, we discard the experimental subject s and use the updated teacher to distill into the student s on the same training batches. the use of meta learning allows the teacher model to receive feedback from the student in a completely differentiable way. we provide a simple and intuitive approach to explicitly optimize the teacher using the student's quiz performance as a proxy. to test the effectiveness of metadistil, we conduct extensive experiments on text and image classification tasks. metadistil outperforms knowledge distillation by a large margin, verifying the effectiveness and versatility of our method. also, our method achieves state-of-the-art performance compressing bert @xcite on the glue benchmark @xcite and shows competitive results compressing resnet @xcite and vgg @xcite on cifar-100 @xcite . additionally, we design experiments to analyze and explain the improvement. ablation studies show the effectiveness of our proposed pilot update and dynamic distillation. also, compared to conventional kd, metadistil is more robust to different student capacity and hyperparameters, which is probably because of its ability to adjust the parameters of the teacher model. ## knowledge distillation with meta learning an overview of metadistil is presented in figure 1 . metadistil includes two major components. first, the meta update enables the teacher model to receive the student model's feedback on the distillation process, allowing the teacher model to "learn to teach" and provide distillation signals that are more suitable for the student model's current capacity. the pilot update mechanism ensures a finergrained match between the student model and the meta-updated teacher model. ## discussion in this paper, we present metadistil, a knowledge distillation algorithm powered by meta learning that explicitly optimizes the teacher network to better transfer its knowledge to the student network. the extensive experiments verify the effectiveness and robustness of metadistil.