Pro- KD : Progressive Distillation by Following the Footsteps of the Teacher (2022)

Abstract: With the ever growing scale of neural models, knowledge distillation (KD) attracts more attention as a prominent tool for neural model compression. However, there are counter intuitive observations in the literature showing some challenging limitations of KD. A case in point is that the best performing checkpoint of the teacher might not necessarily be the best teacher for training the student in KD. Therefore, one important question would be how to find the best checkpoint of the teacher for distillation? Searching through the checkpoints of the teacher would be a very tedious and computationally expensive process, which we refer to as the checkpoint-search problem . Moreover, another observation is that larger teachers might not necessarily be better teachers in KD, which is referred to as the capacity-gap problem. To address these challenging problems, in this work, we introduce our progressive knowledge distillation (Pro-KD) technique which defines a smoother training path for the student by following the training footprints of the teacher instead of solely relying on distilling from a single mature fully-trained teacher. We demonstrate that our technique is quite effective in mitigating the capacity-gap problem and the checkpoint search problem. We evaluate our technique using a comprehensive set of experiments on different tasks such as image classification (CIFAR-10 and CIFAR-100), natural language understanding tasks of the GLUE benchmark, and question answering (SQuAD 1.1 and 2.0) using BERT-based models and consistently got superior results over state-of-the-art techniques.

## introduction knowledge distillation (kd) @xcite has gained a lot of attention in different deep learn- task model best #chpt epoch acc ∆ ↑ mrpc bert large (♠t) 6 88.22 -bert small (♣s) 3 85.29 +0.7 distilbert(♣s) 5 89.16 +0.8 sst-2 bert large (♠t) 1 92.89 -bert small (♣s) 7 88.76 +0.6 distilbert(♣s) 7 91.63 +0.4 qnli bert large (♠t) 6 92.4 -bert small (♣s) 2 87.08 +0.3 distilbert(♣s) 4 90.46 +0.2 ♠t:teacher, ♣s: student, #chpt: checkpoint number, acc: accuracy, ∆ ↑: accuracy improvement compared to the best checkpoint of the teacher kd adds a new loss term to the regular crossentropy classification loss. this new loss encourages the student model to mimic the output of a pre-trained teacher network. the teacher network is usually a higher capacity model which is able to learn the underlying function of the training data to a good extent. the output prediction of the teacher is called soft-target for the student model. in contrast to ground-truth labels coming from the training dataset which only carry the information about a single class, soft-targets can provide some information about the relative distribution of different classes for each training data. therefore, a pretrained teacher is able to provide some auxiliary signal besides the labels in the training dataset. kd has been investigated a lot in the literature. @xcite ; @xcite ; @xcite proposed a technique to improve kd by incorporating the intermediate layer matching in the kd loss. @xcite show a two-stage kd with intermediate layer mapping, attention distillation and embedding distillation for bertbased models. mate-kd @xcite and minimax-knn kd @xcite tailor data augmentation for kd, in which augmented samples are generated or selected based on maximum divergence loss between the student and teacher networks. @xcite propose a zero-shot kd technique in nlp in which the student does not need to access the teacher training data for its training. @xcite use kd for multi-task learning in natural language understanding. @xcite propose a sequencelevel kd solution for machine translation. @xcite introduces a collaborative training of students with different capacities with kd. although kd has been successful in many different deep learning tasks, it is subject to some special limitations as well. for example, it is shown in @xcite that, based on vc theory @xcite , the teacher capacity should not be too large in kd. similar observation is given by @xcite using empirical and theoretical justifications that kd will be less effective when the capacity gap between the teacher and student is large. this problem is referred to as the capacity-gap problem. @xcite proposed the takd solution to this problem by adding one or multiple intermediate teacher assistant (ta) networks to learn from the teacher using kd and then train the student using other distillation processes. however, training intermediate networks can be prohibitive in terms of adding to the training time, computational complexity and extra error propagation. the other case in point is that a fullytrained teacher might not be the best teacher for the student @xcite . in other words, an early-stopped teacher can be a better option for kd compared to a fully-trained one. this observation implies that we need to search through the checkpoints of the teacher to find the best model for the distillation process. however, this search can be very expensive especially when we deal with plms. we refer to this problem as the checkpoint-search problem. our investigations of this checkpoint search problem show the significance of this issue especially dealing with plms. in this work, we propose our pro-kd solution to tackle both the capacity-gap and checkpoint-search problems. in pro-kd, the student grows gradually with the growth of the teacher. we hypothesize that the training path of the teacher can be informative for the student and we should not disregard it. therefore, in contrast to the original kd where the student learns from the best pre-trained teacher, in pro-kd the student starts its learning process together with the teacher. furthermore, in contrast to the takd technique which reduces the capacity of the teacher by adding intermediate tas to the distillation process, in our pro-kd, we mitigate the capacity gap by making the training path of the student more smooth and gradual by following the training footsteps of the teacher. moreover, to make the training smooth further for the student, inspired by @xcite , we apply an adaptive temperature factor to the output of the teacher while being trained. this temperature factor is decreased during the training. we will show the effectiveness of our solution using theoretical justification and empirical evaluations. we evaluate pro-kd by performing experiments on both nlp (the glue @xcite ) benchmark and squad 1.1 and 2.0) and image classification tasks (cifar-10, cifar-100). the contributions of this paper are summarized in the following: in pro-kd, the student follows the training footsteps of the teacher. moreover, it introduces a dynamic temperature function to the output of the teacher when it is distilled to the student to make the student training gradual and smooth. ## related work in this section, we review the most related works to our paper in the literature. ## methodology: pro-kd methodology of this paper concerns addressing the capacity-gap and checkpoint-search problems in kd. we propose our pro-kd technique which is a progressive training procedure in which the student is trained together with the teacher. therefore in pro-kd, the student is not exposed to the output of a fully trained teacher from beginning, and instead, it learns from the teacher at the same time while the the teacher is being trained. this training process for the student is more gradual and also the student can learn from the training path of the teacher. inspired by @xcite , we define the training in two phases: first, the student is only supervised by an adaptive smoothed version of the teacher; second, the student will be only trained on ground-truth labels. the detail of each phase is explained in the following. ## experiments and results in this section, we evaluate our pro-kd on 3 different sets of experiments: on image classification, natural language understanding and question answering tasks. in all these three experiments, we compare pro-kd with the state-of-theart kd techniques which can address the capacity gap problem such as takd @xcite , and annealing-kd @xcite , or rco @xcite technique which does not require searching teacher checkpoints for kd. we also include baselines such as the original kd method @xcite and also the regular training without kd. ## conclusion in this paper, we highlighted the importance of the capacity-gap problem in kd. we used the vc-dimension analysis to show that dealing with larger teachers may discount the benefit of using knowledge distillation in training. to address this capacity-gap problem, we introduced our pro-kd technique. our technique was based on defining a smoother training journey for the student by following the training footprints of the teacher instead of solely relying on distilling from a mature fullytrained teacher. in other words, in contrast to the regular knowledge distillation recipe for model compression in which the student model learns only from a fix fully-trained teacher, in our pro-kd method, the student learns even from the training path of the teacher. we believe that following the training footsteps of the teacher can be quite influential in improving the student training and can lead to mitigating the capacity-gap problem in kd as well. we showed our theoretical analysis and justifications to support this idea; moreover, we evaluated our technique using a comprehensive set of experiments on different tasks such as image classification (cifar-10 and cifar-100), nlp language understanding tasks of the glue benchmark, and question answering (squad 1.1 and 2.0) using bert-based models and consistently got superior results. arguments provided by @xcite state that a neural network learns more useful information faster than non-principal pieces of information about the input samples; however, it does not explain why continuing training teacher will decrease the level of that "dark knowledge". to justify this part, we use results from tishby, naftali, and noga zaslavsky. "deep learning and the information bottleneck principle." @xcite and saxe, andrew m., et al. "on the information bottleneck theory of deep learning." @xcite . in @xcite and @xcite , authors show that there are two stages during training a deep neural network, initial fitting phase and compression phase. these two stages have some important characteristics: based on the aforementioned facts, a teacher which has stopped learning at the optimum epoch provides more informative pieces of information to the student. our method exploits the dark knowledge of the best teacher by training teacher and students together; in our method, during training teacher and student together, at some point teacher will have the highest level of dark knowledge which leads to the best student. the important point here is that we do not need to find the optimum epoch at which this highest level of dark knowledge will be achieved. since we are performing knowledge distillation in fixed intervals and storing intermediate checkpoints, the student trained by guidance of the best teacher will be one of these checkpoints and we have access to that. in contrast with usual kd methods, pre-trained teacher networks are not used in our methods; in fact, we train teacher alongside with the student. based on the aforementioned arguments, pro-kd outperforms other kd methods as it gives us the opportunity of exploiting the experience and the dark knowledge provided by the best teacher, the teacher which has stopped learning at the optimum epoch.