Optimizing Knowledge Distillation in Transformers: Enabling Power of Multi-Head Attention without Alignment Barriers

Authors: Zhaodong Bing, Linze Li, Jiajun Liang
Primary area: generative models
Keywords: Knowledge distillation; Multi-head Attention; Transformers; Generative models;

Abstract:
Knowledge distillation has been proven effective for compressing transformer architectures by transferring knowledge from teacher to student models. Logits-based methods of knowledge distillation cannot fully capture the intermediate representations and features within the teacher model, which may result in the student model not fully learning all the knowledge from the teacher model. Thus, previous work focuses on transferring knowledge through intermediate features or attention maps. However, leveraging multi-head attention maps in transformers for knowledge distillation presents challenges due to head misalignment and suboptimal feature alignment, often requiring projectors to align features or special modifications to the model architecture. To address above limitations, we propose the Squeezing-Heads Distillation (SHD) method. This method reduces the number of attention maps to any desired number through linear approximation, without requiring additional projectors or parameters. This facilitates better alignment and knowledge transfer between models with different numbers of heads, enhancing both flexibility and efficiency. Experimental results demonstrate significant improvements in both language and vision generative models, validating the effectiveness of our method.

Decision: reject
Reviewer scores: 5, 5, 5, 8

Reviews and discussion:
<thread>
creation date:
1730396374438

thread_type:
official_review

<thread_object>
creator: reviewer_aruu

<content>
summary:
The paper proposes a new method which tries to solve the issue of head misalignment in multi-head attention mechanism for knowledge distillation. The main idea of the paper is to compress multiple attention heads into a single one, and distill this knowledge. The authors show  experimental results in image generation, (small) LLM pretraining and finetuning, outperforming several other methods. Furthermore, they present an excellent collection of ablation studies, in regards to relation with other distillation models, main hyperparameters, alternatives to their method and the loss function.

strengths:
The paper has the following strengths:

1) Motivation - the paper is naturally well-motivated, starting with some clear observations: i) the head redundancy in Transformers; ii) alignment missmatch where teacher and student have different dimensions (which they almost always do); 3) attention matrices being rank deficient. While all of these have been observed in the past, it is still a good way of formalizing them and then giving the solution by proposing the attention head compression using the linear optimization method (section 4.2).

2) The method provided here to merge the attention heads is elegant, mathematically justified, and to me makes sense. Its derivations is also quite complete and just.

3) The method is shown to improve over standard knowledge distillation models in different settings: Image generation, small LLM pretraining, and small LLM finetuning. The results in all cases are quite significant, convincingly outperform the baseline and in some cases are (almost) as good as the much larger teacher.

4) The ablation study at the end of the paper is excellent, showing the robustness of the method under different hyperparameter settings, comparison with other methods (hard selection, constant merging of heads) and comparing the KD loss with MSE.

weaknesses:
I think the paper has some room for improvement:

1) While the paper explains how 2 attention heads can be merged, the paper mentions that multiple attention heads can be merged. That part is quite unclear to me and would be nice if the authors can clarify the following points:

a) How this is done, is it simply merging each pair of attention heads, thus reducing them by half, or if there is some iterative procedure.

b) If there is some iterative procedure, how this is done in the first place, is it done for several heads alltogether on in pairs. If for several heads, how does the procedure in 4.2 is extended, and furthermore do the alphas need to sum up to 1?

c) How are the heads to be selected chosen in the first place?

2) The size misalignment where the teacher has larger hidden dimensions than the student, is a real issue. Simply projecting the features of the teacher to smaller dimensions, or those of students to larger dimensions, usually destroys information making the KD process not work. However, there are some remedies of this that completely mitigate this issue, for example by using orthogonal projections, as the following paper, making the standard KD techniques work.

[A] Miles et al., VkD: Improving Knowledge Distillation using Orthogonal Projections, CVPR 2024.

While this paper and [A] target the problem quite differently, the motivations are quite similar, so it would be interesting if the paper can be discussed and compared with under some setting (image generation being the closest setting in them).

3) All the results are shown in generative tasks. It would be interesting to see if the method works in discriminative tasks (e.g., image classification and/or object detection), showing the method to be more general.

4) All the NLP results are shown in relatively small LLMs. It would be interesting to see if the method scales well to larger real-world LLMs. From my experience, this is always not the case, and what works in 'small' (tens to hundreds of millions of weights) does not neccesarily work in 'large' language models (billions of weights). If possible, it would be interesting to see how the method works when distilling the knowledge in LLama models. The ideal case would be Llama 13B -> Llama 7B, if that is not possible then something like Llama 7B -> TinyLlama 1.1B. Would be fine to use different alternatives to Llama (Mistral, Qwen-2), but it would be nice if they are a) relatively modern, b) relatively large.

5) Missing citation. The observation in equation 3 and 4, the idea that value-output multiplication can be considered as a single circuit (similar to query-key) has been done before, and it would be nice for the authors to cite it.

[B] Elhage et al., A Mathematical Framework for Transformer Circuits, Anthropic 2021, https://transformer-circuits.pub/2021/framework/index.html

6) Some clarification in Figure 1 is needed. If the results are from a single example, then that is very obviously expected. Of course that most attention heads are not going to produce anything meaningful (get 'activated'), that is by design. If not:

a) How did the authors choose the example?

b) Wouldn't it be better to check lots of inference examples, and merge the results?

7) Some writing issues:

a) Please rewrite the contributions in active instead of passive.

b) Figure 1 cannot be well read, and cannot be read at all if it is in printed paper. The authors should redesign it so it is more readable (probably by making it larger)

c) It is unclear to me what precision and recall means in the context of image generation (Table 1).

questions:
I have already included several questions in the weaknesses. In particular:

1) The authors should ideally clarify the questions under weakness (1) and (6)

2) The authors should give credit to papers [A] and [B] and ideally compare with paper [A] under some setting.

3) If possible, the authors should compare their method on other settings: a) discriminative tasks; b) larger LLMs.

I am initially positive, the paper makes sense, seems to work in different settings, and reach good results. Thus, I am scoring it initially as (6) Borderline accept. However, I acknowledge that the paper has some issues with clarity, more comparisons, writing, and putting it in context of other works. Providing that the authors would well-address the mentioned issues, I will switch to (8) Accept.

---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

After rebuttal:

I thank the authors for their excellent and detailed rebuttal.

1) I expressed some confusion about the merging process. The authors have adequately addressed this and provided additional results.

2) I mentioned VkD as missing reference and a recent paper that addresses a similar problem. The author have integrated the comparison with it in the manuscript, showing that their method outperforms VkD.

3) I mentioned lack of discriminative tasks. The authors have compared their method with other methods in a discriminative benchmark, showing that it outperforms the other methods.

4) I mentioned lack of experiments in real LLMs. The authors have shown an additional experiments in LLama13B-LLama7B setting.

5-6-7) The authors have adequately addressed my concerns about clarify, missing citation and metrics.

I was initially positive about the paper. After the rebuttal, the authors have adequately addressed all of my concerns, and have provided results that further improve the paper. Furthermore, checking the criticism from other reviewers, and the response by the authors, I do not think that the paper has any weakness that might warrant blocking the publications.

I think this is a very good to excellent paper. As promised, I am increasing my score to 8 and I am happy to further champion it.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730606570695

thread_type:
official_review

<thread_object>
creator: reviewer_sz2t

<content>
summary:
This work investigates the knowledge distillation specially tailored for transformer architectures. Specifically, the authors study the misalignment of attention heads between the teacher and student. They propose the squeezing-head module to compress multiple attention heads into single one, alleviating the mismatch of attention heads between different architectures. More specifically, the authors formulate the attention heads as a convex hull and aim to seek a substitute with minimal reconstruction error. Experiments on image generation and natural language processing show effectiveness.

strengths:
1. The motivation is important to KD in transformer architectures. Due to computation budget, different models vary from depth to width. This work proposes to address the problem of misalignment of attention heads between teacher model and student model without learnable projectors.
2. The experiments are thorough. The authors show competitive results of their method on various tasks and dataset.

weaknesses:
1. I have some concerns about the formulation starting from line 258. The authors formulate the *random convex hull* among multiple attention heads (*vertices*) and aim to find a substitute with minimal reconstruction error through **convex combination**. The problem is that no evidence reveals and supports the linearity (Eq. 5, line 266)  between two attention maps. In other words, it is less meaningful if we use the linear combination of two attention maps to approximate the other attention map. As an alternative, I think PCA, optimal transport, and kernel method can be the better tools.
2. It seems like Eq. 9 (line 291) does not guarantee global minimum. 
3. The experiments mainly focus on the heterogeneous architecture where two models have different attention heads. I wonder about the result of homogeneous architectures.

questions:
1. If teacher and student has the same attention heads, can this approach consistently beats existing KD methods?
2. In addition to efficiency and well alignment, one advantage of multi-head attention over single-head attention is their descriptive power. How is Figure 1 after you combine multi-head into single head?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730616663399

thread_type:
official_review

<thread_object>
creator: reviewer_znbn

<content>
summary:
This paper examines knowledge distillation in transformer architectures. Based on three observations in transformer attention maps, the proposed method investigates attention compression through linear approximation, using the squeezed attention maps for knowledge transfer between teacher and student models. Experiments are conducted on image generation and language model training tasks to demonstrate the effectiveness of the approach.

strengths:
- This paper addresses a critical challenge in the era of large language models (LLMs): compressing transformer models using knowledge distillation techniques.

- The proposed method is easy-to-follow.

weaknesses:
- Observations on transformer attention maps lack clarity and sufficient justification. For instance, the statement that “the number of heads is redundant to some degree” is not well-supported.

- The experiments do not adequately demonstrate the effectiveness of the proposed method. Key baseline and comparison methods are missing, which makes it difficult to evaluate the method’s impact on improving knowledge distillation in transformer architectures.

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730807138341

thread_type:
official_review

<thread_object>
creator: reviewer_up1v

<content>
summary:
The paper proposes a new feature-based knowledge distillation method, specially tailored to Transformers. Specifically, the method reduces the number of attention maps to any desired number through linear approximation, without requiring additional projectors or parameters. The experiments are conducted on both vision tasks and language tasks, validating the effectiveness of the method.

strengths:
1. The method is straightforward and reasonable, which is easy to understand. 
2. The experiments are performed on both vision tasks and language tasks.

weaknesses:
1. Fundamentally, the paper criticizes previous feature distillation methods because they often require projectors to align features (or special modifications to the model architecture), which introduce additional training costs. However, to the best of my knowledge, the prevalent of feature projectors (such as linear projectors, conv projectors, and MHA projectors) are rather lightweight, compared with the overhead of the teacher model and the student model.  The paper does not provide the FLOPs and training speed comparison between FD and SHD to show the necessity of designing SHD.

2. Furthermore, although the paper criticizes previous FD methods, it only compares with the most basic FD (with the basic linear projector) in the ablation study. The paper should cover the performance comparisons with various Top FD methods to enhance this claim. 

3. In these experiments, it seems that SHD must be applied with logit-based methods and the improvement seems rather limited.  I wish to see the results that SHD independently applied to the network pair.


3. Although SHD does not have special designs for generative tasks,  all experiments are performed on image-generative tasks and text-generative tasks. I wish to see more results on image classification tasks, which are feasible to make comparisons with various Top FD methods.

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732096578112

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
We greatly appreciate the reviewers' detailed responses and insightful suggestions. Your reviews truly help us polish our paper better. We sincerely appreciate that all reviewers believed our method is easy-to-follow,  straightforward and our experiments are thorough.  Additionally, Reviewer  ARUU affirmed our motivation and mathematics, improvement in all kinds of experiment settings, and the completeness of our ablation studies.  Reviewer Sz2t also appreciates the importance of solving the problem of misalignment of attention heads between the teacher and the student without learnable projectors. We will respond to each reviewer's questions and comments, respectively. We also revised our manuscript and updated it according to the reviewers' suggestions. Major modifications are highlighted in blue color.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732350946102

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
Dear Reviewers,

The authors sincerely appreciate the time and effort you are dedicating to reviewing our paper, "Optimizing Knowledge Distillation in Transformers: Enabling Power of Multi-Head Attention without Alignment Barriers".

As the rebuttal phase progresses, we wanted to kindly follow up on the status of your feedback. One of the reviewers **ARUU** has already shared detailed and encouraging comments, which we are truly grateful for. The reviewer **ARUU** thinks that our revision addresses the problems of clarity, more comparisons, writing, and putting it in context of other works. **The reviewer ARUU also believes that “our paper has no any further weakness that might warrant blocking the publication” by checking the criticism from other reviewers and the response by the authors.** 

We cannot be more thankful for all the precious comments. We would greatly value your input as well to ensure a thorough and balanced evaluation of our work.

**Given the approaching deadline for final decisions, we kindly ask if it would be possible to share your feedback at your earliest convenience. Please do not hesitate to let us know if there are specific points we can clarify or address further.**

Thank you again for your valuable contributions to this process.

Best Wishes,

Authors of Submision 1443
</content>
</thread_object>

</thread>
<thread>
creation date:
1733112211513

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
The authors want to thank the reviewers for the efforts in the reviews.
As the deadline of rebuttal disccusion is approaching, we kindly ask if it would be possible to share your feedback at your earliest convenience. Please do not hesitate to let us know if there are specific points we can clarify or address further. The authors have added promising results and replied to the concerns and weaknesses. We are confident with our method and experiments results and williing to answer any further questions. 

We believe your suggestions and our efforts get our paper improved and meet the high standard of ICLR.

Best regards,

Authors of submission 1443
</content>
</thread_object>

</thread>
<thread>
creation date:
1733157543801

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
Rebuttal Summary
We sincerely thank all the reviewers for their detailed and constructive feedback,**especially reviewer ARUU, who has the most insightful suggestions for the papers, the fastest responses to our rebuttal, and the most paper-related questions. The reviewer ARUU also believes that “our paper has no any further weakness that might warrant blocking the publication” by checking the criticism from other reviewers and the response by the authors.** Below, we summarize our key responses to address the main weaknesses raised by the reviewers.
***

## **Regarding Weakness 1: On computational efficiency and feature distillation (FD)**



●Feature Projectors: We clarified that our ablation studies involve comparisons across all layers of the student model, contrasting FD with VkD [2] and our method, SHD. Results show that SHD achieves comparable efficiency while avoiding the overhead of projectors and enhancing training speed and parameter efficiency.

●Computational Complexity: We provided a detailed analysis demonstrating that SHD introduces significantly less computational overhead (1.28%) compared to linear projectors (15.38%).

●Mathematical Justification: Convex combinations, a foundation of SHD, ensure attention map properties (non-negativity, row-sum-to-one) while maintaining computational practicality and theoretical rigor.

***

## **Regarding Weakness 2: Compatibility with other methods and performance comparisons and Lack of discriminative tasks**

**We added SHD results on image classification of ImageNet-1k.** 

●State-of-the-Art Comparisons: SHD was compared with VkD [2] and ViTKD [3] across tasks like image classification, LLM pretraining, and supervised fine-tuning, **consistently outperforming these baselines.**

●Independent Effectiveness: Ablation studies showed SHD improves performance even without logits-based KD, proving its standalone effectiveness (e.g., **+0.95%** Top-1 accuracy on DeiT-Tiny).

●Compatibility: We demonstrated SHD's compatibility with other feature distillation methods, showing it improves performance even when combined with strong baselines e.g. FD methods and logits-KD methods, achieving a **0.42%** improvement over NKD+ViTKD, **proving our compatibility with traditional FD methods and logits-KD methods.**

***

## **Regarding Weakness 3: Scalability**

●Scalability: SHD showed superior generality with modern, large-scale language models (e.g., LLaMA 13B → 7B), outperforming baselines under constrained computational resources.

***

## **Regarding Weakness 4: Linear Assumption in SHD**

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

**Mathematical Soundness**: Our method's use of convex combinations ensures that the probabilistic properties of attention maps are preserved, unlike PCA or kernel methods which may not maintain non-negativity and row-sum-to-one properties. The coefficients used in our method are mathematically justified, as they always fall within the [0,1] range, which is crucial for effective knowledge distillation.

**Computational Efficiency**: Our method offers linear complexity (O(N^2)), making it practical and scalable compared to more computationally demanding alternatives like optimal transport, which requires O(N^3) complexity. This efficiency is a significant advantage in practical training scenarios.

**Effectiveness**: The effectiveness of our method is demonstrated through extensive experiments across four different tasks, where our method (SHD) consistently outperforms baselines and projector-based methods in both training speed and performance. These results validate the effectiveness of our linear assumption and show that our approach offers a superior balance between theoretical rigor, practical efficiency, effectiveness, and interpretability compared to PCA, optimal transport, and kernel methods.


***

## **Regarding Weakness 5: Missing citations, Writing issues**

●We addressed the missing citations and added references to ensure proper acknowledgment of prior work. We also revised the manuscript thoroughly to address all writing concerns.

***

# **Conclusion and Final Appeal**

**We are pleased to confirm that all issues raised by the reviewers have been fully addressed through extensive experiments, theoretical analysis, and clarifications. The revisions strengthen the manuscript significantly, substantiating SHD's practicality, robustness, and effectiveness across diverse tasks.**

We deeply appreciate the reviewers’ thoughtful feedback and respectfully request them to kindly reconsider their initial scores in light of the comprehensive updates. **We believe the resolved issues, coupled with our clarified contributions and rigorous evaluations, warrant a positive reassessment of our work.**

We remain open to further questions or suggestions and sincerely thank the reviewers for reviewing our submission.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a distillation method for transformer-based models. The main idea is to reduce the number of attention maps to any desired number through linear approximation, which is simple but effective. However, the authors have not sufficiently demonstrated why the linear assumption is valid or clearly articulated the motivation and advantages of the method. Additionally, it would be beneficial to illustrate the main idea with figures to enhance understanding. Based on these strengths and weaknesses, the decision is not to recommend acceptance at this time. We encourage the authors to carefully consider the reviewers' comments when revising the paper for submission elsewhere, particularly in appropriately articulating the motivation and advantages of the method.