MoH: Multi-Head Attention as Mixture-of-Head Attention

Authors: Peng Jin, Bo Zhu, Li Yuan, Shuicheng YAN
Primary area: foundation or frontier models, including LLMs
Keywords: Multi-Head Attention, Mixture of Experts, Foundation Models

Abstract:
In this work, we upgrade the multi-head attention mechanism, the core of the Transformer model, to improve efficiency while maintaining or surpassing the previous accuracy level. We show that multi-head attention can be expressed in the summation form. Drawing on the insight that not all attention heads hold equal significance, we propose Mixture-of-Head attention (MoH), a new architecture that treats attention heads as experts in the Mixture-of-Experts (MoE) mechanism. MoH has two significant advantages: First, MoH enables each token to select the appropriate attention heads, enhancing inference efficiency without compromising accuracy or increasing the number of parameters. Second, MoH replaces the standard summation in multi-head attention with a weighted summation, introducing flexibility to the attention mechanism and unlocking extra performance potential. Extensive experiments on ViT, DiT, and LLMs demonstrate that MoH outperforms multi-head attention by using only 50\%$\sim$90\% of the attention heads. Moreover, we demonstrate that pre-trained multi-head attention models, such as LLaMA3-8B, can be further continue-tuned into our MoH models. Notably, MoH-LLaMA3-8B achieves an average accuracy of 64.0\% across 14 benchmarks, outperforming LLaMA3-8B by 2.4\% by utilizing only 75\% of the attention heads. We believe the proposed MoH is a promising alternative to multi-head attention and provides a strong foundation for developing advanced and efficient attention-based models.

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

Reviews and discussion:
<thread>
creation date:
1730019915833

thread_type:
official_review

<thread_object>
creator: reviewer_3gij

<content>
summary:
This paper observes that each attention head in multi-head attention operates in parallel. By formulating the multi-head attention in summation form, it builds a dynamic mixture-of-head attention without increasing the number of parameters. Meanwhile, it introduces shared heads and a two-stage routing mechanism to enhance the standard MoH method. Extensive experiments across popular model frameworks, including ViT, DiT, LLMs and continue-tuning LLMs demonstrate strong performance and applicability.

strengths:
1. Without altering the number of parameters, this work treats standard multi-head attention as Mixture-of-Head attention, which enhances the flexibility of the attention mechanism and shows improved performance.
2. To consolidate common knowledge, it introduces shared heads along with a corresponding routing mechanism. The ablation study results have validated the effectiveness of these designs.
3. Popular tasks including ViT for image classification, DiT for image generation, and LLM for language generation, demonstrate superior performance. Furthermore, the proposed MoH attention can continue-tune pre-trained standard multi-head LLaMA3-8B, significantly enhancing its applicability.

weaknesses:
1. This work claims enhanced inference efficiency multiple times. However, there is a lack of experimental evidence to support this claim. Upon reviewing the implementation provided in the Supplementary Material, it appears that you simply mask the useless heads after obtaining the whole multi-head results, which may not genuinely improve inference efficiency. Additionally, the process of dynamically routing each token to the appropriate heads could potentially increase inference costs.
2. The cited works that show multi-head attention contains redundant attention heads primarily focus on natural language processing, it would be better incorporate additional studies from the field of computer vision to provide a more comprehensive perspective.
3. In Equation 5 on line 190, the dimension of $W_r x_t $ is $h-h_s$, so for indices where $h_s+1<i \leq h$, $i$ in $(W_r x_t)_i$ should be 
${i-h_s}$.
4. In Table 5 on line 383, does the baseline LLaMA3-8B refer to the model after continue-tuning with standard multi-head attention, consistent with the configuration of MoH models, or does it represent the starting point for continue-tuning? If the baseline is the starting point, it would be better to add a baseline that reflects the results after continue-tuning with multi-head attention, as continue-tuning is likely to improve performace.

questions:
Please see weakness section
</content>
</thread_object>

</thread>
<thread>
creation date:
1730524946365

thread_type:
official_review

<thread_object>
creator: reviewer_fpak

<content>
summary:
The paper introduces Mixture-of-Head Attention (MoH) to the multi-head attention mechanism in Transformer models, incorporating a routing mechanism that activates the most relevant attention heads for each token. Extensive experiments across diverse model architectures demonstrate that MoH achieves comparable or better performance with fewer attention heads than traditional multi-head attention.

strengths:
1. The idea of applying the mixture-of-experts paradigm to attention heads is novel.
2. MoH shows clear effectiveness for reducing computational overhead by activating fewer attention heads without sacrificing accuracy.
3. The paper presents a wide range of experiments across different model types, demonstrating the effectiveness of MoH. The ability to fine-tune pre-trained multi-head attention models like LLaMA adds practical value to the method.

weaknesses:
1. The contribution is incremental. Replacing the summation of heads with a weighted sum and using expert selection are not entirely new ideas in machine learning, and their application here may not be sufficiently ground-breaking to warrant significant attention without a stronger theoretical basis.
2. The ablation studies are limited in scope and fail to deeply explore the design choices behind MoH. For example, there is little discussion on the impact of different numbers of activated heads beyond the experiments shown. The use of shared heads is also not well-motivated, and the reported improvements may be marginally due to tuning specific hyperparameters.
3. The paper could benefit from evaluations on more diverse and challenging tasks, such as object detection and instance segmentation, in line with prior research on ViT designs.
4. MoH introduces additional complexity with its routing mechanism. The added complexity is not fully justified by the performance gains, especially given that the gains appear marginal in some cases (e.g., DiT models).

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

</thread>
<thread>
creation date:
1730697797532

thread_type:
official_review

<thread_object>
creator: reviewer_z5ku

<content>
summary:
The authors introduce an approach that mitigates redundancy among attention heads through the use of the MoE, which adaptively selects attention heads according to input tokens. This method enhances inference efficiency by employing only those heads that are crucial for feature extraction during the inference process. The MoH demonstrates enhanced performance, even when utilizing a limited number of heads, as evidenced by comprehensive validation experiments.


I appreciate the authors' response. The response provided to my comment appears to be satisfactory, so I keep my score.

strengths:
- The author conducts comprehensive verification experiments to assess the performance of MoH, demonstrating results that are equal to or surpass previous methods. 
- MoH can significantly reduce the head's resources, which can tackle the most important problem of heavy MHSA operations.

weaknesses:
- The author performed ablation studies utilizing different ratios to determine the optimal configuration of shared heads or activated heads; however, this approach is heuristic. The ablation study concerning the ratio of shared heads presented in Table 7 indicates that identifying the optimal head ratio shows significant challenges.
- Given that the primary focus of this paper is to enhance inference efficiency through the reduction of multi-head ratios, it is essential to conduct an experiment that compares this approach with prior methods aimed at decreasing multi-head ratios.
- The author asserts that the shared head acquires common knowledge in Line [180-183], yet this paper is limited to providing evidence that the shared head genuinely learns common knowledge.

questions:
The author claims that the inference efficiency is improved by MoH. However, there is limited ground for this claim in the experiment. Is there any more evidence to support this claim? If further experiments are difficult, even a theoretical interpretation should be presented.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730717563554

thread_type:
official_review

<thread_object>
creator: reviewer_tm7z

<content>
summary:
In the field of deep learning, multi-head attention mechanism has always been a core component of Transformer models, achieving great success in natural language processing and computer vision tasks. However, research has shown that not all attention heads are equally important, and many attention heads can be pruned without affecting model accuracy. Based on this insight, this paper proposes a new architecture called Mixture of Head Attention (MoH) aimed at improving the efficiency of attention mechanisms while maintaining or surpassing previous accuracy levels.

strengths:
MoH can achieve competitive performance while using fewer attention heads.By introducing shared heads and a two-stage routing mechanism, MoH enhances the standard Mixture-of-Experts (MoE) method, enabling the model to capture shared knowledge more effectively across different contexts.

MoH can be fine-tuned from pre-trained multi-head attention models, such as LLaMA3-8B, significantly enhancing the applicability of the model.

The method has been validated across various popular model frameworks, including Vision Transformers (ViT), Diffusion Models (DiT), and Large Language Models (LLMs), demonstrating superior performance in both image classification and language tasks.

weaknesses:
It is suggested to provide more evidence about the diversity within the selected heads. Visualizations and statistics of the distribution may provide more insights.

What is the effect of MoH on multi-task joint learning. More discussions or experiments are welcomed.

What is the `density' in Figure.3. Is it a weight used to select whether to activate？

The discussion section indicate that MoA only suitable for encoder-decoder architecture.  It requires more evidence and explanation.

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

</thread>
<thread>
creation date:
1731944570215

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We sincerely thank all PCs, SACs, ACs, and Reviewers for their time and efforts when handling our paper.

All reviewers appreciate the contributions of our method:
* All reviewers pointed out that **MoH achieves competitive or superior performance while utilizing fewer attention heads**.
* All reviewers commented that we conduct **a wide range of experiments across different model types**, including Vision Transformers (**ViT**), Diffusion Models (**DiT**), and Large Language Models (**LLMs**), demonstrating **superior performance in both image classification and language tasks**.
* Reviewers Tm7Z, fpAk, and 3giJ all highlighted that **the proposed MoH can continue-tune pre-trained standard multi-head LLaMA3-8B, significantly enhancing its applicability**.

As suggested by the reviewers, we have revised the manuscript as follows:
* Additional visualizations of the head load distribution for MoH-ViT (Fig. A left on Page 20), MoH-DiT (Fig. A right on Page 20), and MoH-LLaMA3 (Fig. B on Page 21).
* Additional visualizations of the head routing score distribution for MoH-ViT (Fig. C on Page 22), MoH-DiT (Fig. D on Page 23), and MoH-LLM (Fig. E on Page 24).
* The ablation study on the impact of the activated head ratio (Tab.F on Page 20).
* Comparison between MoH-LLaMA3-8B and LLaMA3-8B-stage1 (Tab.E on Page 19).
* Discussion of related works about redundant attention heads in computer vision in the Introduction.
* Modifying the inappropriate indices in Equation 5 and Equation 7.
* Modifying the inappropriate description of MoA.
* Modifying the caption of Table 5 and Figure 3.

We highlight all the modifications using red color.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732551566957

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
The Discussion Period Ends Soon

comment:
Dear Reviewers,

We sincerely appreciate the invaluable feedback provided by reviewers. Since the discussion period ends soon (November 26, AoE), we would appreciate hearing whether our rebuttal has addressed your concerns. We would be happy to provide more input if needed.

Thanks for your attention!

Best regards,

Authors of Paper#2876
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper investigates neural architecture design and proposes a new architecture called a Mixture of Head Attention aimed at improving the efficiency of attention mechanisms while maintaining accuracy. The manuscript was reviewed by four experts in the field. The recommendations are ("5: marginally below the acceptance threshold", 3 x "6: marginally above the acceptance threshold"). The reviewers raised several concerns regarding the paper, e.g., incremental technical contribution, unconvincing experimental evaluations, etc. Considering the reviewers' concerns, we regret that the paper cannot be recommended for acceptance at this time. The authors are encouraged to consider the reviewers' comments when revising the paper for submission elsewhere.