Annealing Knowledge Distillation (2021)

Abstract: Significant memory and computational requirements of large deep neural networks restricts their application on edge devices. Knowledge distillation (KD) is a prominent model compression technique for deep neural networks in which the knowledge of a trained large teacher model is transferred to a smaller student model. The success of knowledge distillation is mainly attributed to its training objective function, which exploits the soft-target information (also known as “dark knowledge”) besides the given regular hard labels in a training set. However, it is shown in the literature that the larger the gap between the teacher and the student networks, the more difficult is their training using knowledge distillation. To address this shortcoming, we propose an improved knowledge distillation method (called Annealing-KD) by feeding the rich information provided by teacher’s soft-targets incrementally and more efficiently. Our Annealing-KD technique is based on a gradual transition over annealed soft-targets generated by the teacher at different temperatures in an iterative process; and therefore, the student is trained to follow the annealed teacher output in a step-by-step manner. This paper includes theoretical and empirical evidence as well as practical experiments to support the effectiveness of our Annealing-KD method. We did a comprehensive set of experiments on different tasks such as image classification (CIFAR-10 and 100) and NLP language inference with BERT-based models on the GLUE benchmark and consistently got superior results.

## introduction despite the great success of deep neural networks in many challenging tasks such as natural language processing @xcite , computer vision @xcite , and speech processing @xcite , these state-of-the-art networks are usually heavy to be deployed on edge devices with limited computational power @xcite . a case in point is the bert model @xcite which can be comprised of more than a hundred million parameters. the problem of network over-parameterization and expensive computational complexity of deep networks can be addressed by neural model compression. there are abundant of neural model compression techniques in the literature @xcite @xcite , among which knowledge distillation (kd) is one of the most prominent techniques @xcite . kd is tailored a lot to serve different applications and different network architectures @xcite . for instance, patient kd @xcite , tiny-bert @xcite , and mobilebert @xcite are designed particularly for distilling the knowledge of bert-based teachers to a smaller student. the success of kd is mainly attributed to its training objective function, which exploits the softtarget information (also known as "dark knowledge") besides the given regular hard labels in the training set @xcite . previous studies in the literature @xcite show that when the gap between the student and teacher models increases, training models with kd becomes more difficult. we refer to this problem as kd's capacity gap problem in this paper. for example, @xcite show that if we gradually increase the capacity of the teacher, first the performance of student model improves for a while, but after a certain point, it starts to drop. therefore, although increasing the capacity of a teacher network usually boosts its performance, it does not necessarily lead to a better teacher for the student network in kd. in other words, it would be more difficult for kd to transfer the knowledge of this enhanced teacher to the student. a similar scenario happens when originally the gap between the teacher and student network is large. @xcite proposed their takd solution to this problem which makes the kd process more smooth by filling the gap between the teacher and student networks using an intermediate auxiliary network (referred to as "teacher assistant"). the size of this ta network is between the size of the student and the teacher; and it is trained by the teacher first. then, the student is trained using kd when the ta network is playing the role of its teacher. this way, the training gap (between the teacher and the student) would be less significant compared to the original kd. however, takd suffers from the high computational complexity demand since it requires training the ta network separately. moreover, the training error of the ta network can be propagated to the student during the kd training process. in this paper, we want to solve the kd capacity gap problem from a different perspective. we propose our annealing-kd technique to bridges the gap between the student and teacher models by introducing a new kd loss with a dynamic temperature term. this way, annealing-kd is able to transfer the knowledge of the teacher smoothly to the student model via a gradual transition over soft-labels generated by the teacher at different temperatures. we can summarize the contributions of this paper in the following: ## method: annealing knowledge distillation in this section, we describe our annealing-kd technique and show the rationale behind it. first, we start by formulating the problem and visualizing our technique using an example for a better presentation. then, we use vc-dimension theory to understand why our technique improves knowledge distillation. we wrap up this section by visualizing the loss landscape of annealing kd for a resnet network in order to investigate the impact of our method on the kd loss function. kd defines a two-objective loss function (i.e. the l kd and l ce terms in equation 1 ) to minimize the distance between student predictions and soft labels and hard labels simultaneously. without adding to the computational needs of the kd algorithm, our annealing-kd model breaks the kd training into two stages: stage i, gradually training the student to mimic the teacher using our annealing-kd loss l annealing kd ; stage ii, fine-tuning the student with hard labels using l ce . we can define the loss function of our method as following. therefore, our annealing-kd bridges the gap between the student and teacher models by introducing the dynamic temperature term (that is the annealing function φ(t )) in the stage i of training. this way our annealing-kd method is able to smoothly transfer the teacher's knowledge to the student model via a gradual transition over softlabels generated by the teacher at different temperatures. to summarize, our annealing-kd technique is different from kd in following aspects: algorithm 1 explains the proposed method in more detail. in this section, we proposed an approach to alleviate the gap between the teacher and student models as well as reducing the sharpness of the kd loss function. in our model, instead of pushing the student network to learn a complex teacher function from scratch, we start training the student from a softened version of the teacher and we gradually move toward the original teacher outputs through our annealing process. ## experiments in this section, we describe the experimental evaluation of our proposed annealing kd method. we evaluate our technique on both image classification and natural language inference tasks. in all of our experiments, we compare the annealing kd results with takd, standard kd, and training student without kd results. ## discussion in image classification experiments, the improvement gap between the annealing kd results and the other baselines in cifar-100 experiments is larger than cifar-10 ones. we can observe similar conditions for the nlp experiments between bert-small and distilroberta students (the performance gap of bert-small is larger). in both of these cases, the problem for the student was more difficult. cifar-100 dataset is more complex than cifar-10 dataset. so the teacher has learned a more complex function that should be transferred to the student. in nlp experiments, on the other hand, the tasks are the same but bertsmall student has a smaller capacity in compare with distilroberta. therefore the problem is more difficult for bert-small. from this observation, we can conclude, whenever the gap between the teacher and student is larger, the annealing kd performs better than the other baselines and lever-age the acquired knowledge by the teacher to train the student. ## conclusion and future work in this work, we discussed that the difference between the capacity of the teacher and student models in knowledge distillation may hamper its performance. on the other hand, in most cases, larger neural networks can be trained better and get more accurate results. if we consider better teachers can train better students, then larger teachers with better accuracy would be more favourable for knowledge distillation training. in this paper, we proposed an improved knowledge distillation method called annealing kd to alleviate this problem and leverage the knowledge acquired by more complex teachers to guide the small student models better during their training. this happened by feeding the rich information provided by the teacher's soft-targets incrementally and more efficiently. our annealing-kd technique was based on a gradual transition over annealed soft-targets generated by the teacher at different temperatures in an iterative process; and therefore, the student was trained to follow the annealed teacher output in a step-by-step manner.