CIFD: Controlled Information Flow to Enhance Knowledge Distillation

We thank the reviewer for their time and effort. Detailed comments follow the summary.

Summary

• We showed that our proposed method outperforms TAKD and DGKD even when there is only 1 RDM. This ensures fairness. However, in terms of training cost, our method with 3 RDMs which is far superior is also far cheaper (as seen Fig. 1).
• We clarified that the low cost of our method is due to the relatively cheap computation cost of RDMs compared to the student/teacher model, ability to train RDMs in parallel, and train them for lesser epochs. We also extended our numerical training cost analysis in appendix C.3 and provided more information.
• We directly addressed the concern on the efficacy of CIFD over large-student teacher gap. We distilled ResNet18 using ResNet152, ResNet101, ResNet50 (Table 13) as teachers and showed that proposed CIFD shows increasing student performance with teacher size. We also analyzed CLIP models (Table 16, 17) and found that it provides more improvement over baseline when the teacher student capacity gap is large.

Detailed response

Difference between TAKD, DGKD and reported numbers

We wanted to clarify if the reviewer is referring to the difference in performance of TAKD and DGKD in Table 1 and the original numbers in their papers? If so, the problem is both those papers used the test set for parameter fine-tuning [A] and [B]. After rectifying these errors, we got the above numbers.

[A] TAKD first author's comments on GitHub: https://github.com/imirzadeh/Teacher-Assistant-Knowledge-Distillation/issues/19#issuecomment-732454350

[B] DGKD implementation by the first author. Validation function (defined in line 133-137) called in line 143, uses the test set to select the best accuracy across epochs in https://github.com/wonchulSon/DGKD/blob/main/train.py

Fairness between TAKD, DGKD, and CIFD

Coming to the fairness between comparing TAKD and DGKD in Table 3, one fairness is along the axis of training cost. In which case, as seen in Fig. 1, we incur lower training cost despite using three RDMs and significantly outperform them. Further as seen below, even when comparing our 1 RDM results with theirs, our proposed method is better.

| RN34 to RN18 on IN1k | Top-1 | Top-5 |
+----------------------+-------+-------+
|        TAKD          | 71.37 | 90.27 |
|        DGKD          | 71.73 | 90.82 |
|   Proposed (1 RDM)   | 72.05 | 90.70 |
|   Proposed (3 RDMs)  | 72.32 | 90.88 |

Computation costs

We would like to clarify that NormKD, DistKD, IPWD, IFD, and KD are indeed slightly lower cost than our method. However, our algorithm is on the pareto front of the training cost v/s performance curve.

We use the number of MACs in the forward pass as a measure of computation. There are three reasons. Firstly, the number of MACs consumed by RDMs is insignificant compared to a forward pass of the teacher model during RDM training. For e.g., one forward pass of the RDM per image is $1.31 \times 10^6$ MACs (RDM for RN34), whereas RN34 forward pass is $3.7 \times 10^9$ MACs. Thus, during RDM training the teacher model's forward pass is the biggest cost. Secondly, the RDM is trained only for 30 epochs. Finally, the RDMs are trained in parallel amortizing the cost of the forward pass of the teacher model during training. These in total reduce the cost of the training. Below we give the cost computations (also presented in Appendix C.3).

Our method in total costs 762 PMACs (Peta or $10^{15}$ MACs) for distilling RN18 from RN34. First, we compute the RDM training cost. RDMs consist of a three layer fully Connected Network, with a computational complexity of $1.31$ MMACs ($10^{6}$ MACs) per image per forward pass. A 30 epoch training results in total cost of $151$ TMACs ($10^{12}$ MACs) for three RDMs, the significant chunk of computation in training RDMs is coming from running the teacher model in inference mode ($142$ PMACs). The rest of the compute comes from the student model training when the three RDMs, student, and teacher are being used. Note that the RDM forward passes are computationally insignificant at this stage, only $0.38$ PMACs. Finally, we compute the training cost for existing methods like NormKD, DistKD, IPWD, IFD, and KD as $704$ PMACs, using the cost of the student and teacher models.

Larger teacher experiments

We trained distilled RN18 from RN152, RN101, RN50 teachers. DistKD [5] showed that ResNet152, ResNet101 to ResNet18 faces the issue of large capacity gap (Table 3 of DistKD). The parameter ratio of from largest teacher to student is 5.12. Experimental results are in Table 13. First we observed that student performance increased with teacher size. Secondly, our results outperformed that of DistKD. This directly addresses the reviewer's concern about large teacher student capacity gap.

Additionally, we also compare CLIP models. Due to resource constraints, we compared the improvement provided by proposed CIFD over the nearest baseline CLIPKD [31]. Specifically, we compared the difference when a ViT-L-14 teacher was used to train ViT-B-16 and ViT-S-16 students. The parameter ratios are 2.8 and 6.9, respectively. Results are in Table 16 and 17. CIFD shows greater improvement over baseline when the teacher student ratio is large for 3 out of the 5 zero-shot classification datasets and almost always for the two zero-shot retrieval datasets. The capacity gap in CLIP-like models is much larger than the capacity gap in traditional KD settings such as RN34-18 and RN50-MobileNet V1.

Layer for teacher embedding

This is based on the idea that the penultimate layer output is usually considered the embedding of the network. Additionally, the embeddings holds all the important information for classification and we can easily remove that information using our RDM.

Related work position

Thank you for the feedback, absolutely we will move it

Thank you for providing a detailed response that resolved some of my doubts. The range of uncertainty in hyperparameter R is still one of my concerns, but experiments based on CLIP are indeed very distinctive. Based on the opinions of other reviewers and my understanding, I will increase the score to 5.

Thank you for your time and effort, we greatly appreciate your response. We also wanted to acknowledge your inputs in helping us improve the paper, thank you.

We apologize for missing the point on the hyper-parameter R. From our analysis, we found that the values of R are inversely proportional to the capacity gap between the teacher and the student. Since $R^{-1}$ controls the weight of the bottleneck rate, too small an $R$ and the RDM will focus more on compression than retaining features crucial for accuracy, resulting in a poor RDM. Since the RN34 to RN18 gap is small (both in size and performance), smaller $R$ values were more appropriate and hence the selection. We plan to add this insight to our paper in the final version.

If you have any further questions or concerns to improve our paper, please let us know and we are happy to discuss. Thank you again.

We thank the reviewer for their time and effort. Detailed comments follow the summary.

Summary

• By using one RDM without IBM, we showed that the key to our superior performance compared to Factor Transfer (FT) [28] is the principled loss function used to train the RDM. This is in addition to the multiple RDMs and IBM which have no equivalents in FT.

• We showed superior performance against newer works like MLKD, CTKD, and LSKD.

Detailed response

Motivation and comparison with FT [28]

Before we respond let us quickly summarize the work of FT [28]. Instead of using teacher's logits for KD, FT proposed using teacher's embeddings. They transfer the knowledge using a paraphraser network with a convolutional autoencoder architecture. The paraphraser encoder outputs a vector with dimension that is $k$ times the dimension of its input and the decoder used to reconstruct the encoder's input, is discarded after training the encoder. Kim et. al. explored both dimensionality reduction ($k=0.5$) and expansion ($k=4$), and settled on $k=0.5$ as the best performing. During the training of the student model, they use a translater network, similar to the encoder of the paraphraser network and task it to map the student embeddings to the output domain of the paraphraser encoder. Below, we will highlight the main difference between FT and our method, and explain that this difference is crucial in obtaining better performance.

While both paraphraser and our RDM aim to do compression, the mechanism of training the compression module is different. Kim et. al. train the paraphraser network like an autoencoder. The paraphraser training limits information by dimensionality reduction (when $k<1$). We train our network using the Rate-Distortion criterion, the mathematically principled way of performing lossy compression. The crucial difference is that the RDM limits the rate of information flow by minimizing the mutual information between the latent representation ($\hat{Y}$ in Fig. 2(a)) and the input ($X$) in addition to the reconstruction error (see the difference between the loss function in our eqn. (7) and FT's eqn. (1)). This training scheme makes a significant difference in performance as shown in Table 15. In the table, we compare our method with only one RDM with the work of FT, and we see that our method is superior, despite removing all our other innovations, including multiple RDMs, and IBM. Thus, it shows that the principled training of the RDM plays a significant role in improving the performance of the student network. As a further comparison, we also look at IFD [20], which uses an ensemble of three paraphraser networks. However, our method with 3 RDMs and no IBM still outperforms it. This shows that just relying on diversity of initialization of paraphrasers does not yield benefits. However, training in RDMs in a principled manner to focus on different levels of information (fine to coarse) helps train the student better.

To summarize our differences w.r.t. FT

1. We propose a principled way of compressing the features from the teacher model using insights from Rate-Distortion theory. This principled training loss leads to performance gains even in a standalone manner (71.83 > 71.43).

2. We propose multiple RDMs to help student network learn at different compression ratios. This has no equivalent in Kim et. al.

3. We propose the use of IBM in the student network to control the information flow in the student during training. This has no equivalent in Kim et. al.

While at a certain abstraction we can regard RDMs as adapters, the key here is the loss function used to train the adapters. As the new results show, the principled loss function is critical for good performance.

Comparing with MLKD, CTKD, and LSKD

Thank you for pointing them out. Our performance is superior to MLDK, CTKD, and LSKD works as seen in Table 14.

How is $q(\hat{Y})$ computed?

$q(\hat{Y})$ is not produced by a network. Instead it is produced using a non-parametric distribution approximation process proposed in \cite{balle2017end}. The method also computes an upper-bound on the entropy as detailed in \cite{balle2017end}.

W3 and Q2

Thank you for the inputs, we will rectify them.

We thank the reviewer for their time and effort in providing feedback. Detailed response after summary.

Summary

• We directly addressed the concern on the efficacy of CIFD over large-student teacher gap. We distilled RN18 using RN152, RN101, RN50 (Table 13) as teachers and showed that proposed CIFD shows increasing student performance with teacher size. We also analyzed CLIP models (Table 16, 17) and found that it provides more improvement over baseline when the teacher student capacity gap is large.
• We extended our discussion in Appendix C.1 to showcase that the Masked-Image-Modeling (MIM) objective is an approximation of an upper bound on the IB objective, whereas ours is a direct approximation of the IB objective. We compared and showed superior performance compared to Masked Generative Distillation (MGD).
• We compared with DiffKD and showed our performance is superior. Further, it takes two methods, DistKD and DiffKD to achieve similar performance.
• We clarify that we did not modify the student architecture to ensure fairness

Details

Large teacher distillation experiments

We trained a distillation from RN152, RN101, RN50 teachers to RN18 distillation using CIFD. DistKD [5] showed that RN152, RN101 to RN18 face the issue of large capacity gap (Table 3 of DistKD). The parameter ratio from largest teacher to student is 5.12. Experimental results are in Table 13. We observed that student performance increased consistently with teacher size and also outperformed the performance of DistKD. This directly addresses the reviewer's concern about large teacher student capacity gap.

Additionally, we also compare CLIP models. Due to resource constraints, we compared the improvement provided by proposed CIFD over the nearest baseline CLIPKD [31]. Specifically, we compared the difference when a ViT-L-14 teacher was used to train ViT-B-16 and ViT-S-16 students. The parameter ratios are 2.8 and 6.9, respectively. Results are in Table 16 and 17. CIFD shows greater improvement over baseline when the teacher student ratio is large for 3 out of the 5 zero-shot classification datasets and almost always for the two zero-shot retrieval datasets. The capacity gap in CLIP-like models is much larger than the capacity gap in traditional KD settings such as RN34-18 and RN50-MNV1.

Connections of IBP to MIM and DiffKD

We called the RDM modules as such because they capture different levels of information based on the rate constraint. Given existing work in learning based data compression [9, 34, 35], which explicitly work on Rate-Distortion objective (like we do), we felt appropriate to call this the RDM.

We were aware of the similarity between MIM and IBP, hence we included appendix C.1 in our submission to address this. Building upon our insights in appendix C.1, we show that MIM and MGD objectives are upper bounds on the IBP objective. Consider MIM shown in Fig. 7a. The objective from the information bottleneck principle is $\min -I(X;\hat{U}) + \lambda I(T;\hat{U})$. The first term is usually approximated as a reconstruction error between $X$ and $\hat{U}$. Looking at the second term, we can write, $I(T;\hat{U}) \leq H(T)$, where $H$ signifies the entropy. Since $H(T)$ is not dependent on the neural network parameters, it can be dropped. This is the MIM objective, $\min -I(X;\hat{U})$ or minimize the reconstruction error. Similar simplification also holds for MGD. So both MIM and MGD minimize an upper-bound on the IBP got by dropping the second term. In our case, we compute an approximation for the second term, thus making it tighter to the original IBP objective. IBP works better because it forces the student to focus on those features necessary to predict the teacher embedding (first term) and remove information not useful in the prediction (second term), whereas in MIM and MGD the removal of information is not present. There have recently been works studying how IBP and generalization errors are connected ("How does information bottleneck help deep learning?," ICML 23) which provides support to the proposed idea. Further, our method outperforms the performance gains obtained from MGD [Yang et. al., 22] as seen in Table 14

We have cited the work of DiffKD in related work. Based on our understanding (Fig. 3 in DiffKD), it appears extra modules are added during inference which changes the student architecture, which is not fair. However for completeness, we compare against it in Table 14. Proposed CIFD outperforms DiffKD except for top-1 accuracy in RN50 to MNV1 distillation. It takes the combination of two works DiffKD [25] and DistKD [5] to obtain better top-1 performance than CIFD in both RN34 to RN18 and RN50 to MNV1. However, our top-5 accuracy is better for both cases, indicating the competitiveness of proposed CIFD.

We respectfully disagree with the characterization that CIFD is just noise addition and removal process and thus similar to MGD & DiffKD, latter only removes noise. The key differentiator are the principled loss functions used to train the models in the presence of noise. If we were to just add and remove noise from student's embeddings, then it would result in a simple, one-stage procedure, similar to MGD/DiffKD. The central theme of the paper is to control information flow during distillation and one of the best tools to accomplish it is by adding noise. Further, the task accomplished by the noise in the teacher (mimic TAs) is different from the student model. Finally, the effect of this is seen in the performance gains obtained by our proposed method

Student architecture changed?

We do not modify the student architecture. We designate a layer in the unmodified student model as the bottleneck (usually the penultimate layer). The model's layers preceding the bottleneck acts as the IB encoder, and the layer(s) following as the decoder. Noise is added to the output of the bottleneck layer during training and noise addition is disabled during inference

We again thank the reviewer again for their time and response. We greatly appreciate their interest in helping us improve the paper.

Moving discussion between IBP and MIM to main paper

• We absolutely agree that having this discussion in the main paper will enrich it. In fact, it was in the main body of the initial draft. However, the space requirements proved very restrictive and we were forced to move the discussion to the appendix. We agree with you that it should be put back into the main paper and plan to do so in the final version.

Although the motivations seem novel, the main improvements are from multiple RDMs.

Difference between [1], [3] and our work.

While the multiple projection approaches has existed as pointed out in [1] (which is also cited as [20] in our paper), we performed comparisons specifically with [1] in the rebuttal pdf (Table 15). Note that [1] and [3] use multiple projectors at the student (see Fig. 1(c) in [1] and Fig. 1, right subfigure in [3]) and our method uses multiple RDMs at the teacher. Thus, these methods are quite different. For completeness we include comparisons with these methods in the table below. However, our method with 1 RDM and IBM performs close to [3] and the with 3 RDMs outperforms it. Additionally, our method with 3 RDMs and IBM outperforms [3] with 8 projectors at the student.

To further answer your concern about whether the claimed method of RDM training is giving the performance improvement, we compare with Factor Transfer Kim et. al [28] in Table 15 (rebuttal PDF). Our proposed method with one RDM and no IBM outperforms FT [28], where the difference is the loss function used to train the RDM. This shows that the improvement is coming from the new principled loss function used to train the RDM network.

Difference between [2] and ours.

We have to note that [2] is for Online Distillation without a teacher. This setting itself is different from ours where use the presence of a teacher. While [2] does use an ensemble of projection modules at the teacher, the difference in setting used to train makes it hard to compare them. Specifically they use a gated ensemble of branch modules to simulate a teacher using the student's backbone. In the below table we include [2] for reference only, the setting is different and they should not be compared directly.

+------------------------------+-------+-------+
|     RN34 to RN18 on IN1k     | Top-1 | Top-5 |
+------------------------------+-------+-------+
|       FT (Kim et. al.)       | 71.43 | 90.29 |
|           IFD [1]            | 71.94 | 90.68 |
|  ONE [2] (No RN34 teacher)   | 70.55 | 89.59 |
|           NORM [3]           | 72.14 |  ---  |
|   ------------------------   | ----- | ----- |
| Proposed (1 RDM) and no IBM  | 71.83 | 90.69 |
| Proposed (1 RDM) and wt IBM  | 72.05 | 90.70 |
| Proposed (3 RDMs) and wt IBM | 72.32 | 90.88 |
+------------------------------+-------+-------+

Summary

We have shown over multiple experiments the efficacy of our proposed method. See Table 2 on CIFAR-100, Table 3 and Table 14 on ImageNet, Table 13 for large student-teacher gap in ImageNet, Table 5 on CLIP models. Specifically we have addressed the question if our proposed method is reason for performance improvement by disabling various parts of the proposed idea and show improved performance with existing methods (see Table 15). These specifically showcase the importance of the proposed loss functions and their improved performance.

We again thank the reviewer for their time and effort. If you have any further questions, please feel free to reach out to us. Thank you!

Thanks for your detailed response, which addresses my partial concerns.

Here are my remaining concerns.

I personally believe that lacking deep discussions about related techniques in the main paper is inappropriate. I encourage the author to add this content in the main paper, not just in the Appendix.

Although the motivations seem novel, the main improvements are from multiple RDMs. Similar distillation designs have been explored in [1-3]. On the other hand, when reading the paper and the authors' response, I find no obvious technical flaws. Therefore, I'm a little bit confused about whether the effectiveness comes from the theoretical framework that the author claimed.

As I think the above arguments are not strong enough to reject the paper, I would consider my final rating according to both the response and other reviewers' comments.

[1] Chen, Yudong, et al. "Improved feature distillation via projector ensemble." Advances in Neural Information Processing Systems 35 (2022): 12084-12095.

[2] Zhu, Xiatian, and Shaogang Gong. "Knowledge distillation by on-the-fly native ensemble." Advances in neural information processing systems 31 (2018).

[3] Liu, Xiaolong, et al. "Norm: Knowledge distillation via n-to-one representation matching." arXiv preprint arXiv:2305.13803 (2023).