Talking Models: Distill Pre-trained Knowledge to Downstream Models via Interactive Communication

To answer reviewer's questions:

1. How dose this approach perform on natural language processing tasks?

We actually conducted additional experiments on language modeling tasks using T5 to show by asking questions to T5 Large, we can significantly improve the T5 small student on downstream benchmarks such as SuperGLUE. However, we decided not to include this result due to the design of encoding/decoding (through text prompt) and the communication process is entirely different to our current method, even though the main idea is the same. We want to focus on the general concept of interactive communication in this paper.

2. How is the approach compared with the most recent knowledge distillation methods?

As discussed above, we compared with the canonical ones from which our method builds upon and can be extended to our application. However, many most recent knowledge distillation for specific domains can be extended and integrated into our interactive communication framework.

3. How do you determine w_1, w_2, w_3?

We conducted hyper-parameter search on our method and baseline methods, the details of hyper-paramter search and results are discussed in Appendix 6.3 and 6.5.

[1] https://arxiv.org/abs/2106.01548

[2] https://openreview.net/forum?id=q6bZruC3dWJ

[3] https://openreview.net/pdf?id=0ltDq6SjrfW

[4] https://openreview.net/pdf?id=157Usp_kbi

We thank the reviewer for the on-point comments and insightful questions. We address the reviewer’s questions below:

Weaknesses:

1. Limited Comparison Baselines

We thank the reviewer for providing an additional baseline method. One thing we want to point out here is that our goal is not to improve knowledge distillation in general, instead, we want to extend knowledge distillation to the case where it can be used to help transfer knowledge efficiently from pre-trained large models to smaller models specific for downstream tasks. Therefore, we don’t compare with many KD baselines, we carefully select the canonical ones from which our method builds upon and can be extended to our application. Since we use intermediate representations, many state-of-the-art feature distillation methods (such as [4]) can be used in our framework by modifying the distance metric d in section 3.3. We carefully examine the paper the reviewer provided. It proposed a very scalable and general method for transfer learning with the kernel method. It seems non-trial to apply this method to our problem setup, which is knowledge distillation between pre-trained teacher model and downstream student.

2. Limited tasks.

We totally agree that expanding the scope of evaluation with more tasks would provide a more robust assessment. This is also the reason why we picked two datasets from totally different domains, i.e., image classification and recommendation systems. We believe that these two different domains can cover many popular use cases. We also studied natural language tasks, however, the task being formulated as generative language modeling tasks totally changed the paradigm of knowledge distillation. Therefore, we focus on verifying our method on the CV and recsys domains.

3. Insufficient Comparison with IAKD.

We thank the reviewer for pointing out IAKD. IAKD indeed shares similar motivation as our work. But the key idea of IAKD is very different and cannot be directly applied in our case. IAKD directly swapped blocks of models, which require models to be similar in model architecture. In the IAKD paper, the authors used resnet without a pre-training dataset. In our case, we use ViT pre-trained on ImageNet21k, which is much larger and cannot be directly swapped to students due to layer and shape mismatch.

4. Underwhelming Performance

We thank the reviewer for pointing out the SOTA results. And we want to emphasize that by using smaller models, it is still not realistic to achieve results better than or close to SOTA. This is actually one of the key discussions we want to raise in our paper (Section 5 and Appendix 6.5). In our humble opinion, we don’t think closing the gap between student and teacher model’s performance (fine-tuned teacher in this case, which is SOTA) is realistic for any distillation techniques. The gap caused by capacity and training resources (training extensively on pre-trained data) cannot be overcomed by simply letting one teach the other, as pointed out by many related works [2], [3]. But we still think it’s hopeful that using large pre-trained models can improve downstream models (compared to without distillation), especially when fine-tuning is not realistic in training (limited resource and access to pre-trained model) and serving (serving cost limitation).

Especially the training from scratch ViT model is not as good as some of the SOTA results, as discussed in [1], ViT models can be relatively hard to train from scratch. However, we want to highlight the relevant improvement that different distillation methods can bring with the help of a large pre-trained ViT model. We believe this is an important use case, because not only is fine-tuning a large pretrained model costly, there will be much more cost when actually serving the model. And in many use cases, the limitation on serving makes it impossible to use a large model hence smaller downstream model is our only choice. In this case, the proper baseline is the training from scratch smaller model, not the fine-tuned teacher model.

We thank the reviewer for the careful examination and insightful comments for our paper. We address the reviewer’s comments below: Weaknesses

1. The experiments on movie prediction only cover a narrow scope, and the teacher/student tasks are quite similar with student task is to predict movie from one genre. The results could be more convincing if more varied tasks are involved, and if the "gap" between teacher and student is larger.

We agree with the reviewer that our method can be further verified with more distribution and task differences. In fact, this is also one of the reasons we choose movielens to be one of our evaluation datasets. The user behaviors on movies from different genres can be entirely different, especially on genres with very few user interactions, such as documentary, it can have totally different user behaviors patterns compared to popular genres such as comedy. We intentionally split the genres to two categories, i.e., dense and sparse during our evaluation. We reported their results separately, where we can see improvement of our methods on both categories. Then on image classification, the distribution between CIFAR10, CIFAR100 and ImageNet can also be very different. We believe that our experiments covered different scenarios. We think it can definitely be improved via experiments in more extreme case, e.g., distilling language modeling teacher to movie lens students, we included this as one of the scenarios we want to verify in our future work.

2. The approach makes sense but quite straightforward by adding teacher receiving messages. It's worth more discussion on insights of this effect to the teacher model (if not frozen).

We thank the reviewer for the great insights. It is indeed very interesting to discuss the effect on the teacher model if not frozen. In our paper, we conduct analysis by comparing the representation between teacher and student on input images with different classes, where we have shown by introducing the encoder/decoder, we can better align their representations. It can be different when we enable fine-tuning. This is indeed one of our future work, however, some early results suggest fine-tuning doesn’t always align the representations between teacher and student. We added some discussion in our future work section.

We thank the reviewer for the comments, careful examination of our paper, and understanding of our key insights. And we address the reviewer’s comments below:

Weaknesses:

1. direct evidence is lacking to support its advantages. For example, why should the messages of the teacher and the student be consistent (L_{MC}), considering that they are produced by the two models sequentially? I doubt that the performance improvement largely comes from the four additional modules as they bring more parameters.

We thank the reviewer for the question. We provided some evidence that our interactive communication method works better than other one-way communication based methods by some case study included in Section 4.4. We showed that the ``undistillable’’ class between teacher and student can be distillable after the training of encoder and decoder. We believe the training of encoder and decoder with the consistency loss can provide two benefits: (1) Prevent inefficient knowledge transfer from domain shift between pre-trained data and downstream tasks. (2) Enable the teacher to understand what information students can provide, because the student cannot get complex representation from the input as the teacher due to capacity gap, this will allow the teacher to teach with limited representation suitable for the student’s capacity. We provided a detailed discussion on the current design regarding the reviewer’s questions in Appendix 6.1 (for design simplicity and future work). To verify our explanation and insights, we also show that (figure 3) using interactive communication, we can bridge the capacity gap between teacher and student for “undistillable classes”[1].

For the question on performance improved due to additional parameters: our method didn’t add significant computation costs compared to existing KD methods. Costs added by encoder and decoder are comparable to feature distillation methods. But our costs are (k-1 times) higher when we are doing interactive communication in k iterations. Besides discussion in section 4.3, we will add a quantitative analysis in our revision.

2. The experiments are not enough to support the advantage of the proposed method.

One thing we want to highlight is that our paper is not targeted at improving knowledge distillation in general. Our focus is to extend knowledge distillation to the application of efficient knowledge transfer from pre-trained large models to smaller models for downstream tasks. Therefore, different from general KD methods, our proposed method considers two additional challenges: (1) the teacher model is trained differently using pre-trained tasks. (2) fine-tuning the teacher to arbitrary downstream tasks is not practical. We believe this is a very important application as pre-trained foundation models have shown great advantages in different domains such as NLP and CV, and directly using such a large model is not realistic as fine-tuning is not practical, nor serving the fine-tuned large model.

Therefore, we don’t compare our methods to the KD algorithms which cannot be used in our application without non-trivial modification. Another reason is that since our application is on efficient knowledge transfer from large pre-trained models to smaller downstream models, we don’t want our method to be tied to a specific domain (e.g., CV), instead, we evaluate our method on various domains including recommendation system and CV. Therefore, many KD methods targeting a single domain, e.g., CV, or NLP, cannot be directly used to compare with our method.

However, with our proposed interactive communication framework, we can integrate many feature distillation methods to further improve KD in different domains. As an example, the two papers the reviewer pointed out can be very useful to improve KD for ViT and they can be integrated by modifying the distance function d in section 3.3 (for example, mimincing and generating in ViTKD can be used to improve state and message consistency losses).

3. The presentation is poor.

We will make significant improvements to make our presentation clear in our revision.

4. Some technical details are missing.

We thank the reviewer for pointing this out! We include our experimental designs in Appendix 6.2 and 6.3. For the losses introduced in interaction, only the interactive loss “L_{interact}” are added multiple times with equal loss. We made this clear in our revision.

[1] https://openreview.net/forum?id=q6bZruC3dWJ

We thank the reviewer for the thorough comments and insightful questions, especially pointing out the weak connection between our method and distribution shifts. We address the reviewer’s comments below:

Weaknesses:

1. it would be nice to make the paper more formal, e.g. by defining different task distributions and the problem you are to tackle.

Thanks for the suggestion, we added a formal definition of our problem in our Appendix, along with the formal description of the algorithm. However, the definition of task/data distribution seems hard to be added in the formal definition since we only empirically verify our hypothesis on the distribution shift between pre-training and downstream datasets. The theoretical quantification of the differences is still challenging.

2. The link between the method and different task distributions does not seem very clear.

The connection between our method and task distribution lies in the pre-training fine-tuning paradigm of large foundation models. In this paper, we are not improving the knowledge distillation algorithm in general, but extending knowledge distillation to a very specific use-case: leveraging the knowledge of pre-trained large models to teach smaller student models for downstream tasks, without extensive or at-all fine-tuning. We believe this is a very important application as pre-trained foundation models have shown great advantages in different domains such as NLP and CV. In this case, the teacher cannot adapt to the downstream tasks’ distribution, which is often different from the pre-training distribution. Therefore, we propose to use the encoder and decoder to help the teacher learn the downstream tasks’ distributions without modifying the teacher model. We strengthened our introduction to make this clear.

3. I do not find the first half of the paper (section 1 and 2) very informative. I think empirical studies in 4.2-4.4 are especially useful for justifying such type of approach, but regrettably they are not highlighted in the main paper.

Thanks for pointing this out! Following your suggestion, we modified our paper to mention our empirical study results in our introduction to help justify our approach.

Minors: better use vector graphics such as pdf rather than bitmap for figures. Thanks for the suggestion, we replaced our figures to pdfs.

To answer reviewer's questions:

1. How is hg and lg chosen? There are also many other hyperparameters, how are they tuned?

We don’t tune h_g and l_g, instead, l_g is just after the first input projection layer and h_g is before the last classification layer. Therefore, we enable teacher and student to align their input and label space via the encoder and decoder. We totally agree with the review that deciding which layers are lower and higher is tricky. The point we want to highlight is that we don’t tune this as hyper-parameters. We believe (and use the setup) that the lower should be as low as possible and higher can be as high as possible. Because this is to solve the challenge of domain and task differences between teacher and student. The question of what proper layer and techniques to distill can be partially solved using existing feature distillation techniques, as they can be extended into our framework by modifying the distance function d in the equations in section 3.3. There are discussions in some recent feature distillation works [1] and [2] on which layers and losses to use for feature distillation. We will add more discussion and clarification in our revision.

As for other hyper-parameters, we performed grid search of some of hyper-parameters such as loss weights for our method and baseline methods. We included the detailed hyper-parameter tuning in Appendix 6.3.

2. Can you provide more insights on the question in weakness 2?

Discussed above.

3. How is the approach related with foundation models, as teacher models are just some pretrained models, the same as standard KD setting?

The key difference here is that we assume fine-tuning teachers for downstream tasks is costly. And also the pretraining dataset can be very general and broad covering various distributions but the downstream tasks are relatively specialized and different from the pre-training dataset and tasks. We believe this differentiates our problem setup with the standard KD setting.

4. How distribution shifts are reflected in experiments?

We designed our experiments as the pre-trained teacher model is not fine-tuned for downstream tasks. For example, on Movielens, the pretrained model sees data from user ratings on all movie genres while the student is trained on a specific genre. We can see that students for different genres benefit differently from the teacher model using different baseline methods.

5. Can you discuss the connection with existing cross-task KD approaches?

Thanks the reviewer for pointing out this related literature. By comparing our method with some recent cross-task KDD approaches [3,4,5], our proposed method and problem setup have the following differences: (1) The data for training the teacher and student is largely different. Existing cross-task KD don’t assume the distribution differences on the input side between pre-training data and fine-tuning data. (2) Limited ability in modifying the teacher. (3) Method: These work adopt typical or modified feature distillation paradigm, where our method introduced the inter-active communication paradigm.

[1] https://arxiv.org/pdf/2209.02432.pdf

[2] https://arxiv.org/pdf/2104.09044.pdf

[3] Ye et al., Distilling Cross-Task Knowledge via Relationship Matching

[4] Li et al., Prototype-guided Cross-task Knowledge Distillation for Large-scale Models

[5] Yang et al., Cross-Task Knowledge Distillation in Multi-Task Recommendation