DHA: Learning Decoupled-Head Attention from Transformer Checkpoints via Adaptive Heads Fusion

Summary

The paper proposes a novel method to convert the Transformer attention mechanism from multi-head attention (MHA) to the proposed decoupled-head attention (DHA). DHA employs fewer unique KV heads than MHA or GQA by representing the heads as a linear combination of unique heads within the group, where each group contains similar heads. This conversion process is performed at the parameter level, so basically any MHA model can be converted to a DHA model. After conversion, continued pre-training helps mitigate performance drops. The core novelty of the proposed conversion mechanism comes from its head fusion strategy: first, it searches for similar head groups, optimizes the fusion weights to obtain many-to-one mapping, and then allocates a different number of heads to different layers within the resource budget. Experiments show that DHA requires considerably fewer training resources while achieving lower LM loss compared to GQA. Furthermore, DHA is beneficial for reducing the KV cache size.

Strengths

* The paper identifies significant redundancies in the attention KV heads and proposes a practical solution to address the problem. Empirical evidence for the motivation is also provided. * The learning-based fusion optimization and adaptive budget allocation offer a systematic approach to finding optimal mappings, adding non-trivial value to the proposed method. * Related works are sufficiently addressed.

Weaknesses

* It is unclear whether the paper uses inter-layer grouping of heads or only intra-layer grouping. The introduction and Figure 1 suggest that KV heads from different layers can also be merged, but this does not seem to be reflected in the method and the models used in the experiments. * If the grouping is only performed in an intra-layer manner, the paper should be revised accordingly for clarity. * If the grouping is performed in an inter-layer manner, the optimization process would become much more complex. Also, during inference, each layer would need to call a large number of parameters (essentially all weights of the corresponding group across layers), which would slow down the inference due to the memory-bound nature of LLMs. * If the grouping is performed in an inter-layer manner, a visualization similar to Figure 2 across layers could strengthen the claim. * It is assumed that "GQA" in all experiments refers “MHA -> GQA” conversion using simple mean pooling. Is this a correct understanding? Is it a common approach to convert MHA to GQA instead of training a GQA model from scratch? Can we consider the uniform division of heads and apply mean pooling for GQA construction as a solid and strong baseline? For example, using CKA to identify similar heads before mean pooling could improve GQA’s performance. * It is difficult to say that DHA shows "much less training overhead" because it starts from MHA (line 294), especially compared to the models that begin with GQA from scratch.

Questions

* Although the paper introduces the method as performing MHA to DHA conversion at the parameter level (checkpoint-level), the Lagrangian min-max training requires LM loss computation for training. This means the fusion step in the conversion process requires task-oriented training using real data. Is this correct? * How long does it take for the searching process? The paper mentions 100 search steps (line 197) for each layer. * Providing a demonstration of how merging weights change after each process step would help in understanding the method better. * Figure 2 visualizes the weight similarity instead of CKA, which makes sense because similar weights would generate similar outputs, and the method merges similar parameters. Then, is CKA-based analysis necessary (also there is no CKA visualization)? * (minor) Typo: Adaptvie --> Adaptive in the caption of Table 2.

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The paper adequately addresses the limitations.

Summary

This paper introduces a novel mechanism to optimize large language models (LLMs) by addressing the computational and memory costs associated with the Multi-Head Attention (MHA) mechanism. The authors propose Decoupled-Head Attention (DHA), which adaptively configures the sharing of key and value heads across layers, aiming to strike a balance between performance and efficiency. The transformation from MHA to DHA is achieved through a process of head parameter clustering and linear fusion, which retains the knowledge from the original model with minimal performance degradation.

Strengths

-The paper demonstrates a remarkable balance between performance and efficiency, showing that DHA can achieve high performance with significantly reduced computational resources. -The proposed method is not limited to specific model size, indicating that it could be broadly applied to various existing MHA Transformer models.

Weaknesses

The effectiveness of DHA on a wider range of model architecture remains unexplored. Could the method work well on models beyond LLAMA? The process of transforming MHA to DHA involves several stages of training and fusion that might be complex to implement and reproduce for researchers without access to similar computational resources. The paper relies solely on linear methods for parameter fusion. There might be room for improvement by exploring non-linear fusion techniques that could potentially offer better optimization.

Questions

See the weaknesses.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes

Summary

The paper introduces Decoupled-Head Attention (DHA), a new efficient attention mechanism for large language models. DHA adaptively configures group sharing for key and value heads across layers, transforming Multi-Head Attention checkpoints through a three-stage process. It achieves 97.6% of the original performance while using only 0.25% of the pre-training budget and reducing KV cache by 75%. DHA outperforms other efficient attention methods in terms of training speed and performance under limited resources.

Strengths

1. Novel approach: Introduces Decoupled-Head Attention (DHA), an innovative method to optimize attention mechanisms in large language models. 2. Efficiency gains: Achieves significant reductions in computational and memory costs (75% KV cache reduction) with small performance loss. 3. Rapid adaptation: Requires only 0.25% of the original pre-training budget to achieve 97.6% of performance.

Weaknesses

1. While the proposed method seems to be effective for MHA, the authors did not demonstrate its compatibility with GQA, which is increasingly popular recently and starts to replace MHA in almost all state-of-the-art models. 2. While efficient adaptation is important, it is more important to see whether the performance of DHA is similar to GQA/MQA in the end of the training process, when the entire model converges. If DHA just converges faster, I think its value is rather limited. 3. While the accuracy loss is not as large as other KV compression methods, a ~5% performance drop for Llama2-7B is still too large. It's almost equivalent to the performance gap between Llama-2-7B and Sheared-Llama-2.7B. This indicates that the process is extremely lossy, and the resulting model might not be useful after DHA process. 4. It is unclear how DHA combines with existing KV cache compression techniques, such as KV cache quantization.

Questions

Please answer my questions listed in "Weaknesses"

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

The paper introduces a novel attention mechanism, Decoupled-Head Attention (DHA), designed to enhance the efficiency of large language models (LLMs) with minimal performance loss. DHA adaptively configures key and value heads across layers by leveraging insights from attention redundancy, leading to significant savings in computational and memory costs. The authors propose a method to transform a standard multi-head attention (MHA) model into DHA with a minimal pre-training budget, achieving near-original performance while reducing the key-value (KV) cache requirements substantially. The experiments demonstrate the effectiveness of DHA in maintaining performance with reduced resources.

Strengths

The paper proposes a novel Decoupled-Head Attention (DHA) mechanism that addresses the computational and memory costs associated with Multi-Head Attention (MHA) in large language models (LLMs). The idea of adaptively configuring group sharing for key and value heads is innovative. DHA initialization can recover the performance with a very small amount of restorative pre-training, compared to Group-Query Attention.

Weaknesses

1. **Comparison with GQA and MHA Training Data Requirements:** It would be beneficial for the authors to report the performance of the DHA model when continue pretrain with a larger dataset, compared to GQA and MHA models. 2. There is a noticeable accuracy gap between DHA and MHA models on academic datasets.

Questions

1. Could the authors clarify why an increase in the fusion phase token budget leads to a decrease in average accuracy post-CT as observed in Table 3? 2. Are there any experimental results available that demonstrate how the DHA model performs on larger datasets?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The author has mentioned limitations in Appendix

Please start reviewer-author open discussion

During August 7-13 there will be an open discussion between the reviewers and authors. I observed that the author responded to reviewer questions elaboratively. Reviewer FLzu, wgtQ, Txec, aEta --- Please use this opportunity to further clarify remaining concerns and confusions. Thank you! AC

Response to Reviewer aEta’s Feedback

Thank you very much for reading our response and for all your efforts in the review process! Your valuable comments have greatly helped us improve the paper. If you have any other concerns, we are happy to answer them.