Regarding the Weakness 4:

Under limited computational resources, we added the result of Llama 13B-7B teacher-student pair in the SFT tasks. The batch size is set to 16 due to our limited computational resources. Our SHD surpasses the MiniLLM by 1.1% on UnNI. We showed the general ability over relatively large and relatively modern language models.

Regarding the Weakness 5:

We are very sorry for the missing citations. We have added them to our citations.

Regarding the Weakness 6:

We chose the example randomly. We watched around 1000 inference samples and they all behave like this simple example. The redundancy in the behavior of heads is similar. We will add some visualized results of this phenomenon to our project site due to the limitation of pages of ICLR paper.

Regarding the Weakness 7:

We thank the reviewer for all the writing problems mentioned and we fixed them in the manuscript. Thanks so much for pointing them out. The precision and recall metrics are from [1]. The basic concept is similar to precision and recall in object detection.

[1]:improved precision and recall metric for assessing generative models. NIPS2019

The authors sincerely thank the reviewer for the very complete and detailed comment. The suggestions of the reviewer are very helpful and insightful. We added and modified some contents in the paper as the reviewer suggested.

Regarding the Weakness 1:

We are sorry for the confusion we caused in the part of the merging process. As the reviewer said, we did the merging process progressively since each time of the merging process can make the number of heads halved. The result in Table 3 of the GPT2-XL teacher(25 heads) and GPT2-small student(12 heads) used SHD to merge heads progressively twice.

We also did the experiments of matching heads with each other before training and computed the costs of head matching, so that the attention maps of merging heads have the highest similarity of attention maps on the total training dataset. We added the performance of head matching in Table 4.

Regarding the Weakness 2:

We reproduced VkD in the supervised fine-tuning tasks of LLM and migrated it from their official implementation as the reviewer suggested. We searched for the best hyper-parameter for VkD and reported the best result in the ablation study of Table 5. Please notice that our implementation uses fp32 to train VkD since the cuda version of the orthogonal projector in PyTorch only supports fp32. The performance of our SHD is reported in fp16 training and still has improvement over VkD. Besides, VkD is a simple and cost-effective constrained feature distillation pipeline. Our approach, however, is even more cheap—it's practically free.

Regarding the Weakness 3:

We added the results of discriminative tasks as the reviewer suggested. We compared SHD with NKD, and ViTKD which focuses on the knowledge distillation of ViT-ViT teacher-student pairs. Our SHD still gains improvement whether independently or over ViTKD and shows compatibility with feature distillation and its generability.

The result of ViTKD+NKD+SHD also shows the compatibility of SHD with other FD methods, improving the performance of ViTKD+NKD by 0.42% on a strong baseline.

The authors want to thank the reviewer for raising our score to 8 and admiring our revision again.

We noticed that the reviewer's confidence score is still 3. Please let us know if there are any further questions we can answer to raise the confidence of the review. We are very happy to answer them.

Best regards, Authors of Submission 1443.

We want to thank the reviewer for the very detailed comment.

Regarding the Weakness 1:

The authors agree with the reviewer that PCA, optimal transport, and kernel methods can be better tools. However, all these tools are time-consuming for training a model, since we need to do it online. We need a practical and useful method to compute the supervision of attention maps. And convex combination is a great way of being practical and it can significantly improve the performance in our experiments as well.

Regarding the Weakness 2:

There seems to be a misunderstanding. It is the global minimum in the convex combination once the heads’ order to squeeze is determined. The merging coefficients in the range of [0,1] can guarantee the positiveness of attention maps, and they always fall into this range as we depicted in Fig. 2 in the appendix. We also reported the results of using the training dataset and teacher models to match teachers’ heads with each other as "head_matching". This should be the global minimum overall.

All other SHD experiments we reported, if not specified, did not use head matching since it requires evaluating the whole training dataset before training. We want our method to be plug-and-play and easy-to-follow.

Regarding the Weakness 3 and Question 2:

This is a practical problem that we always use the heterogeneous architecture where two models have different attention heads. Large teachers have large heads and small students have small heads in the knowledge distillation of transformers. Our work aims to transfer knowledge to train better models with smaller sizes. Also, our method aims to tackle the problem of misalignment of teacher heads and student heads. It does not include the situation where teacher and student have the same heads, since it's pretty rare.

The authors thank the reviewer for telling us that most of the reviewer's concerns are addressed.

We want to address the concern of the linear assumption from three terms: 1)the mathematical term, 2)the efficiency term, and 3)the effectiveness term.

●1)Mathematical Soundness:

Convex combinations maintain the essential probabilistic properties of attention maps, which is crucial for effective knowledge distillation.While methods like PCA or kernel methods can capture correlations, they may not preserve the non-negativity and row-sum-to-one properties of attention maps. Therefore, convex combination of merging coefficients between [0,1] is mathematical important. We also showed that the distribution of coefficients we computed in SHD in the Appendix. The global minimum coefficients always fall into this range, proving our assumption to be rational.

●2)Computational Efficiency:

The linear complexity of our method ensures practicality and scalability, unlike computationally intensive alternatives like optimal transport. Our linear approximation algorithm operates with a complexity of O(N^2) whereas optimal transport solvers, such as Interior Point Methods, are computationally more expensive, requiring O(N^3） complexity, which is unacceptable in practical training.

●3)Effectiveness:

We want to thank the reviewer for admiring the thoroughness of our experiments on four different tasks, including Image Classification (discriminative task), LLM SFT, LLM Pretraining, and Image Generation (generative tasks), where SHD shows consistent improvement over baselines. We also compared with projector-based methods, which can be seen as a variant type for kernel methods, and outperformed these projector-based methods both on training speed and performance.The effectiveness of the linear assumption is further proved by these experiments.

While PCA, optimal transport, and kernel methods have their own merits, they introduce complexities and limitations that make them less suitable for our specific context according to the explanations above. We believe that our approach strikes an optimal balance between theoretical rigor, practical efficiency, effectiveness, and interpretability.

We sincerely and kindly ask the reviewer to reconcile based on the explanation above since it is the only issue as the reviewer mentioned. We are open to any further requirements by the reviewer.

We want to thank the reviewer for the comment.

Regarding the Weakness 1:

The observations that the reviewer referred to were mentioned in previous papers[1] and are not what we mainly discussed in the paper. Our paper mainly focuses on our method: SHD. [1] Are Sixteen Heads Really Better than One? (NIPS2019)

Regarding the Weakness 2:

We added the comparison between SHD, ViTKD, NKD, and V_kD, FD with projector and FD with self_correlation.

We thank the reviewer for the thorough response.

Regarding the Weakness 1:

As the reviewer said, the feature projectors in the knowledge distillation of image classification are quite lightweight since most FD methods only focus on a single layer of features in the backbone network (mostly features of the last layer before the classification head) due to the limitation of computational costs and performance. We also want to mention that most FD methods experimented on CNN-CNN/CNN-ViT teacher-student pairs in the image classification or detection tasks. When it comes to generative models, some research has suggested that the supervision of FD from each layer is a better solution in the training [1]. Our ablation studies between FD and SHD are computed on all layers of the student model. The training parameters and the training speed of FD, VkD[2], and SHD are reported and added to our new manuscript. Here is the result:

Here we present an overall result of FLOPs. Consider a training scenario where the batch size is denoted by $b$, the hidden layer dimension by $h$, and the sequence length by $s$. The total FLOPs for the self-attention and MLP layers during both the forward and backward propagation processes for each layer is given by $(24bsh^2 + 4bsh^2) \times 3$. Taking the linear projector as an example, the additional computational cost is $4bsh^2 \times 3$. Our method does not incur any extra overhead for backpropagation. The additional computational cost introduced by our method is $2bsh^2$. Therefore, the proportion of additional computation due to the linear projector (relative to the student network) is $\frac{1}{6 + \frac{s}{h}}$, and the proportion of additional computation due to our method is $\frac{1}{6 + \frac{36h}{s}}$. Taking $h = 1024$ and $s = 512$ as an example, the proportion of additional computation due to the linear projector relative to the original model is 15.38%, while our method is approximately 1.28%.

[1]Patient Knowledge Distillation for BERT Model Compression, EMNLP2019 [2]VkD: Improving Knowledge Distillation using Orthogonal Projections, CVPR 2024

Regarding the Weakness 2:

As Reviewer ARUU suggested, we further compared VkD (CVPR 2024) with SHD, one of the state-of-the-art methods in FD. Our method outperforms VkD on all five test datasets.

Also, we want to mention that SHD is not incompatible with most FD methods. The further discriminative experiments we added to the manuscript have clarified that. We added SHD to ViTKD[3], which mainly focuses on distilling features for ViT-ViT teacher-student pairs. The SHD can still improve the accuracy with a strong FD baseline.

[3]ViTKD: Practical Guidelines for ViT feature knowledge distillation

Regarding the Weakness 3&4:

As mentioned before, we added discriminative experiments to our paper thanks to all reviewers. We did the single ablation study without logits-based KD in the discriminative tasks to show that our SHD can still improve performance significantly.

The reason why we reported performance with logits-kd methods for most experiments is that we want to show that even strong baselines can benefit from SHD and logits-kd methods are mostly near-cost-free. We did experiments on image classification, image generation, LLM pretraining, and LLM SFT, on which SHD shows consistent improvement.