What Matters in Transformers? Not All Attention is Needed

Authors: Shwai He, Guoheng Sun, Zheyu Shen, Ang Li
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Transformer, Attention, Model Compression, Efficiency

Abstract:
While scaling Transformer-based large language models (LLMs) has demonstrated
promising performance across various tasks, it also introduces redundant archi-
tectures, posing efficiency challenges for real-world deployment. Despite some
recognition of redundancy in LLMs, the variability of redundancy across different
architectures in transformers, such as MLP and Attention layers, is under-explored.
In this work, we investigate redundancy across different modules within Trans-
formers, including Blocks, MLP, and Attention layers, using a similarity-based
metric. Surprisingly, despite the critical role of attention layers in distinguishing
transformers from other architectures, we found that a large portion of these layers
exhibit excessively high similarity and can be pruned without degrading perfor-
mance. For instance, Llama-2-70B achieved a 48.4% speedup with only a 2.4%
performance drop by pruning half of the attention layers. Furthermore, by tracing
model checkpoints throughout the training process, we observed that attention
layer redundancy is inherent and consistent across training stages. Additionally,
we further propose a method that jointly drops Attention and MLP layers, allowing
us to more aggressively drop additional layers. For instance, when dropping 31
layers (Attention + MLP), Llama-2-13B still retains 90% of the performance on the
MMLU task. Our work provides valuable insights for future network architecture
design. The code will be released upon acceptance.

Decision: reject
Reviewer scores: 6, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1729974604473

thread_type:
official_review

<thread_object>
creator: reviewer_pbgt

<content>
summary:
This paper describes a training-free method for pruning transformers models.

A similarity measure is proposed to estimate the amount of transformation a block (or a layer) apply to their inputs. The similarity measure can be used to define an importance measure which will guide the pruning process.

The importance measure is first applied to prune entire transformer blocks (attention+MLP) from the model. The paper then explains how to apply this measure to estimate the importance of each sublock (attention or MLP) by looking at the residual element within the sublock.
The paper provides empirical block-wise importance data for several popular models, and shows that the earlier and last layer of these models typically exhibit greater importance than the shallow layers.

The paper introduces a Speedup Degradation Ratio metric to help assess the tradeoff between speed and accuracy while pruning a model.

The paper provides KV cache and latency measurements, as a function of how many blocks/layers are pruned.

strengths:
The paper is well written, it is easy to follow and the claims are backed with empirical results, latency measurements and visual insights.

The method is simple and looks easy enough to implement and replicate.

weaknesses:
The novelty of the paper is limited. The paper's main findings (that attention is more easily pruned than MLP, and that shallow layers are more easily pruned than first and last layers) are known, see for example "A deeper look at depth pruning of LLMs" (https://arxiv.org/pdf/2407.16286).

The importance measure is using the scale-invariant cosine similarity measure. It could be argued that this fails to capture magnitude information. Since the cosine similarity measure only depends on the orientation of the vectors, a block that doesn't change the orientation of its input but changes its magnitude could be deemed unimportant. In that case, it might be considered to replace said block with a scaling factor.

There are no comparisons against other methods in the main part of the paper. The appendix has a comparison against ShortenedLLaMA and Wanda. There is no comparison against methods such as LLMPruner, SliceGPT, ShearedLLaMA, or Minitron.

The paper does not show a study on the potential benefits of further fine-tuning after pruning, which could help recover some of the accuracy loss.

The paper does not study how to achieve a finer level of pruning, for example pruning individual attention heads.

The paper does not study the importance of the calibration dataset, or whether domain-specific datasets could be used to improve task-specific benchmarks.

questions:
Can you add a comparison against one or more of LLMPruner, SliceGPT, ShearedLLaMA, or Minitron?

In section 5, it is stated that "For instance, in Llama-2-13B, the KV-cache is reduced from 52GB to 26GB, a 50% reduction. This memory reduction is even more pronounced in larger models like Llama-2-70B, where the KV-cache decreases from 20GB to 10GB.", and the data in the table match this statement. Is this not the other way around, that the 52GB-to-26GB reduction relates to LLama-2-70B?

Did you try to evaluate whether you can prune individual attention heads, as opposed to the whole attention layers?

Did you try to evaluate whether fine-tuning after pruning helps?

Did you try to evaluate whether a scaling factor could be introduced in lieu of a pruned block?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730603449433

thread_type:
official_review

<thread_object>
creator: reviewer_rs71

<content>
summary:
- investigates the redundancy within different modules of Transformer-based large language models, including Blocks, MLP, and Attention layers.
- The redundancy is evaluated using a similarity-based metric.
- This paper validates that some model components can be pruned with obvious speedup and minor performance drop.

strengths:
- The paper is clearly written and well organized, it provides a systematic exploration of redundancy in Transformers, focusing on Blocks, MLP, and Attention layers.
- This paper provides several useful insights, for example:
	- FFNs seem more important and Attention modules can be dropped with minimal performance impact with high efficiency
	- deeper layers seem less important compared t the shallower ones, which indicated the model has obtained anwsers in early layers.
- The findings in this paper have practical implications for deploying LLMs more efficiently in real-world applications by reducing deployment costs and resource demands.
- The experiments are extensive and the efficiency of the method is validated to be consistent on different tasks and models.

weaknesses:
- All experiments in this paper are conducted on a group of datasets, however these datasets are still limited and cannot represent the real-world applications and validate the generalization ability of the pruned model. For example, if the input sequence is not short, early layer attention modules can model the token-wise relationships and predict correct anwsers, but long sequence tasks such as needle in a haystack might be seriously affected by the dropping.
- In addition, the importance scores rely on calibration datasets, and the paper does not extensively explore how variations in these datasets might affect the pruning results. More details should be given, why performances degradation on some datasets are extensive but ignorable on others.

questions:
I wonder the results on more complex tasks and long sequence tasks such as longbench and needle in a haystack.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693353452

thread_type:
official_review

<thread_object>
creator: reviewer_gguv

<content>
summary:
This paper research into the redundancy of attention and MLP layers of transformer-based LLMs. The paper proposes a layer pruning methods with similarity-based layer redundancy measurement. By tracing model checkpoints throughout the training process, it is shown that the layer redundancy is inherent and consistent across training stages. The experiments demonstrate that the pruning method speedup the methods while preserve the performance to some extend.

strengths:
1. The paper proposed an detailed method for Attention and MLP pruning.
2. The paper is well-written and easy to understand.

weaknesses:
1. The comtribution of this paper is limited. There are many previous similar methods in layer pruning [1][2][3], and this paper is a simple extension to the MLP and attention layers. 
2. The author did not analyze the reason for the layer redundancy. Although extensive experiments are provided, the reason why transformer-based LLMs exhibit redundency on the MLP and Attention layers are not explained.
3. The author did not compare the experimental performance with previous block or layer pruning methods [1][2][3]. 


[1] Song J, Oh K, Kim T, et al. SLEB: Streamlining LLMs through Redundancy Verification and Elimination of Transformer Blocks[J]. arXiv preprint arXiv:2402.09025, 2024.

[2] Men X, Xu M, Zhang Q, et al. Shortgpt: Layers in large language models are more redundant than you expect[J]. arXiv preprint arXiv:2403.03853, 2024.

[3] Zhang Y, Li Y, Wang X, et al. FinerCut: Finer-grained Interpretable Layer Pruning for Large Language Models[J]. arXiv preprint arXiv:2405.18218, 2024.

questions:
1. Why a normally trained Transformer model exhibits redundancy in layer levels? Since there are no explicit inductive bias of layers, it is more likely that each layer learns distinct information for the final prediction. Why pruning several layers do not affect the performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702634997

thread_type:
official_review

<thread_object>
creator: reviewer_pgkh

<content>
summary:
This paper studies the redundancy across different modules within Transformers,
including Blocks, MLP, and Attention layers. The study shows that the redundancy mainly comes from the Attention. The authors thus propose an “Attention Drop” method for parameter pruning in LLMs. Experiments on Llama and Mistral are conducted to demonstrate the effectiveness of the proposed method.

strengths:
- The paper is well-organized and easy to follow.
- The study on Attention redundancy is interesting. The proposed approach is reasonable.

weaknesses:
-	The results show that Attention layers and last layers have more redundancy. However, the experiments are mainly conducted on Llama and Mistral. Is this finding valid for other LLMs? It would be interesting to show more LLMs.
-	Please further clarify how the dropped layers are selected in the pruning process. If the target is to drop 4 layers, the attention layers are tested one by one to find 4 layers with lowest importance score?

questions:
Please see above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731735917681

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Reviewers

comment:
Dear Reviewers,  

Thank you for your time and effort in reviewing our paper. We appreciate your feedback and are grateful for the opportunity to refine our work further. **We have noticed some recurring questions and concerns, and we would like to address them collectively here before providing detailed responses to individual reviews**.  

**Due to space limitations, we had to move certain content to the Appendix**, including comparisons with other compression techniques, ablation studies on calibration datasets, and experiments on knowledge-intensive and long-context tasks, among others. **We kindly request that you revisit these sections in the Appendix, as they directly address some of the concerns raised**. It appears that certain claims regarding weaknesses may stem from these sections being overlooked during the review process.  

In addition to the proposed techniques, **we focus on the observation and analysis of the redundancy of attention layers in Section 6, which represents one of our key contributions**. This analysis forms the foundation of our proposed techniques and provides key insights into their effectiveness. **We invite you to revisit this section, as it may help clarify several points raised in your feedback**.  

Once again, we sincerely thank you for your constructive feedback and the time you have invested in evaluating our submission. We hope these clarifications provide a better context for understanding our work. Your insights are invaluable, and we greatly appreciate your assistance in refining our paper further.  

Sincerely,  
Authors of Submission 8274
</content>
</thread_object>

</thread>
<thread>
creation date:
1732640285846

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response on Clarifying the Novelty of Our Paper

comment:
Dear Reviewers,  

Thank you for the time and effort each of you has dedicated to reviewing our submission. We appreciate your insights, but we believe there are some misunderstandings regarding our contributions that merit clarification.  

First, while prior works have identified layer redundancy, **many of them are either concurrent or predate the era of large language models. These works do not diminish our contributions.** Instead, we offer a comprehensive investigation into the redundancy of attention layers, including how to represent it, address it, and explore potential underlying reasons. **Our contribution is both insightful and holistic**, rather than narrow or incremental.  

Second, our work specifically focuses on **architectural redundancy through the lens of layer/block analysis**, utilizing layer drop and block drop as the most straightforward techniques to address this issue. **These techniques offer greater efficiency** compared to finer-granularity pruning approaches. Therefore, **claiming that layer pruning lacks novelty is unwarranted**.  

**Last but not least, while our methods are simple, they are highly effective, and their novelty should not be underestimated.** We have highlighted the unique challenges posed by layer drop as distinct from block drop. Importantly, our approach maintains a complexity of O(L) compared to the O(L^2) or higher complexity of other methods. Experimental results have already validated the effectiveness of our techniques, and **introducing unnecessary complexity on algorithms would detract from their elegance and efficiency**.  

We sincerely appreciate your time and effort in reviewing our work and respectfully request that you reconsider your evaluation. Should you have further questions or require additional clarification, please do not hesitate to reach out.  

Best regards,  
Submission8274 Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733412900112

thread_type:
public_comment

<thread_object>
creator: ~weilin_cai1

<content>
comment:
I find it very interesting that there are many attention modules in the transformer model that can be dropped.
This approach is simple yet effective, offering valuable insights for my exploration of more efficient transformer model structures.
Additionally, the code provided in the Supplementary Material is very user-friendly and indeed reproduces the optimization results consistent with those described in the paper.
</content>
</thread_object>

</thread>


Metareview: