Parameter-Efficient and Student-Friendly Knowledge Distillation

Introduction

With the continuous growth of data scale and computing resources, the scale of deep neural networks has rapidly increased, from millions of parameters to billions of parameters . Although these huge models, such as BERT , GPT-3 , and CLIP , have powerful multitasking capabilities, they cannot be deployed on some edge devices with limited computing resources, which further limits the environmentally friendly and efficient deployment of deep learning applications.

Knowledge distillation (KD) , as an important method for model compression, has been widely used in various fields [51; 20; 41; 49] of deep learning. This paradigm utilizes a pre-trained teacher network to obtain a student network that is close to the teacher network but with fewer parameters. In the traditional distillation framework [41; 31; 30], the prediction output (soft label) is produced by the fixed teacher, as shown in Figure 1(a) . However, limited by the capacity of the student model, the teacher’s knowledge cannot be well transferred to such small models, which makes traditional distillation methods suboptimal [20; 27]. Based on this phenomenon, we propose the first research question: RQ1: How to transfer teachers’ knowledge to students more effectively?

In recent works [29; 20; 45; 35], the teacher network and the student network are jointly trained to make the teacher’s knowledge more friendly to the students. The online distillation methods, such as DML and KDCL , usually update most or even all the parameters of the teacher by using the real labels and the feedback information (soft labels) of the students, as shown in Figure 1(b). However, in the setting of online distillation, a large teacher network and a student network need to be trained simultaneously for each new downstream task from scratch, which is too time-consuming and inefficient. This motivates us to consider the following research question: RQ2: How to make the friendly transfer of knowledge more parameter-efficiently?

In this paper, we propose a new distillation framework that is parameter-efficient and student-friendly, as shown in Figure 1(c). Our paper forges a connection between two literatures that have evolved independently: knowledge distillation and parameter-efficient learning. This allows us to leverage the powerful methods of parameter-efficient learning to ensure the friendliness of teachers’ knowledge to students, and to update teachers’ knowledge representation more efficiently. Based on extensive experiments and analyses, we show that our framework can utilize the information from ground-truth labels and student supervision to train the adapter modules, and further narrow the gap between the teacher and student models, which makes knowledge transfer easier.

We propose a parameter-efficient and student-friendly distillation (PESF-KD) framework, where an adapter is deliberately designed for the teacher.

We provide theoretical and experimental evidences to show PESF-KD can facilitate the knowledge transfer by reducing the gap between teacher and student.

We empirically validate the effectiveness and efficiency of our PESF-KD upon several vision and language models compared to existing knowledge distillation methods.

Background

Gradient descent is mostly used to optimize the cross entropy of the hard ground-truth labels and predicted values in deep learning. And in recent years, Szegedy et al. introduced Label Smoothing to soften and weight the traditional hard labels with the uniform distribution. This approach has successfully improved the effectiveness of several deep learning models and has been widely validated in natural language processing (NLP) and computer vision (CV). And to date, this approach has also been used as a training trick to improve the training of models. We provide a mathematical description of the label smoothing process. First we show the original cross-entropy:

where $y_{k}$ is "1" for the correct class and "0" for the rest. Then the label smoothing is achieved by increasing the smoothing parameter $\alpha$ to change $y_{k}$ to $y_{k}^{LS}$:

When a network is trained with label smoothing, the differences between the logits of the correct and incorrect classes become a constant that is dependent on $\alpha$, while knowledge distillation provides dynamic soft labels to let the network learn the distribution of teachers.

where $\tau$ is the temperature used in KD, which controls how much to rely on the teacher’s soft predictions. For simplicity, we use $\mathcal{L}_{KL}$ represent $\mathcal{L}_{KL}\left(\bm{p}(\tau|\theta^{s}),\bm{p}(\tau|\theta^{t})\right)$ in the following sections.

2 Parameter-Efficient Training

Transfer learning from pre-trained models is a general learning paradigm that has been applied to a variety of tasks [15; 9; 26]. The most popular transfer learning method nowadays is to fine-tune all of the model’s parameters (full fine-tuning). But full fine-tuning is relatively expensive because parameter adjustment correlates to the need to completely retrain the entire model for each downstream task. Artificially constructed modules with a modest number of parameters are used in approaches like Prompt and Adapters [18; 25] to achieve a fit between the pre-trained model and the downstream task with parameter-efficient training.

Methodology

Motivation. Compared with label smoothing, knowledge distillation can improve the training of the network for the following two reasons: First, teachers can understand the nuances of different classes, and such inter-class information brings more information than label smoothing and helps students generalize some unseen data. Second, the soft distribution of teachers constrains students’ learning directions such that they will not be mistakenly overconfident. For a student network, the information it learns is obtained through this soft label. Previous works [29; 52; 20; 27] pointed out that the mismatch of teacher and student network capacity can cause the knowledge transfer difficulty of such soft label . One of the solutions to reduce transfer difficulty is smoothing the teacher’s output by adjusting the temperature . Müller et al. indicates that the label smoothness of the target provided by the teacher exerts a great influence on the student network, and the difference information between classes determines whether the student’s performance can be improved. But manual conditioning of the label smoothness by the temperature is quite difficult and may cause the loss of inter-class information when the temperature is too high. In our preliminary experiments, we found that if the teacher network continues to fine-tune through the ground truth labels with the rest of the settings as the valina KD , teacher’s labels will be smoother, and the accuracy of the distilled student network is better, where we leave this competitive setting as the strong benchmark, namely “KD w/ fine-tuned teacher”, in the main experiments (Section 5). This phenomenon urges us to explore the relationship between the teacher network and the student network in Section 4. Training Objectives. As shown in Figure 1(b), different from vanilla KD , our distillation method requires fine-tuning the parameters of the teacher network. The teacher network needs to output soft labels to supervise the student network, it also requires ground-truth labels to training itself. Take classification task as instance, the corresponding teacher’s loss is:

where $c$ is a class label and $C$ denotes the set of class labels.

In vanilla KD , students receive soft label supervision (in Equation 4) from the teacher as well as label supervision as the following formulation:

where the task loss $\mathcal{L}_{task}$ follows the same format as the teacher network.

2 Parameter-Efficient and Student-Friendly Knowledge Distillation

A Closer Look at Teacher-Student Relationship in Distillation

The relationship between the teacher network and the student network in KD is rarely discussed. When independently training the model from scratch, the larger model is more likely to output sharper values and obtain better accuracy, while the smaller model is more likely to output smoother values and obtain poorer accuracy [5; 52; 27; 6]. This phenomenon (i.e., capacity mismatch) significantly exacerbates the difficulty of knowledge transfer in knowledge distillation. An optional solution to measure the difference in the sharpness of the network output between teachers and students is to directly use the entropy of the network output to get the output value of the maximum confidence score. Differently, we use a simple and intuitive sharpness metric to get a smooth approximation to the maximum function considering the overall information of each class. If we use $K$ to denote $K$ classes, the Sharpness is defined as the logarithm of the exponential sum of logits:

For the difference between teacher and student networks, we use the sharp gap between the two networks as a metric:

2 How to Narrow the Gap?

As shown in Equation 8, we get the expression for the gap between teachers and students. In this section, we will explore what factors affect this gap. We first approximate this expression using a Taylor second expansion:

Following Hindon’s assumption and also through experimental phenomena , it can be known that the logits of each training sample are approximately zero-meaned so that $\sum_{j}^{K}z_{j}(\theta^{s})=\sum_{j}^{K}z_{j}(\theta^{t})=0$. So the gap can be rewritten as:

where the $\sigma^{2}=\frac{1}{K}\sum_{j}^{K}z_{j}(\theta)^{2}$ is the variance of logits.

Naturally, assuming that the temperature $\tau$ is fixed. Then the change of the gap comes from the change in the variance of teachers’ logits and students’ logits. Once the these logits become smooth then the corresponding variance becomes smaller, if the logits become sharp then the variance becomes larger. The smoothness of the final logits results in a change in the gap. In the following sections, we compared three distillation methods, namely, vanilla KD (vanilla) with different temperature and our proposed methods – fine-tuned teacher (train with finetune teacher) and adaptive teacher (train with adapter) to check out how temperature and respective methods affect the gap.

In order to more intuitively show the effect of logits output smoothness on the gap, we show the average major logit distribution in the Figure 4. In most cases, the logit has the largest value that represents the model’s category prediction, and other tiny values show that the input image is similar to those other categories. In the real situation, the variance of students’ logits will be affected by the variance of teachers’ logits, and these two statistics change simultaneously. It is difficult to quantitatively analyze the gap between a fine-tuned teacher and an adaptive teacher, but we can get inspiration from Figure 4 and 4. Obviously, vanilla KD brings more sharp logit output, that is, greater variance. The logit output of fine-tuned teacher is in the middle position compared other two methods while PESF-KD has the smallest variance. It is worth noting that our methods greatly reduce the variance of logits compared to the vanilla KD, and both our proposed methods have relatively small variance. In Figure 4, it shows the degree of sharpness for different temperatures in vanilla KD and our methods. Consistent with the analysis of Equation 10, since the teacher’s output is unchanged, that is, the variance is unchanged, by increasing the temperature, the student’s output is indeed smoother, resulting in a larger gap. Figure 4 shows the Top-1 accuracy of the corresponding methods. In general, lower gaps are associated with higher accuracy (lower temperature, lower gap value), but very close gaps may introduce anomalies (see the case of temperature 2 and 4). This phenomenon shows that smoothing the output of teachers and students within a certain range can improve the friendly transferability of knowledge (in terms of improving the accuracy and reducing the gap). In the face of each sample, it is impossible to simply manually adjust the common temperature to achieve a dynamic smooth output to improve the students’ ability to accept the teacher’s knowledge. With our methods, the smoothness of the output can be dynamically adjusted according to the gradient of each sample to achieve better knowledge transfer (compared to the vanilla KD with the best temperature setting, our methods have improved a lot about accuracy, and a large reduction in the gap).

In this part, we show that trainable part in teacher models can narrow the gap between student and teacher and make knowledge transfer process more student-friendly thus achieving better accuracy. We further explore the relationship between network consistency (sharpness gap in Figure 4, KL-Divergence in Figure 8 & 8 and CKA in Figure 8 & 8 ) and accuracy in Figure 4. The three metrics mentioned above measure the final degree of consistency of the teacher-student network from different perspectives. A lower sharpness gap represents a closer knowledge representation of the teacher-student, and a lower KL represents the final convergence degree of the lower bound through distillation learning, while the a larger CKA represents larger similarity of the students and teachers. We get the following interesting findings: 1) From Equation 10, it is clear that the gap also decreases when the student network trained with the vanilla KD. This reduction comes from the fact that the output of the student network becomes sharper i.e., more similar to the output of the larger teacher network (both kl-divergence and gap decrease). 2) Both our proposed methods can reduce the gap, and model trained with fine-tuned teacher can bring a great reduction (from 36.3 to 16.8). Our methods also make the output of teachers and students more consistent (the KL-Divergence, gap of the two methods are significantly lower and the CKA is higher than those of vanilla KD with different temperature). 3) The accuracy of the final student network trained by our two methods is similar, although the gap between them is quite different (32.4 vs. 16.8). The KL-Divergence (0.42 vs. 0.23), CKA of logits and predictions (almost same) of these two proposed methods are closer, which shows that our methods can guarantee the consistency of teacher and student characteristics. It also illustrates that gap and KL-Divergence interpret the similarity of output distributions from different perspectives.

To explore the changes in class distribution of logits, we visualize the penultimate layer representation of student model ResNet20 on dataset CIFAR-100 as shown in Figure 9. We randomly choose 7 classes and record the penultimate layer representation in the 230-th epoch and then use T-SNE to project the data in the 2D plane. Classes are marked with different colors. We observe that clusters in our proposed approach are tighter because the student model is encouraged to learn more information from all other class templates in the training data set by narrowing the sharpness gap between teacher and student networks. Besides, when looking at the projections, some clusters, i.e., crimson and dark blue ones, are more discernible in our proposed models than in Vanilla KD.

Experiments

Datasets Two types of tasks incluing image classification (CIFAR-100 and ImageNet ) and natural language understanding (GLUE ) are adopted for a series of experiments. Specifically, we test on two types of NLU datasets according to the amount of training resource. Low resource: MRPC and STS-B for Paraphrase Similarity Matching; SST-2 for Sentiment Classification ; RTE for the Natural Language Inference. High Resource :QNLI for the Natural Language Inference and QQP https://www.quora.com/q/quoradata/First-Quora-Dataset-Release-Question-Pairs for Paraphrase Similarity Matching. Baselines We report several knowledge distillation methods for comparison, including vanilla KD , knowledge distillation via collaborative learning (KDCL) , deep mutual learning (DML) , contrastive representation distillation (CRD) , relational knowledge distillation (RKD) and probabilistic knowledge transfer(PKT) . According to Gou et al. , KD methods can be divided into two groups, online distillation and offline distillation. For a more fine-grained comparison, We further split them into three different kinds, online KD (DML, KDCL), offline KD (vanilla KD, PKT) and representation KD (CRD, RKD). Besides these methods, we report the result of our method “w/ Fine-tuned teacher” and “PESF-KD” to support our argument about less sharpness gap helps student to perform better to absorb the knowledge of the teacher.

2 Experimental Setup

For CV tasks, we follow previous works [41; 51] using various combination of student & teacher networks. Each pair of student & teacher networks are from different capacity and architecture, including ResNet and VGG . We run two types of distillations, isomorphic distillation and isomeric distillation. For isomorphic distillation, we run three different combinations (ResNet56-ResNet20, ResNet 110-ResNet32 and VGG13-VGG8). For isomerism distillation, the results of ResNet-56 to VGG-8 are reported. For NLU tasks, we first fine-tune the pre-trained teacher (12-layers version of BERT-Base) and then train student model (6-layers version of BERT-Base) on each downstream task. We report Top 1 Accuracy (@1) for image classification experiments as a network performance metric. For NLU tasks, we report the same format as on the GLUE leaderboard. For MRPC and QQP we report F1. For STS-B, we report Pearson and Spearman correlation. For other tasks we report accuracy. Besides, to measure network consistency between teachers and students , we report average predicted KL divergence (KL), CKA consistency (CKA) and the sharpness gap (GAP) as similiarity metrics.

3 Results

This section compares different KD methods, which can be grouped as offline KD, representation KD, and online KD, and their training running time (batch time), need for updated parameters (params), and the corresponding accuracy., as shown in Table 1. Encouragingly, our experiments fit our two previous hypotheses. On the one hand, through the joint training of the teacher network and the student network, the problem of difficult knowledge transfer when the teacher-student capacity does not match can be alleviated. Specifically, the distillation results of the online distillation methods are significantly higher than those of the offline distillation methods, especially in the case of excessive difference in network capacity between teachers and students (vgg13-8) get a huge improvement. And our method “w/ Fine-tuned teacher” achieved the best results among almost all distillation methods, such as 71.65 for resnet56-20 and 73.90 for resnet110-32 and 74.40 for resnet56-vgg8.

On the other hand, through the adapter module, the amount of parameters that teachers need to update can be reduced on the premise of alleviating the problem of capacity mismatch, thereby reducing the cost of training. As shown in Table 1, the method of online distillation will greatly increase the training cost due to the need to update the parameters of the teacher network and the student network synchronously (see the batch time change). While PESF-KD significantly reduces the number of parameters that need to be updated for training, i.e., vgg13-8 reduces the number of parameters from the $\sim$ 54M to $\sim$ 16M. Obviously, it also significantly reduces the time required for training. Whereas in most combinations of teacher-student architectures on distillation, even if using adaptive teacher has a slight reduction in the best distillation results, such reduction is negligible. To further illustrate the superiority of our methods, we further compare the current typical distillation methods like previous work , as shown in Table 2. Results of baselines in Table 2 are reported by . The results show that our two proposed methods achieve the best results among these methods under a broad combination of teacher-student architectures in distillation.

3.2 Results on ImageNet

We further conduct experiments for ResNet (ResNet50, 102.2M to ResNet18, 46.8M) distillation on a lager dataset ImageNet. As shown in Table 3, our methods achieves comparable Top-1 accuracy results compared to other baselines, including SOTA representation KD method CRD (70.00 for fine-tuned teacher; 69.96 for PESF-KD). Notably, compared with other online KD method KDCL (70.17), student trained with fine-tuned teacher show similar top-1 accuracy (fall behind 0.17 point) and the same KL loss. That’s partly because both the methods have similar working mechanisms which allow part of the teacher model to be trainable during the student training process. But KDCL fails to fill the sharpness gap between teacher and student (31.8), while our methods are more capable of narrowing the gap (14.9 and 19.2). Besides, the student trained with the adaptive teacher reach similar top-1 accuracy (69.96) compared with fine-tuned teacher (70.00) and only require less than one-third the time (0.47 batch time for adaptive teacher and 1.44 batch time for fine-tuned teacher). That fits our assumption of PESF-KD to improve performance and reduce training time.

3.3 Results on GLUE

In Table 4, we further explore the effect of the PESF-KD in NLP dataset, and other baselines are reported by . Time in the Table 4 refer to training resources cost, which is the lowest consumption with our PESF-KD compared with other baselines except for vanilla KD. In , they find that the overly strict regularization of KD is unnecessary in NLP datasets, which will actually reduce the final performance using KD in most cases of GLUE. However, PESF-KD can even bring further improvements to most datasets with KD, probably due to better student learning from teachers. Compared to vanilla KD , PESF-KD achieve consistent improvement in low-resource situations, showing that the same effect of low-resource adaptability . Similar to the results of CV datasets, our method also obtains similar results and requires minimal training cost compared with baselines on NLP datasets, and even surpasses other baselines on SST-2, which verifies the generalizability of our methods.

4 Analysis

The adapter module is our recipe for success in the above performance comparisons. To explore the influence of adapter structure on classification results, we compare three classic adapter structures and report their top-1 accuracy and CKA consistency. We use the teacher-student combination of vgg13-8 and conduct comparative experiments on CIFAR-100. As can be seen from Table 6, the simple and efficient adapter achieve second performance in top-1 accuracy, best CKA scores and time consuming. Therefore we use simple adapter module in all experiments.

The hyper-parameter $\alpha$ decides the percent of cross entropy loss between the ground truth labels and student models prediction, and the KL loss between student and teacher models. Bigger $\alpha$ means a higher percent of kl loss. Figure 11 shows that bigger $\alpha$ benefits student models in PESF-KD. That’s partly support our hypothesis of PESF-KD to promote a easy-learning teacher.

Figure 11 shows the performance of student models in differ temperature. In the experiments vanilla KD and our proposed approach has similar trend in different temperatures. And PESF-KD shows better performance in all experiments.

We also test PESF-KD in other methods to verify the potential improvement on other knowledge distillation methods. Table 6 shows the results of four knowledge distillation approaches (vanilla KD , PKT , CRD and RKD) combined with PESF-KD on CIFAR-100. Compared with standard knowledge distillation approaches, results with PESF-KD get better performance on top-1 accuracy and indicate that PESF-KD has the potential to become a plug-in method on top of different KD methods.

Conclusion

In this paper, we present PESF-KD, a novel knowledge distillation framework by applying adapters to optimize the teacher network for better knowledge transfer to the student network. Through detailed analysis, we point out that the decline in sharpness and a better ability to distinguish within classes lead to better knowledge transfer, which leads to better results. Extensive experiments demonstrate the robustness and effectiveness of our method. Future works of embedding PESF-KD into existing distillation framework is expected to generalize our methods.

References

Appendix A Appendix

Table 7 shows the statistics used in our experiments. The training set, testing set, and development set refer to the scale of data in each part of the dataset. Table 7 shows that ImageNet is much bigger than CIFAR-100 in scale. And that’s the reason we choose the two datasets, to verify our proposed approach in datasets of different scales. GLUE benchmark can be naturally split into low resource tasks and high resource ones for the similar propose. We divide low-resource tasks and high-resource ones by 100k on a data scale.

We are also concerned about people’s consent and privacy in datasets. From the information search from the website, faces in ImageNet have been blurred in order to protect privacy after being blamed to collect information from people without consent. But we do not find a similar announcement on the website of CIFAR-100 and GLUE. Maybe it’s because they have pre-processed the data or the data itself contains little privacy information. It’s better for the websites to make the announcement of not containing personally identifiable information or offensive content.

We implement our models with Pytorch, and our experiments are as follows:

Tools: Python3.7, tensorflow2,2,0, sklearn 0.23.2

In CV experiment, we follow previous works , the settings on CIFAR-100 and ImageNet dataset are the same as these works. In CIFAR-100 we train the student model by SGD optimizer with a momentum of 0.9, a batch size of 64 and weight decay of $5\times 10^{-4}$. The learning rate starts from 0.05 and decays by 10 every 30 epochs after 150 epochs. And on ImageNet we train the student model by SGD optimizer with a momentum of 0.9, a batch size of 256 and weight decay of $1\times 10^{-4}$. The learning rate starts from 0.1 and decays by 10 every 30 epochs after 30 epochs. In the experiment of students trained with the fine-tuned teachers, we train the teacher model along with the student model during the training process with a learning rate of $1\times 10^{-3}$. And in the experiment of a student trained with a teacher with the adapter module, also called adaptive teacher, the learning rate of the trainable part in the teacher model is set to $1\times 10^{-4}$. Notably, classification loss from the teacher model is appended to the loss of students by multiplying a hyper-parameter of 0.5.

For NLP, we inherit parameters like maximum sequence length, temperature and batch size according to setting from previous works . We also perform grid search over the sets of the student learning rate $\lambda$ from {1e-5, 2e-5, 3e-5, 4e-5, 5e-5}, teacher learning rate $\mu$ from {2e-6, 1e-5, 2e-5}, batch size from {32, 64}, the weight of KD loss from {0.3, 0.5, 0.8, 0.9} for a better performance. We evaluate our methods on dev set of GLUE benchmark.