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

We sincerely appreciate your thoughtful comments and for recognizing our work as "intuitive and effective," acknowledging that "MoH can achieve competitive performance," and highlighting our "detailed ablation study and well-organized paper structure." Below, we address your questions in detail.

Q1: The idea is not entirely new. It shares a similar idea with MoA.

A1: We explain this problem in three aspects:

• In terms of motivation, MoH aims to make the attention mechanism more efficient and effective without adding extra parameters. In contrast, MoA, like MoE, focuses on increasing model size while keeping inference costs low. Therefore, the model settings of MoH are more stringent than those of MoA.
• In terms of methodology, we maintain the original structure of multi-head attention as much as possible, allowing MoH to seamlessly replace standard multi-head attention across different tasks without extra tuning. In contrast, MoA introduces shared keys and values to add heads without increasing the KV cache, which disrupts the original multi-head attention design.
• In terms of flexibility, we show that pre-trained multi-head attention models can be further continue-tuned into our MoH models, making MoH highly practical. In contrast, MoA integrates multi-head attention with MoE but relies on shared keys and values, requiring training from scratch, which reduces its flexibility.

Q2: No experiments compared with MoA.

A2: Thanks for your insightful advice. As suggested, we compare MoA and MoH on the translation task and additionally provide results for image classification. As shown in the table below, MoH outperforms MoA, primarily for two reasons: (1) MoA's use of shared keys and values across heads limits its expressiveness; (2) MoH's shared heads and two-stage routing improve the model's ability to capture general knowledge (please refer to our response (A1) to Reviewer 6acz).

Q3: How to find a good balance between specialization and generalization.

A3: There are two ways to choose shared heads and routed heads: (1) manual configuration and (2) learning through learnable masks. In the manuscript, we show the results of manual configuration. For the learnable mask approach, we follow a method similar to [1]. Specifically, we introduce a mask module with binary values {0,1} applied to the heads. These masks are dynamically learned rather than statically assigned, allowing the model to determine which heads should be shared. The remaining heads are designated as routed heads, and we set K based on a predefined ratio. For example, if the ratio is 1/4 and there are 8 routed heads, then K is set to 2. As shown in the table below, our latest experimental results show that learning through learnable masks is generally better than manual configuration.

[1] Liu, Peiyu, et al. "MOEfication by Experts as Masks."

Q4: Table 5 presents the Ablation study on the impact of each component of the proposed MoH, is the first row equal to MoA?

A4: The first row in Table 5 follows a structure similar to MoA, but without shared keys and values or the use of additional z-loss.

Q5: Providing state-of-the-art mechanism explanation presented in 2024.

A5: Thanks for your valuable suggestion. We have added your additional references to the Introduction to better explain that not all attention heads hold equal significance.

Q6: TYPO: Sec 4.3 "Please refer to the Appendix for detailed hyper-parameter settings (Tab. C)" link to wrong table of Table 3.

A6: Thank you for your thorough review. This issue may be a bug in LaTeX, and we will work on fixing it.

Q7: Does the implementation of MoH affect GPU memory usage? if so, it increases of decreases the total memory usage? In that case, can we support larger model or small model by using MoH?

A7: As shown in our response (A2) to Reviewer 6acz, MoH slightly reduces GPU memory usage, though the difference is not significant. This is because GPU memory is primarily used to store model parameters, gradients, and the KV cache. Since MoH only optimizes attention computation, it does not substantially reduce the GPU memory of these three components. Although MoH doesn't allow training larger models, it can make training and inference faster.

We sincerely appreciate your thoughtful comments and for recognizing that "the idea of adding a MoE-like routing in the attention mechanism is fascinating and worth investigating, especially for efficiency purposes." Below, we address your questions in detail.

Q1: In most cases, performance enhancements provided by MoH are marginal.

A1: We explain this problem in two aspects:

• The motivation for our work is to upgrade the multi-head attention mechanism to reduce computational costs while maintaining or surpassing the previous accuracy level. Therefore, our method activates fewer parameters than multi-head attention, and our method performs better if it activates the same level of parameters.
• To demonstrate the robustness of our method, we only replace multi-head attention with MoH in various structures while keeping the original training parameters unchanged. Our latest experimental results indicate that, with tuning, our method can achieve even higher performance.

Q2: The accessibility of the method could be improved by rephrasing the paragraph on two-step routing.

A2: Thanks for your insightful advice. As suggested, we have rewritten the paragraph on two-step routing to make it easier for readers to understand.

Q3: The difference between the setting of ViT and DiT w.r.t. Llama.

A3: We explain our experimental setup in detail:

• ViT for Image Classification. We trained a MoH model from scratch based on TransNeXt. To ensure a fair comparison, we only replace the standard multi-head attention with the proposed MoH, while keeping all other training parameters identical to TransNeXt.
• DiT for Class-Conditional Image Generation. We trained a MoH model from scratch based on DiT by replacing the standard multi-head attention with the proposed MoH. We also keep all other training parameters identical to DiT.
• Training LLMs from Scratch. We trained the LLMs from scratch, maintaining the multi-head attention baseline with exactly the same training parameters as the MoH model.
• Continue-Tuning LLaMA3-8B. To significantly enhance the applicability of the proposed MoH method, we attempt to further continue-tune pre-trained multi-head attention models, such as LLaMA3-8B, into MoH models.

Q4: What strategy is used to choose shared heads?

A4: All MoH models contain shared heads. There are two ways to choose shared heads: (1) manual configuration and (2) learning through learnable masks. In the manuscript, we show the results of manual configuration. For the learnable mask approach, we follow a method similar to [1]. Specifically, we introduce a mask module with binary values {0,1} applied to the heads. These masks are dynamically learned rather than statically assigned, allowing the model to determine which heads should be shared. The remaining heads are designated as routed heads, and we set K based on a predefined ratio. For example, if the ratio is 1/4 and there are 8 routed heads, then K is set to 2. As shown in the table below, our latest experimental results show that learning through learnable masks is generally better than manual configuration.

[1] Liu, Peiyu, et al. "MOEfication by Experts as Masks."

Q5: Essential references not discussed.

A5: Thanks for your advice. We have expanded [1,2,3] and improved the discussion of related works.

Q6: The dynamic routing works at the single token level.

A6: We input the sentence "Give me a short introduction to large language model." into MoH-LLaMA3-8B. We find the tokens forming the phrase share most of the activated heads. For example, all three tokens in "large language model" activate heads {6,7,8,9,11,12}. For other tokens, no clear pattern emerged in the activation of routed heads. Notably, besides routed heads, the shared heads create a stable semantic interaction between all tokens.

We sincerely appreciate your thoughtful comments and for recognizing that our method is a "practical contribution," "can be fine-tuned from pre-trained multi-head attention models," and "the results are promising and suggest broad applicability." Below, we address your questions in detail.

Q1: There are inconsistencies between the theoretical formulation of MoH and its implementation in the LLaMA3-8B experiment.

A1: We explain this problem in three aspects:

• Experiments on ViT, DiT, and LLMs are conducted with models trained from scratch, giving us more flexibility to modify the architecture. However, LLaMA3-8B is a pre-trained model trained on 15T tokens, while we had only 400B tokens (about 3% of the pre-training data) for continue-tuning. This limited our ability to alter the model structure significantly. To address this, we adjusted our original formula to maximize weight reuse from the pre-trained model while preserving its output distribution.

• We have conducted experiments on LLMs trained from scratch, demonstrating the advantages of our method. Since much of the current research focuses on fine-tuning open-source LLMs, we introduced additional experiments to continue-train LLaMA3-8B. This contribution has been acknowledged by all other reviewers.

• To the best of our knowledge, our work is the first to attempt to reduce the computational amount of the attention mechanism without degrading the model performance by continue-tuning on a pre-trained model. We consider our proposed training techniques to be a valuable additional contribution.

Q2: The paper does not provide a clear rationale for how the value of K in the Top-K selection is determined.

A2: Thanks for your insightful advice. There are two ways to set this up: (1) manual configuration and (2) learning through learnable masks. In the manuscript, we show the results of manual configuration. For the learnable mask approach, we follow a method similar to [1]. Specifically, we introduce a mask module with binary values {0,1} applied to the heads. These masks are dynamically learned rather than statically assigned, allowing the model to determine which heads should be shared. The remaining heads are designated as routed heads, and we set K based on a predefined ratio. For example, if the ratio is 1/4 and there are 8 routed heads, then K is set to 2. As shown in the table below, our latest experimental results show that learning through learnable masks is generally better than manual configuration.

[1] Liu, Peiyu, et al. "MOEfication by Experts as Masks."

Q3: The relationship between the K value and the ratio of shared heads is not discussed.

A3: Thanks for your valuable suggestion. It is worth noting that when using learnable masks to determine the number of shared heads, we only need to set the ratio of K among the routed heads. As shown in the table below, this ratio is a trade off parameter. If this ratio is too small, the number of activated heads will be insufficient, potentially degrading performance. Conversely, if this ratio is too large, it reduces the model’s sparsity, limiting efficiency improvements.

We sincerely thank you for your constructive comments. We will add the above important discussions in the final manuscript and highlight them. Thanks again for taking the time and effort on our paper.

We sincerely appreciate your thoughtful comments and your recognition that our method "provides a promising suggestion to advance attention-based models." Below, we provide detailed responses to your questions.

Q1: While prior work supports this claim when heads are combined through concatenation, the proposed method instead uses summation.

A1: First, in the forward pass, the summation and concatenation forms give the same results, i.e., $\textrm{MultiHead}({X},{X}')=\sum_{i=1}^{h} {H}^{i}{W}_O^{i}=\textrm{Concat}({H}^{1}, {H}^{2}, ..., {H}^{h}){W}_O$.

Second, the gradients are also identical, i.e., $\frac{ \partial \textrm{MultiHead}({X},{X}') }{\partial H_{i} }={W}_O^{i}$. Therefore, the summation and concatenation forms are mathematically equivalent.

Q2: Theoretical guarantee for the effectiveness of the proposed methods.

A2: Thanks for your valuable suggestion. As mentioned by Reviewer xCL1, we put the theoretical guarantee that MoH is superior to multi-head attention in Appendix A. Specifically, we proved that MoH not only improves efficiency and model performance but also helps different attention heads to specialize better compared to multi-head attention, from both theoretical and experimental perspectives.

Q3: The methodology is incremental.

A3: The motivation for our work is to upgrade the multi-head attention mechanism, the core of the Transformer model, to reduce computational costs while maintaining or surpassing the previous accuracy level, rather than making improvements to the MoE. We propose combining multi-head attention and MoE, so we adopt the MoE technique:

• For the two-stage routing, as shown in our response (A1) to Reviewer 6acz, the reason for our two-stage routing is to capture general knowledge. Besides, we compare the gradients and training data distribution of shared heads and routed heads in Appendix Table A to further demonstrate that shared heads play a key role in capturing general knowledge.

• For the load-balancing loss, since we decompose multi-head attention into a summation form, which is similar to MoE structure, we directly adopt the auxiliary loss used in MoE.

Q4: The lack of discussion of the new experts defined.

A4: As suggested, we have defined the expert as $H^i W^i_O$ to avoid possible misunderstanding.

Q5: Further theoretical discussion regarding the effectiveness, the convergence rate and the optimization scheme.

A5: Mathematically, we prove that the summation and concatenation forms are equivalent. Besides, we show in Appendix A that the gradient per head in the MoH differs from the gradient of multi-head attention by only a single weight. Finally, we replace multi-head attention with MoH in various structures while keeping the original training parameters unchanged. The experimental results demonstrate that MoH enhances the performance of multi-head attention, providing experimental evidence of its effectiveness.

Q6: The abuse of notation $W$.

A6: As suggested, we have replaced the notation $W$ in the router with $E$, referring to them as expert embeddings.

Q7: Concerns about the claim of the author that the new methods can be fine-tuned from pretrained models.

A7: We explain this problem in three aspects:

• We simply select the first 16 attention heads of each layer as shared heads. Even if the structure is not optimal, the experimental results show that MoH-LLaMA3-8B has a significant advantage over LLaMA3-8B. This result shows the robustness of our method.
• Unlike the MoE upcycling technique, which copies the FFN to increase the model size, our MoH prunes the original model to reduce the activation parameters, making it more challenging.
• To the best of our knowledge, our work is the first to attempt to reduce the computational amount of the attention mechanism without degrading the model performance by continue-tuning on a pre-trained model.

Q8: Acknowledge the base code they adapted.

A8: Thanks for your advice. We will acknowledge the contributors in our official code, and release our trained models.

Q9: Existing works are not thoroughly discussed.

A9: As suggested, we have expanded and improved the discussion of related works.

Q10: What are the motivations to define $\alpha_1$ and $\alpha_2$ as in Eq. 6 rather than considering them as hyperparameters?

A10: We choose to predict $\alpha_1$ and $\alpha_2$ based on the input so that different tokens can dynamically combine general knowledge from the shared heads and specialized knowledge from the routed heads.

Q11: Explain the parameter-free router mentioned in Section 4.4 (line 319)?

A11: We use the $l_2$ norm (which measures the magnitude of a vector in Euclidean space) of each head to represent its importance. We then normalize the $l_2$ norm of all heads using SoftMax. This simple router achieves results comparable to learnable routers.

We sincerely thank the reviewer for the constructive comments, and for noting that our method provides "a new perspective on structured sparsity in Transformers" and that "the significance is high." We address the questions as below.

Q1: If MoH is inspired by Mixture-of-Experts, then why does it manually designate some heads as "always active" instead of letting the router handle all head selection dynamically?

A1: We explain this problem in three aspects:

• Load balance loss (also called MoE loss) pushes experts to focus on specific areas rather than general knowledge. As shown below, load balance loss ensures all experts are chosen equally and become specialized. However, this goes against the idea that essential experts should be selected more often and should learn broader, general knowledge. Due to the load balance loss, even though some heads are critical enough to remain active at all times, the routing function cannot naturally prioritize them.

$L_b = \sum_{i=h_s+1}^{h} P_i f_i,$

$P_i = \frac{1}{T} \sum_{t=1}^{T} \text{Softmax}(W_{r} x_{t})_{i-h_s},$

$f_i = \frac{1}{T} \sum_{t=1}^{T} {1}(\text{Token } {x}_t \text{ selects Head } i)$.

• From the perspective of training stability, keeping some heads active at all times helps maintain a stable gradient. If all heads are selected freely, the gradients and loss spikes can increase significantly, reducing training efficiency.

• In recent MoE work [1], some experts are also selected as shared experts to extract general knowledge. Besides, in attention mechanisms, some heads may capture common knowledge across different contexts, such as grammatical rules in language. Inspired by this idea, we designate a subset of heads as shared heads that remain always activated. We also compare the gradients and training data distribution of shared heads and routed heads in Appendix Table A to further demonstrate that shared heads play a key role in capturing general knowledge.

[1] Dai, Damai, et al. "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models." arXiv preprint arXiv:2401.06066 (2024).

Q2: Claims about computational efficiency are not supported by FLOPs/memory benchmarks.

A2: Thanks for your valuable suggestion. We have provided the comparison of the efficiency between MoH and multi-head attention mechanisms in Table 7, and we have added FLOPs and memory as suggested. Specifically, MoH surpasses standard multi-head attention mechanisms, with its advantage becoming more pronounced as the input sequence length increases.

We find that MoH slightly reduces GPU memory usage, though the difference is not significant. This is because GPU memory is primarily used to store model parameters, gradients, and the KV cache. Since MoH only optimizes attention computation, it does not substantially reduce the GPU memory of these three components.

Q3: No comparisons are made against alternative dynamic attention mechanisms (e.g., MoE-based attention, sparse routing).

A3: Thanks for your insightful advice. As shown in the table below, MoH outperforms both sparse attention and MoE-based attention. These results also demonstrate the benefits of shared heads from an experimental perspective.

We sincerely thank you for your constructive comments. We will add the above important discussions in the final manuscript and highlight them. Thanks again for taking the time and effort on our paper.

We sincerely appreciate your thoughtful comments and for recognizing that "The experiments in the paper are well-conducted." Below, we address your questions in detail.

Q1: One weakness of the method is its relatively high activation rate.

A1: To demonstrate the robustness of our method, we only replace multi-head attention with MoH in various structures while keeping the original training parameters unchanged. Therefore, the results presented in the manuscript may not be optimal. Our latest experimental results indicate that, with tuning, our method can achieve even higher performance with a lower activation rate. Besides, we are developing deep learning-based methods to automatically determine the optimal activation ratio. We believe that our proposed MoH is promising and can be further optimized for even better performance.

Q2: Lack of comparison with mixture of feedforward experts.

A2: Thanks for your insightful advice. We explain the difference between MoH and MoE from the following three aspects:

• Attention and FFNs are the core components of Transformers. While MoE applies sparse activation at the FFN level, MoH introduces sparsity at the attention level. MoH not only extends the scope of MoE but also offers a more effective approach to reducing Transformer computation. This is particularly significant because, as the input sequence length increases, FFN computation grows linearly, while attention computation scales quadratically. Consequently, MoH has greater potential to alleviate the computational burden of Transformers.
• MoH presents greater technical challenges than MoE. Unlike the MoE upcycling technique, which copies the FFN to increase the model size, our MoH prunes the original model to reduce the activation parameters, making it more challenging.
• MoH naturally leverages the multi-head structure in Transformers, while MoE requires additional FFN replication. From this perspective, MoH matches the original Transformer design better.

Q3: Lack of integration with mixture of feedforward experts.

A3: Thanks for your valuable suggestion. Due to the time constraints of the rebuttal, we designed a 28M-sized small model as a baseline, with a training budget of 100 epochs on ImageNet-1K classification. As shown in the table below, MoH can be combined with MoE, where MoE enhances the model's performance by replicating the FFN, while MoH optimizes computational efficiency by the dynamic activation of attention heads.

We sincerely thank you for your constructive comments. We will add the above important discussions in the final manuscript and highlight them. Thanks again for taking the time and effort on our paper.