The representation gap between teacher and student is an emerging topic in knowledge distillation (KD). To reduce the gap and improve the performance, current methods often resort to complicated training schemes, loss functions, and feature alignments, which are task-specific and feature-specific. In this paper, we state that the essence of these methods is to discard the noisy information and distill the valuable information in the feature, and propose a novel KD method dubbed DiffKD, to explicitly denoise and match features using diffusion models. Our approach is based on the observation that student features typically contain more noises than teacher features due to the smaller capacity of student model. To address this, we propose to denoise student features using a diffusion model trained by teacher features. This allows us to perform better distillation between the refined clean feature and teacher feature. Additionally, we introduce a light-weight diffusion model with a linear autoencoder to reduce the computation cost and an adaptive noise matching module to improve the denoising performance. Extensive experiments demonstrate that DiffKD is effective across various types of features and achieves state-of-the-art performance consistently on image classification, object detection, and semantic segmentation tasks. Code is available at https://github.com/hunto/DiffKD.
Paper
Similar papers
Peer review
Summary
This authors propose to explicitly eliminate the noises in student feature with a diffusion model to reduce the dicrepancy between student and teacher model for better knowledge distillation. Specifically, they build a lightweight diffusion model to reduce computation cost and introduce an adaptive noise matching module to align the student feature with the approximate noisy level of intermediate diffusion.
Strengths
1.The paper adopts the diffusion model to model the teacher feature and denoise the student feature for reducing the dicrepancy between student and teacher model. 2.Quantitative experiments on multiple high-level tasks are provided to evaluate the effectiveness of the proposed method. 3.The paper is well organized and presented.
Weaknesses
1.For methods, (1)The technique contribution of this work is not significant. (2)It is intuitive to regard the student as a noisy version of the teacher, which lacks the theoretical analysis. And the NFEs of diffusion model also depend on the capacity of student models for different noise levels. (3)The design principle of adaptive noise matching is not explained clearly. For different student models, the final learned γ should be analyzed. (4)It is still confused about how to determine the initial timesteps for the student model to reverse diffusion. The reverse diffusion is started from pure Gaussian noise, how to perform the inference for the feature output from the noise adapter. Would the γ near 0？ 2.For experiments, (1)For Table 2 and Table 3, it seems that the proposed method achieves marginal performance improvement for students trained with strong strategies. The reason should be further discussed and analyzed. Besides, it only obtains comparable performance on object detection task, which did not show obvious advantages than other methods. (2)More visualized results should be provided, such as the affinity matrix. (3)Whether other generative models, e.g., GANs, FLOW, can be adopted to model the dicrepancy between student and teacher model or not? It should be discussed and compared to demonstrate the effectiveness of diffusion model. (4)In Table 6, Params and FLOPs are not shown completely. (5)Since different model capacity determines the dicrepancy between student and teacher model, how about the difference of NFEs for different student models? (6)Some important details are missing, such as the total timesteps for training the diffusion model, the initial timesteps for the student model to reverse diffusion. 3.Some writing typos, e.g., Line 13: adpative -> adaptive Line 136: We -> we
Questions
See Section Weakness
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
2 fair
Presentation
3 good
Contribution
2 fair
Limitations
The paper is incrementally novel .
Additional Experiments with GANs
Dear Reviewer AKTb, We extend our sincere gratitude for your dedicated effort and invaluable feedback aimed at enhancing the quality of our paper. We are truly appreciative of your contributions. We are pleased to provide you with an update on our progress, specifically concerning the incorporation of GANs in the context of KD. As the author-reviewer discussion period approaches its final days, we wish to encourage open and constructive dialogue. If you have any concerns, insights, or suggestions, we are eagerly available for discussion. Your input is highly valuable to us, and we are committed to further refining and advancing our method based on your expertise. --- **Experiments on GANs.** We conducted experiments that integrated mainstream GAN methods in lieu of our DM to transform student features. Specifically, we employed a generator designed with identical architecture to our DM. The generated refined student features were subsequently optimized using a discriminator and an adversarial loss within the GAN framework. The discriminator is trained with refined student feature (fake data) and teacher feature (real data). For a comprehensive evaluation, we compared multiple adversarial losses commonly used in GANs, including DCGAN [1], LSGAN [2], and Hinge loss [3]. We then compared the outcomes of these GAN variants with the results of our DiffKD approach in the table below. |Method|Teacher|Student|ImageNet ACC (%)| |--|:--:|:--:|:--:| |MSE baseline|ResNet-50|MobileNetV1|72.39| |DCGAN|ResNet-50|MobileNetV1|72.89| |LSGAN|ResNet-50|MobileNetV1|72.22| |Hinge Loss GAN|ResNet-50|MobileNetV1|71.24| |DiffKD|ResNet-50|MobileNetV1|**73.62**| Our findings demonstrate that, the GAN-based methods we investigated consistently yielded inferior performance compared to DiffKD. Notably, the performance of these GAN variants exhibited sensitivity to the choice of adversarial loss. This susceptibility can be attributed to the well-recognized challenges of non-convergence and training instability associated with GANs. --- **References** [1] Radford, A., Metz, L., & Chintala, S. (2015). Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434. [2] Mao, X., Li, Q., Xie, H., Lau, R. Y., Wang, Z., & Paul Smolley, S. (2017). Least squares generative adversarial networks. In Proceedings of the IEEE international conference on computer vision (pp. 2794-2802). [3] Lim, J. H., & Ye, J. C. (2017). Geometric gan. arXiv preprint arXiv:1705.02894.
I appreciate the authors carefully answering my questions which covers my concerns. The motivation and contributions had better be polished further. I would upgrade the score.
Summary
This paper proposes a novel method of knowledge distillation. It uses a diffusion model to denoise the student model features, reducing the gap between the teacher and student model. An auto-encoder is also designed to reduce the computational efforts and an adaptive noise module to improve the denoising effect. The method is extensively evaluated and achieves SOTA results in image classification, object detection, and semantic segmentation.
Strengths
This paper is innovative in applying a diffusion model to knowledge distillation to reduce the gap between the teacher and the student model. The method proposed in this paper is highly applicable and can be used for different types of features. It works well for a variety of tasks such as image classification, target detection, and semantic segmentation.
Weaknesses
Some of the hyperparameters are not described clearly. Line 177 says λ1=λ2=1, does this indicate that λ3 is set differently? line 220 says λ1=λ3=1. Distillation methods are usually very sensitive to different backbones/downstream tasks? More evaluations and clarifications of this can be added.
Questions
1. Is the training of the autoencoder and diffusion model done at the same time as the distillation training (according to Equation 9)? Does this cause inadequate training of teacher features and denoising effects at the beginning of training to interfere with student model learning? 2. Why are the methods of comparison in Tables 2 and 3 different?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
3 good
Contribution
3 good
Limitations
n/a
Summary
The paper introduces a new knowledge distillation (KD) method named DiffKD, which aims to bridge the representations between teacher and student features via a diffusion model. The motivation is based on the finding that the student feature is noisier than the teacher feature, and therefore diffusion models can be leveraged to denoise the student feature. Additionally, an efficient diffusion model and a noise matching module are proposed for better efficiency and accuracy. The authors conduct KD on image classification, object detection, and semantic segmentation tasks to validate the method.
Strengths
- The idea of reducing the gap between teacher and student features is an emerging and important topic in KD. This paper discusses the discrepancy between teacher and student, then proposes using diffusion models to reduce the discrepancy, which is interesting and straightforward. - It is good to see that the method can generalize to various feature types and tasks. Unlike existing methods that often focus on specific tasks and design complex loss functions, this method can be used in various tasks and achieve advanced performance with simple loss functions. - The technical contribution is evident. Instead of directly adopting classical diffusion models, this paper introduces a lightweight architecture with an autoencoder to speed up the model and a noise-matching module to improve performance. - The improvements are significant in image classification, detection, and segmentation tasks.
Weaknesses
- The method of denoising the student feature with a teacher-trained diffusion model is a bit strange. According to my understanding, the diffusion model is a generative model. Why don't the authors directly train a diffusion model to take the student model as input and generate a feature that is similar to the teacher for distillation? - In Table 3, compared to DIST, the improvement of DiffKD on Swin-T is not as strong as that on smaller models.
Questions
- In Table 8, why is the performance of DiffKD without AE worse than DiffKD with AE of 512, 1024, 2048 channels?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Applicable.
Summary
This paper presents a novel knowledge distillation (KD) approach. The difference from the existing methods lies in the computation of discrepancy between the teacher and student signals. This paper formulates it using a diffusion model and uses a denoising procedure to reconstruct the teacher features from the student features. Experiments are performed on image classification, object detection, and semantic segmentation datasets.
Strengths
+ The idea of measuring the discrepancy using diffusion models is novel. + The paper is well-written.
Weaknesses
- Although using diffusion model to model the discrepancy is a reasonable idea, this paper lacks sufficient analysis on the essential benefit of using diffusion models for this purpose. Does this paper mean that the diffusion procedure finds the shortest path in a distorted feature space (rather than the plain Euclidean space) which is better in measuring the discrepancy? If yes, are there any validations (metrics, visualizations, etc.) for the statement? - The proposed method mixes a large number of loss terms (Eqn 9). In the ablation part, the contribution of each term is not thoroughly ablated. - The results for image classification, object detection and semantic segmentation are mostly conducted on weak student models. It is questionable whether the method works well on strong student models (because improvement becomes more difficult). - The improvement beyond some competitive results is not strong enough (e.g. 0.1-0.3% gain on object detection, given that the student model is relatively weak).
Questions
Please address the concerns above.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Overall, I think the idea of this paper is interesting. However, there are insufficient validations on either the performance (experiments are not strong enough) or the principle. I am looking forward to further results and/or explanations.
Summary
The authors propose DiffKD, a knowledge distillation technique based on the hypothesis that the student's feature is a noisy version of the teacher's feature. Based on this assumption, they use a diffusion model to iteratively denoise the student's features before matching with the teacher. In addition, they propose to use a linear autoencoder to make diffusion more efficient to compute, and they propose a noise adapter module to predict the student's noise level.
Strengths
1. Treating the student's feature as a noisy version of the teacher's feature is an interesting idea. 2. The authors conducted experiments on three tasks and showed improvements for DiffKD.
Weaknesses
1. The authors did not show strong evidence for the claim that the student's features are the noisy version of the teacher's features. While Figure 2 provides a visualization, it would be best if this can be analyzed quantitatively, as it's the central hypothesis that DiffKD is based on. 2. Figure 4 shows that denoising may not be very important, since the performance is not much better even if more denoising steps are taken. This suggests that it's possible that the noise in the student is very weak, which is further supported by the experiments showing very little improvement over existing methods (less than 1 point), despite the heavy machinery of the diffusion model.
Questions
1. Could you please share some statistics about the predicted noise level from the noise adapter?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
2 fair
Presentation
3 good
Contribution
2 fair
Limitations
Limitations are discussed in the Appendix.
Summary
This work proposes the use of the diffusion model for denoising the noisy feature of student. It tackles the issues of employing the diffusion model for knowledge distillation which are heavy computation and inexact noisy level of student feature. To tackle the issues, it proposes a light-weight diffusion model consisting of two bottleneck blocks in ResNet and adopt a linear auto encoder to compress the teacher feature. It also proposes an adaptive noise matching module which measures the noisy level of each student feature adaptively and specifies a corresponding Gaussian noise to the feature to match the correct noisy level in initialization.
Strengths
The proposed idea is simple. Showed what is the problem of directly applying the diffusion model to knowledge distillation. (Expensive Computation Cost, Inexact noisy level of student feature) It has conducted various experiments on various vision tasks to show that the proposed method can be widely applicable. (classification, object detection, semantic segmentation). It shows good performance gain on several benchmarks.
Weaknesses
Applying diffusion model to denoise the student feature seems somewhat obvious approach, in other words, not novel. The fact that the feature of student is noisier than that of the teacher is already well studied by other papers as mentioned in line 109-117. In my opinion, the contribution of this paper comes from showing what are the practical problems of applying diffusion process to KD. It showed two problems (Expensive Computation Cost, Inexact noisy level of student feature). However, I am not sure if this paper has done enough experiments to show that it solved the above problems. There are not enough analysis to prove the proposed method is efficient. (Table 8 and Figure 4 seems to be the only ones). How is the diffusion process applied on the logit-level? The effect of ANM is not really convincing, there is only a marginal performance gain. It would be more convincing how the ANM is helping to denoise the student feature better by analyzing the student feature with and without the noise adapter. Visualization of student feature with/without ANM could help to better understand the effect of ANM.
Questions
Please refer to the Weaknesses.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Please refer to the Weaknesses.
Thank you for your kind responses. Some of my concerns are now a bit resolved. Another question is raised while reading your response and the paper again. I understood that minimizing eq.(7) would not only optimize the parameters of student but also gamma of ANM, but does it also affect the diffusion model during training? Is DM only trained by eq.(2) or is it also affected by eq.(7) ? Thank you.
Thank you for your kind reply. Below, we have refined my responses to your queries: > 6. I understood that minimizing eq.(7) would not only optimize the parameters of student but also gamma of ANM, but does it also affect the diffusion model during training? Is DM only trained by eq.(2) or is it also affected by eq.(7) ? In our early experiments, we have explored both options: * (a) updating DM with teacher feature (eq. (4)) and distillation loss (eq. (7)); * (b) updating DM with only the teacher feature (eq. (4)), while discarding the gradients produced by the distillation loss (eq. (7)). Interestingly, we observed that both options (a) and (b) yielded almost identical performance. However, for the sake of simplicity and efficiency in implementation, we have chosen option (a) as the approach to be used in our final code. |Option|Teacher|Student|ACC (%)| |:--:|:--:|:--:|:--:| |(a)|ResNet-34|ResNet-18|72.22| |(b)|ResNet-34|ResNet-18|72.18| |(a)|ResNet-50|MobileNetV1|73.62| |(b)|ResNet-50|MobileNetV1|73.68| Thanks, Authors
Post-rebuttal comments
I read the authors' rebuttal and other reviewers' comments. I think the rebuttal addressed part of my questions. The technical contribution of this paper is not so significant but sufficient to get published at NeurIPS. Two more concerns. - Regarding the answer to my Q2, I am a bit curious why removing $L_\mathrm{diff}$ can downgrade the results below the MSE baseline. - For Q3, I meant that a stronger student model (e.g. Swin-B) shall be tested and reported. Swin-T is not acceptable. I choose to keep my original rating for now. If new results on Swin-B (or comparable models) are not reported, I will consider downgrading my score.
Response to post-rebuttal comments
Dear Reviewer pbtu, Thank you for your follow-up discussions. Below, I have refined my responses to your queries: > 5. Regarding the answer to my Q2, I am a bit curious why removing $L_\mathrm{diff}$ can downgrade the results below the MSE baseline. The $L_\mathrm{diff}$ term plays a crucial role in optimizing the diffusion model within our DiffKD framework. By removing $L_\mathrm{diff}$, the parameters of the diffusion model remain in a randomly initialized state. Consequently, the model cannot accurately predict the noise in the student features and may even introduce new noise to the features. This undesired denoising of features significantly degrades the performance of DiffKD, resulting in results below the MSE baseline. > 6. For Q3, I meant that a stronger student model (e.g. Swin-B) shall be tested and reported. Swin-T is not acceptable. We have initiated experiments to train Swin-B with Swin-L. Due to the larger size of these models, the experiments require a longer time to complete. We assure you that we will include the results of these experiments in a new comment as soon as they are available.
Experiments on Swin-B Student
Dear Reviewer pbtu, We trained the Swin-B student with Swin-L teacher using KD, DIST, and our DiffKD, respectively. As the results shown in the following table, with the stronger student Swin-B, our DiffKD can also achieve a 1.2% improvement compared to the independent training. Furthermore, when compared to other knowledge distillation (KD) methods, DiffKD shows promising results. It improves upon vanilla KD by 1% and enhances DIST by 0.4%. |Student|Teacher|Method|Top-1| |:--:|:--:|:--:|:--:| |Swin-B|-|-|83.5| |-|Swin-L|-|86.3| |Swin-B|Swin-L|KD|83.7| |Swin-B|Swin-L|DIST|84.3| |Swin-B|Swin-L|DiffKD|**84.7**|
Dear Reviewers, We wanted to kindly remind you that the author-reviewer discussion period is nearing its halfway mark. We would like to take this opportunity to ensure that our responses have adequately addressed your concerns and to inquire if there are any further questions or clarifications you may require. Your valuable input and expertise are essential in enhancing the quality and impact of our research. We greatly appreciate the time and effort you have already dedicated to reviewing our paper, and we eagerly await your feedback. Thank you for your attention and consideration. Best regards, Authors
Reply to the authors
Firstly, thank you for the detailed reply. I really appreciate it! 1. I think you are right that for any two features, we can say $\boldsymbol{F}^{(s)} - \boldsymbol{F}^{(t)} = \delta$. However, I feel this is different from whether $\delta$ is truly a random Gaussian noise. I think this is also not something that MSE and PSNR can show. 2. The result shows one step diffusion already achieves good performance. This makes me question whether diffusion is truly necessary. On the other hand, a simple linear transformation may be too simple. Do you know how this performs if you apply GAN, for example? 3. I could not find the updated figure about the distribution of the noise levels. I'm sorry if I was looking at the wrong place. Again, thanks for the response!
Thank you for your kind reply and follow-up discussion. Our responses to your queries is as follows. > 4. I feel this is different from whether $\delta$ is truly a random Gaussian noise. We want to clarify that the noise $\delta$ in our diffusion models is not limited to being a Gaussian noise. In fact, it is a mixture of Gaussian noises, which has been shown to have the capability to approximate almost all distributions [1, 2]. This choice allows us to effectively model and denoise complex data distributions. To provide further clarity, in the forward diffusion process, "_the diffusion process is fixed to a Markov chain that gradually adds Gaussian noise to the data according to a variance schedule $\beta_1,...,\beta_T$_" (DDPM [3]). In the diffusion process, it recursively predicts Gaussian condictions for multiple timesteps (it is the same as the timesteps of diffusion process in original DM). Therefore, by utilizing DMs, we can effectively eliminate any distribution of $\delta$ and generate a clean, denoised student feature that aligns with the teacher feature. > 5. The result shows one step diffusion already achieves good performance. This makes me question whether diffusion is truly necessary. On the other hand, a simple linear transformation may be too simple. Do you know how this performs if you apply GAN, for example? We apologize for any confusion caused by our previous response. To clarify, the `transformation` mentioned in our table is not a simple linear transformation. Instead, we adopt the same architecture as the diffusion model to transform the student feature. By making this comparison, we aim to demonstrate that the training manner of diffusion models provides a benefit to the distillation performance, even without the multistep denoising process. Furthermore, we have included results for different NFEs on other models in Table 2 of the rebuttal PDF. These results indicate that with ResNet-18 and EfficientNet-B0 students, the improvement achieved with 3 NFEs compared to 1 NFE is significant. Regarding your suggestion of using GANs for transformation, we have initiated an experiment to transform the student feature using a GAN. We are currently in the process of training and evaluating the results. Once the training is completed, we will report the outcomes and include them in our paper. > 6. I could not find the updated figure about the distribution of the noise levels. I'm sorry if I was looking at the wrong place. In accordance with the guidelines for this year's NeurIPS conference, authors are not allowed to update the submitted manuscript or appendix directly. Instead, we have provided a rebuttal PDF that includes additional information and clarifications. You can find the rebuttal PDF in our common response to all the reviewers, titled "Author Rebuttal by Authors," marked with an orange label. In the rebuttal PDF, Figure 1 presents the statistics of the noisy levels ($\gamma$) that you were looking for. --- **References** [1] Bishop, C. M. (2006). Pattern Recognition and Machine Learning. Springer. Chapter 9 [2] McLachlan, G. J., & Krishnan, T. (2007). The EM Algorithm and Extensions. Wiley. Chapter 8 [3] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33, 6840-6851.
Additional Experiments of Applying GANs
Dear Reviewer tFex, We conducted experiments that integrated mainstream GAN methods in lieu of our DM to transform student features. Specifically, we employed a generator designed with identical architecture to our DM. The generated refined student features were subsequently optimized using a discriminator and an adversarial loss within the GAN framework. The discriminator is trained with refined student feature (fake data) and teacher feature (real data). For a comprehensive evaluation, we compared multiple adversarial losses commonly used in GANs, including DCGAN [1], LSGAN [2], and Hinge loss [3]. We then compared the outcomes of these GAN variants with the results of our DiffKD approach in the table below. |Method|Teacher|Student|ImageNet ACC (%)| |--|:--:|:--:|:--:| |MSE baseline|ResNet-50|MobileNetV1|72.39| |DCGAN|ResNet-50|MobileNetV1|72.89| |LSGAN|ResNet-50|MobileNetV1|72.22| |Hinge Loss GAN|ResNet-50|MobileNetV1|71.24| |DiffKD|ResNet-50|MobileNetV1|**73.62**| Our findings demonstrate that, the GAN-based methods we investigated consistently yielded inferior performance compared to DiffKD. Notably, the performance of these GAN variants exhibited sensitivity to the choice of adversarial loss. This susceptibility can be attributed to the well-recognized challenges of non-convergence and training instability associated with GANs. --- **References** [1] Radford, A., Metz, L., & Chintala, S. (2015). Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434. [2] Mao, X., Li, Q., Xie, H., Lau, R. Y., Wang, Z., & Paul Smolley, S. (2017). Least squares generative adversarial networks. In Proceedings of the IEEE international conference on computer vision (pp. 2794-2802). [3] Lim, J. H., & Ye, J. C. (2017). Geometric gan. arXiv preprint arXiv:1705.02894.
Thank you for the detailed updates! I feel the fact that DMs "have the capability to approximate almost all distributions" weakens the central claim of the paper that student features are noisy versions of the teachers. This feels like DM is only used as a bridge between the two features, and diffusion models are not particularly reasonable in the sense of denoising. However, I feel the results still suggest that DM can be a powerful model to align student and teacher features. Therefore, I would like to maintain my positive score for the submission.
Response to the authors' rebuttal
Thanks for the authors' responses. All my concerns have been solved. I tend to keep my initial rate and vote for accepting this paper.
Decision
Accept (poster)