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

We sincerely appreciate your thoughtful comments and your recognition of our method: "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." Below, we provide detailed responses to your questions.

Q1: This work claims enhanced inference efficiency multiple times. However, there is a lack of experimental evidence to support this claim. Additionally, the process of dynamically routing each token to the appropriate heads could potentially increase inference costs.

A1: Thanks for your insightful comments. The code provided in the supplementary material does not genuinely improve inference efficiency.

Implementing dynamic routing directly at the CUDA level would significantly enhance the efficiency of MoH. However, the support for MoE structures in the open-source community has remained limited, and developing CUDA operators independently is challenging for our small research team.

Although we could not find an efficient open-source MoE routing framework, sparse matrix multiplication can serve as a temporary solution. Specifically, we transform the Q, K, and V features into sparse matrices using the mask generated by the router and replace the dense matrix multiplication in the attention mechanism with sparse matrix multiplication.

As shown in the table below, although dynamic routing introduces additional computational overhead, MoH still outperforms standard multi-head attention mechanisms. Furthermore, as the input sequence gets longer, the advantage of MoH grows.

It is worth noting that, to eliminate the impact of underlying operator optimizations, we replaced all matrix multiplications with sparse matrix multiplication when testing for speed.

Using sparse matrix multiplication to implement MoH is only a temporary solution. The support for sparse matrix multiplication in CUDA operators is still very limited, as current GPUs are heavily optimized for dense matrix computations.

We believe that as the open-source community continues to focus on MoE models, efficient MoE operators will be developed, enabling MoH to achieve even faster performance.

Q2: 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.

A2: Thanks for your insightful advice. Following your suggestion, we have expanded the discussion of related work in computer vision within the Introduction: "In computer vision, some works also identify attention head redundancy. Bhattacharyya et al. (2023) [1] reduces redundancy to boost performance, while Yun & Ro (2024) [2] develops single-head attention for efficiency."

[1] Bhattacharyya, Mayukh, Soumitri Chattopadhyay, and Sayan Nag. "DeCAtt: Efficient Vision Transformers with Decorrelated Attention Heads." in CVPRW. 2023.

[2] Yun, Seokju, and Youngmin Ro. "SHViT: Single-head vision transformer with memory efficient macro design." in CVPR. 2024.

Q3: 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$.

A3: Thank you for kindly pointing out our inappropriate indices $i$ in Equation 5. Following your suggestion, we have revised Equation 5 in the revised manuscript.

Besides, when we carefully proofread the manuscript, we found that the indices in Equation 7 had the same issue. We have also modified Equation 7 in the revised manuscript.

Q4: 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.

A4: The baseline LLaMA3-8B refer to the starting point for continue-tuning.

Due to the page limit, we have included the performance of the model after continued training in Table E (Page 19) in the Appendix. Additionally, we have updated the caption of Table 5 to include the note, "Please refer to Tab. E in the Appendix for the performance of the model at the end of the first stage of training."

As shown in the table below, MoH-LLaMA3-8B quickly recovers the performance of LLaMA3-8B-stage1 within a training budget of 100B tokens. Notably, in English language tasks, MoH-LLaMA3-8B surpasses LLaMA3-8B-stage1 while using only 75% of the attention heads.

However, for Chinese language and math tasks, the recovery performance of the MoH model is not as strong as for English. For example, MoH-LLaMA3-8B achieves an accuracy of 64.4% on CMMLU, compared to 66.0% for LLaMA3-8B-stage1. We attribute this to the fact that the model's Chinese and mathematical capabilities are primarily established during the first training stage. Since the first training stage uses only 300B tokens, significantly less than the 15T tokens in LLaMA3-8B's pre-training, the model's abilities in these areas are not fully stable. In the second training stage, after switching to the MoH model, the model experiences more significant forgetting in Chinese and math tasks.

Overall, as shown in the table below, MoH-LLaMA3-8B achieves an average accuracy of 64.8% across 14 benchmarks, outperforming LLaMA3-8B-stage1 by utilizing only 75% of the attention heads.

It is worth noting that continuing to train the MoH-LLaMA3-8B can further improve its performance. However, due to GPU resource constraints, we were limited to a training budget of 100B tokens.

We sincerely appreciate your thoughtful comments and for acknowledging that our method is "is novel," "shows clear effectiveness for reducing computational overhead by activating fewer attention heads without sacrificing accuracy," and validated through "a wide range of experiments across different model types." Below, we provide detailed responses to your questions.

Q1: 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.

A1: Thanks for your insightful comments. We demonstrate that MoH is superior to vanilla multi-head attention from both theoretical and experimental perspectives.

Specifically, MoH not only improves efficiency and model performance but also helps different attention heads to specialize better compared to multi-head attention.

From the theoretical perspective, in standard multi-head attention, all heads use the same data, which can cause them to learn similar features. Many studies [1,2,3,4] have pointed out that there are redundant heads in multi-head attention.

Given a minibatch of data $D$, the gradient of each attention head in multi-head attention can be written as $E_{x \in D} [ \frac{\partial \mathcal{L(x)}}{\partial h_{i} } ]$.

In contrast, in MoH, routed heads are trained only on smaller subsets of data specifically assigned to them. In MoH's routing mechanism, the data is divided into $h-h_s$ subsets $D_1,D_2,...,D_{h-h_s}$, with each subset corresponding to a routed head. Besides, the routing score for each attention head acts as an adaptive adjustment to the learning rate, enabling the attention heads in MoH to specialize more effectively.

Given a minibatch of data $D$ and the router $G(*)$, the gradient of each routed head in MoH can be written as $E_{x \in D_i } [G(x) \frac{\partial \mathcal{L(x)}}{\partial h_{i} }]$. The gradient of each shared head in MoH can be written as $E_{x \in D } [G(x) \frac{\partial \mathcal{L(x)}}{\partial h_{i} }]$.

As shown in the table below, the routing mechanism and adaptive weights in MoH enable attention heads to specialize more effectively compared to standard multi-head attention.

From the experimental perspective, we calculated the similarity of attention patterns and output features of different attention heads (include routed heads and shared heads).

As shown in the table below, the similarity of attention patterns and output features among attention heads in MoH is lower than in standard multi-head attention, indicating reduced redundancy and greater differentiation among the attention heads in MoH.

Given a pair of attention score matrices A and A', we calculate the similarity of attention patterns as $1 - \frac{1}{2} E_{query}[ ||A-A'||_{1} ]$. Since attention scores form a probability distribution for each query, the similarity is always between 0 to 1.

[1] Voita, Elena, et al. "Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned." in ACL. 2019.

[2] Yun, Seokju, and Youngmin Ro. "SHViT: Single-head vision transformer with memory efficient macro design." in CVPR. 2024.

[3] Michel, Paul, Omer Levy, and Graham Neubig. "Are sixteen heads really better than one?." in NeurIPS. 2019.

[4] Bhattacharyya, Mayukh, Soumitri Chattopadhyay, and Sayan Nag. "DeCAtt: Efficient Vision Transformers with Decorrelated Attention Heads." in CVPRW. 2023.

Q2: There is little discussion on the impact of different numbers of activated heads.

A2: Thanks for your helpful advice. As suggested, we conducted ablation experiments and updated the results in the revised manuscript (Tab.F on Page 20).

The table below shows that activating more attention heads generally leads to improved model performance. These results are intuitive, as activating more attention heads equates to utilizing more parameters and performing additional computations on the input.

Q3: The use of shared heads is not well-motivated.

A3: Thanks for your thoughtful comments. We introduce shared heads to efficiently learn common knowledge.

In attention mechanism, some attention heads may capture common knowledge across different contexts, such as grammatical rules in language. Therefore, we designate a subset of heads as shared heads that remain always activated. By consolidating common knowledge within shared heads, we reduce redundancy among the other dynamically routed heads.

We analyzed the feature rank of shared heads and routed heads. A lower feature rank indicates a higher correlation between features of different samples, suggesting that the features capture more general knowledge.

As shown in the table below, the feature rank of shared heads is significantly lower than that of routed heads. This suggests that shared heads primarily capture common information across samples, while routed heads focus on information unique to individual samples.

In summary, the combination of shared heads and routed heads allows MoH to effectively handle both general knowledge and class-specific knowledge.

Q4: The reported improvements may be marginally due to tuning specific hyperparameters.

A4: To ensure a fair comparison, we only replace the standard multi-head attention with the MoH, while keeping all other training parameters identical to ViT, DiT, and LLM. In other words, the training parameters were fully optimized for multi-head attention, and no additional parameter tuning was performed for the MoH model. Despite no additional parameter tuning, MoH outperforms multi-head attention by using only 50%$\sim$90% of the attention heads.

Q5: 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.

A5: Thanks for your insightful advice. Following your suggestion, we conducted experiments on object detection and instance segmentation tasks. Specifically, we used a Mask R-CNN detection head, trained with a 1× schedule, to evaluate the performance of ImageNet-1K pretrained MoH-ViT-S on the COCO dataset.

As shown in the table below, MoH-ViT-S shows competitive performance in both object detection and instance segmentation.

Q6: MoH introduces additional complexity with its routing mechanism.

A6: Thanks for your thoughtful comments. In our experiment settings, the additional complexity of routing mechanism is small, especially for long input sequences.

In MoH, we transform the Q, K, and V features into sparse matrices using the mask generated by the router and replace the dense matrix multiplication in the attention mechanism with sparse matrix multiplication.

As shown in the table below, although dynamic routing introduces additional computational overhead, MoH still outperforms standard multi-head attention mechanisms. Moreover, as the input sequence length increases, the extra overhead impact from routing mechanism decreases.

To eliminate the impact of underlying operator optimizations, we replaced all matrix multiplications with sparse matrix multiplication when testing for speed.

Using sparse matrix multiplication to implement MoH is only a temporary solution. Implementing dynamic routing directly at the CUDA level would make MoH significantly faster. However, for our small research team, independently developing CUDA operators is challenging.

We believe that as the open-source community continues to focus on MoE models, efficient MoE operators will be developed, enabling MoH to achieve even faster performance.

Q7: 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).

A7: Thanks for your insightful comments. In fact, as shown in the table below, our MoH-DiT-XL/2 achieves significant improvements over DiT.

Dear Reviewer

Could we kindly inquire if the responses have satisfactorily tackled your concerns, or if there is a need for further clarification? Your commitment to reviewing our work is immensely appreciated, and we express our sincere gratitude for your insightful comments and the considerable time you have dedicated to reviewing our paper.

Dear Reviewer

Thank you once again for your thorough review. We hope that our responses and the revised manuscript have effectively addressed your concerns. We would greatly appreciate your feedback. Please feel free to let us know if you have any further questions.

We sincerely appreciate your thoughtful comments, especially noting that we "conduct comprehensive verification experiments" and our method "can tackle the most important problem of heavy MHSA operations." Below, we provide detailed responses to your questions.

Q1: Table 7 indicates that identifying the optimal head ratio shows significant challenges.

A1: Thanks for your insightful comments. Identifying the optimal head ratio remains a challenge that requires further exploration. However, determining the optimal hyper-parameters of a model is a common challenge in deep learning.

For example, even in multi-head self-attention mechanisms, finding the optimal number of heads and their dimensions is not a fully solved problem. As an illustration, increasing the number of heads in a ViT model from 24 to 32 (while reducing the corresponding head dimensions from 32 to 24) improves its accuracy by 0.1 on the ImageNet-1K classification.

We hope future research will provide more insights into determining optimal hyper-parameters for our method.

Q2: A comparison of this approach with prior methods aimed at decreasing multi-head ratios.

A2: Thanks for your valuable advice. As suggested, we trained a new MoH model with a similar number of parameters as SHViT [1] and compared its performance to SHViT. Unlike our method, SHViT uses convolution and single-head attention to address head redundancy in ViT.

As shown in the table below, our MoH model significantly outperforms SHViT in overall performance.

[1] Yun, Seokju, and Youngmin Ro. "SHViT: Single-head vision transformer with memory efficient macro design." in CVPR. 2024.

Q3: Evidence that the shared head genuinely learns common knowledge.

A3: Thanks for your helpful comments. We provide both theoretical and experimental evidence that shared heads learn common knowledge.

• From the theoretical perspective, routed heads are trained only on smaller subsets of data specifically assigned to them, while shared heads are trained on the entire dataset. In MoH's routing mechanism, the dataset is divided into $h-h_s$ subsets ${D_1,D_2,...,D_{h-h_s}}$, with each subset corresponding to a routed head. In contrast, shared heads are trained on the full dataset, $D=D_1 \cup D_2 \cup ... \cup D_{h-h_s}$. Thus, shared heads capture the common knowledge across the entire dataset, while routed heads focus on the specific knowledge of their respective subsets.

• From the experimental perspective, we analyzed the feature rank of shared heads and routed heads. A lower feature rank indicates a higher correlation between features of different samples, suggesting that the features capture more general knowledge. As shown in the table below, the feature rank of shared heads is significantly lower than that of routed heads. This suggests that shared heads primarily capture common information across samples, while routed heads focus on information unique to individual samples.

In summary, whether from theoretical intuition or experimental results, shared heads are more tented to capture general knowledge compared to routed heads.

Q4: Is there any more evidence to support the inference efficiency is improved by MoH? If further experiments are difficult, even a theoretical interpretation should be presented.

A4: Thanks for your insightful advice. Since MoE models gained significant attention following the rise of LLMs, the support for MoE structures in the open-source community has remained limited. While we could not find an efficient open-source MoE routing framework, sparse matrix multiplication can serve as a temporary solution.

Specifically, we transform the Q, K, and V features into sparse matrices using the mask generated by the router and replace the dense matrix multiplication in the attention mechanism with sparse matrix multiplication.

As shown in the table below, although dynamic routing introduces additional computational overhead, MoH still outperforms standard multi-head attention mechanisms. Furthermore, as the input sequence gets longer, the advantage of MoH grows.

It is worth noting that, to eliminate the impact of underlying operator optimizations, we replaced all matrix multiplications with sparse matrix multiplication when testing for speed.

Using sparse matrix multiplication to implement MoH is only a temporary solution. Implementing dynamic routing directly at the CUDA level would make MoH significantly faster. However, for our small research team, independently developing CUDA operators is challenging.

We believe that as the open-source community continues to focus on MoE models, efficient MoE operators will be developed, enabling MoH to achieve even faster performance.

Dear Reviewer

Would it be possible for us to kindly ascertain if the provided responses have satisfactorily tackled any concerns you may have had and if further explanations or clarifications are needed? Your generous investment of time and effort in the evaluation of our work is truly commendable. We extend our heartfelt gratitude for your insightful commentary and the considerable time you have devoted to reviewing our paper.

Dear Reviewer

We truly appreciate your time and effort in helping us improve our work. We hope our responses have addressed your concerns. If anything is still unclear or needs more explanation, we are happy to provide further details. Since the discussion period ends soon, your comments are especially important to us.

Thanks again for taking the time and effort when handling our paper.

We sincerely appreciate your thoughtful comments and for recognizing that our method "achieves competitive performance with fewer attention heads," "can be fine-tuned from pre-trained multi-head attention models," and "has been validated across various popular model frameworks." Below, we address your questions in detail.

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

A1: Thanks for your insightful advice. As suggested, we have included more visualizations in the revision. Specifically, in the appendix of the revised manuscript, we have provided:

• 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).

On the one hand, as shown in the additional visualizations of the head load distribution, there is notable variation in attention head assignments across different categories and task topics. This suggests that the MoH model adapts to a wide range of tasks by utilizing distinct head assignment patterns. This ability enables MoH to allocate attention heads more effectively to specific task types, leading to more efficient parameter utilization than standard multi-head attention.

On the other hand, MoH replaces the standard summation in multi-head attention with a weighted summation. As illustrated in the additional visualizations of the head routing score distribution, these head routing scores also vary across categories and task types. This dynamic weighting mechanism allows MoH to adjust the importance of each head in response to different task requirements, further enhancing its flexibility and performance.

Besides, we also calculated the Coefficient of Variation (CV) of the head routing scores across categories and task topics. As shown in the table below, the routing scores of shared heads change more across categories than those of routing headers. We consider this because routed heads adapt to different categories by adjusting their activation, while shared heads remain activated all the time. Therefore, shared heads primarily rely on changes in routing scores to adapt to different categories.

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

A2: Thanks for your valuable comments. We conducted two experiments:

• Training large language models is naturally a form of multi-task learning because the training data comes from various sources, including CommonCrawl, C4, Github, Wikipedia, Books, ArXiv, StackExchange, and so on. As shown in the table below, MoH achieves higher average accuracy across multiple tasks compared to the multi-head attention model.

• We combined several task-specific datasets (ShareGPT, WizardLM_evol_instruct, and SlimOrca for instruction following, MetaMath for math, and Evol-CodeAlpaca for code). Then, we fine-tuned LLaMA3-8B using this multi-task dataset. As shown in the table below, MoH also achieves higher average accuracy across multiple tasks compared to the multi-head attention model.

Q3: What is the "density" in Figure 3. Is it a weight used to select whether to activate?

A3: In Figure 3, "density" denotes the ratio of head activations to the total number of tokens. For example, if a head is activated 128 times out of 256 tokens, the density would be 0.5.

To make Figure 3 easier to understand, we have added an explanation of "density" in the caption of Figure 3: "'density' denotes the ratio of the number of head activations to the total number of tokens."

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

A4: Thank you for kindly pointing out our inappropriate description. A more precise statement would be that the experiments of MoA are conducted only on the encoder-decoder architecture for language tasks.

In the revised manuscript, we have revised the original sentence "MoA is only validated on the encoder-decoder architecture for language tasks" to "MoA is only validated for language tasks."

Dear Reviewer

May we kindly inquire if the provided responses have adequately addressed any questions you might have had? Does there remain a requirement for further explanations or clarifications? We wish to express our sincere gratitude for your meticulous evaluation and for generously investing a significant amount of your time in reviewing our paper. Your feedback would be greatly valued.