SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

Summary

This paper introduces SwitchHead, a novel MoE architecture for the attention layer. Unlike the existing MoA approach, which computes the output of the attention layer as a weighted average of the top-k heads determined by a learnable routing mechanism, SwitchHead independently applies expert mixtures to the heads' key, value, and output projections. This design allows SwitchHead to achieve comparable or superior performance to MoA and dense Transformers while requiring less computational power. The reduced compute cost is achieved by performing fewer matrix multiplications: since fewer heads need to be instantiated, the attention matrix operations are performed less frequently (same number of params is ensured by setting d_head accordingly). Reduced memory cost is achieved by having to store less activations for the backward pass. The paper points out that only applying MoE to value and output projection in each head performs sufficiently well as compared to applying MoE to the key and query projections as well.

Strengths

Overall, this is a well-crafted paper that presents a straightforward yet effective method for efficiently integrating MoE layers into self-attention blocks. Originality: good. The idea of implementing of transforming self-attention into an MoE is not new, yet the specific method presented and analyzed here is novel enough. Quality: good. The claims are well supported by evidences. Experiments are well designed and seem to be sufficient overall. Clarity: the paper is well organized. In terms of writing, some formulations require further clarifications (see weaknesses/questions) Significance: good. The paper makes a moderate contribution to the field and can potentially impacting future research.

Weaknesses

I cannot identify any fundamental weaknesses with the paper beyond several confusing passages that might require further clarification, and a couple of potentially missing citations (see questions).

Questions

- similar techniques for transforming attention heads into MoE have been proposed in the context of parameter efficient MoEs for the fine-tuning stage where experts are represented with LoRA adapters, authors can consider mentioning these works [2,3,4 inter alia] - ll. 88 - 91: "Intuitively, the above method corresponds to choosing a subset of attention heads based on the destination side alone": given that routing vector s is calculated conditioned on the input x, I do not understand why is it the case that attention heads are selected based on the destination (queries and outputs) side alone. In my understanding the attention head selection is dependent on the input x. - ll. 93 - 96 (related to above point) the explanation is somewhat confusing, why is it the case that in the worst case all possible source projections have to be computed? - l. 439: does the 4 in the calculation of the total MACs for the projections already include the output projection? - ll. 158 - 161: it might be worth noting that the idea of sharing key and value projection across heads has been introduced in earlier work on multi-query attention [1]. Additionally, to enhance the overall clarity, it could be useful to explicitly highlight how exactly MoA differs from the naive attention head mixing technique described in the first part of the Sec. 2.2. - ll. 119 - 120: it would be helpful if authors could provide a brief discussion clarifying how and why exactly parameter-matched setting "better reflects the task of language modelling"? - ll. 225 - 227: the visualization in Fig. 2 is only for 1 layer, and not for all layers as stated in the text - ll. 282 - 287: does not exactly fit into the "Limitations" section, since, according to the authors (ll. 266-269) FlashAttention is orthogonal to SwitchHead - how important is it to have MoE in the attention the the fully-connected block is already an MoE? Could authors add comparison to σ-MoE? Overall solid paper, I am happy to revisit my score upon the rebuttal. [1] Shazeer, Noam. "Fast transformer decoding: One write-head is all you need." arXiv preprint arXiv:1911.02150 (2019). [2] Page-Caccia, Lucas, et al. "Multi-head adapter routing for cross-task generalization." Advances in Neural Information Processing Systems 36 (2024). [3] Zadouri, Ted, et al. "Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning." arXiv preprint arXiv:2309.05444 (2023). [4] Ostapenko, Oleksiy, et al. "Towards modular llms by building and reusing a library of loras." arXiv preprint arXiv:2405.11157 (2024).

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are addressed in section 6.

Summary

This paper introduces SwitchHead, a Mixture of Experts (MoE) method for improving the efficiency of the self-attention layer in Transformers. Unlike traditional MoE methods focused mainly on feedforward layers, SwitchHead effectively reduces both compute and memory requirements, achieving significant wall-clock speed improvements without compromising language modeling performance. This is accomplished by reducing the number of attention matrices needed, offering substantial savings in computational resources, which is validated through extensive experiments across multiple datasets and model sizes.

Strengths

The paper is well-written and easy to follow. The proposed method demonstrates substantial improvements in computational efficiency by reducing the number of attention matrices needed, which leads to lower compute and memory usage compared to standard Transformers.

Weaknesses

The proposed SwitchHead method shares similarities with the Mixture of Attention Heads (MoA), but a key distinction lies in its use of a non-competitive activation function (sigmoid) instead of SoftMax. This design choice is important to understand, and I would appreciate an explanation for opting for the sigmoid activation function over SoftMax in this context.

Questions

1. The authors mention employing a parameter-matched setting to better reflect the task of language modeling. However, the paper lacks explicit justification for why this setting is more representative of language modeling tasks. Additionally, the description of parameters in Table 1 is ambiguous—it should be clarified whether the numbers represent total parameters or only those that are actively used during inference. 2. The paper occasionally uses notations without adequate explanations, which can lead to confusion. For instance, the term “Shared selection” used in Table 4 is not clearly defined. 3. In Table 5, the authors present the training times for the baseline Transformer and SwitchHead models, but the training time for the MoA model is not included. To provide a comprehensive comparison and better evaluate the efficiency of SwitchHead relative to other models, it would be beneficial if the authors could also include the training time data for the MoA model.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitations of their study.

Summary

This paper presents SwitchHead, a Mixture of Experts (MoE) method applied to the self-attention layer in transformer blocks. By applying MoE to the self-attention layer, SwitchHead reduces computational and memory costs while maintaining language modeling performance comparable to traditional dense models. It can be combined with existing MoE methods for feedforward layers, resulting in a fully MoE-based transformer, SwitchAll transformers. Tests on various language modeling datasets (C4, Enwik8, peS2o) show that SwitchHead performs well compared to models with parameter-matched settings.

Strengths

- The proposed MoE method applied to self-attention layer reduces computational and memory requirements while preserving language modeling performance of dense baselines. - The authors make effort to provide a fair comparison between the baselines and the proposed model in parameter-matched settings. They also include both the MAC and wall-clock speedup comparisons, clearly demonstrating the efficiency of the proposed method. - The author performs extensive ablation studies to show that the proposed method performs well with different types of models and datasets. They also provide the hyperparameters used for each experiment, which helps in reproducing the benchmark results.

Weaknesses

- The paper only evaluates the proposed method on language modeling datasets, lacking demonstration of its effectiveness on other important NLP tasks such as document summarization or open-domain question answering. - While the authors claim similarity between SwitchHead and dense baseline attention maps, the provided figures suggest simplification in SwitchHead's maps. The lack of quantitative analysis (e.g., entropy measurements) weakens this claim. - There's a concern that the simplified attention maps produced by SwitchHead could lead to performance degradation when SwitchHead used in encoder-based models due to information bottleneck, which isn't addressed in the current evaluation. - The paper doesn't clearly explain why regularization methods used in σ-MoE (e.g., entropy maximization, expert dropout) are unnecessary for SwitchHead, potentially leaving gaps in the method's theoretical foundation. - The paper suggests that Top-K selection can be treated as a hyperparameter, but doesn't provide a clear analysis of the method's sensitivity to different K values, leaving questions about the necessity of hyperparameter search. - The paper mentions "preliminary experiments" multiple times without proper citation or section indicators, which may confuse readers and reduce the reproducibility of the work.

Questions

- While the paper demonstrates strong performance on language modeling tasks, it would be beneficial to see results on a broader range of NLP tasks. How does SwitchHead perform on tasks such as document summarization[1] or open-domain question answering[2]? This would help demonstrate the method's versatility and potential impact beyond language modeling. - The paper claims that attention maps from SwitchHead and dense baseline models (Transformer XL) are qualitatively similar. However, Figures 2 and 6 suggest that SwitchHead produces simplified attention maps. Could you provide quantitative analysis (e.g., entropy measurements) to more rigorously demonstrate the complexity and quality of SwitchHead's attention maps compared to the dense baseline? - If SwitchHead indeed produces simplified attention maps, there's a concern about potential information bottleneck in encoder-based models. Have you considered evaluating SwitchHead's performance when applied to the query encoder in a retrieval-augmented generation (RAG) setup for open-domain QA tasks[2]? This could help address concerns about information loss in more complex, multi-step tasks. - The paper mentions that SwitchHead doesn't require the extra regularization techniques used in σ-MoE, such as entropy maximization for load balancing or expert dropout. Could you clarify what specific differences between σ-MoE and SwitchHead make these extra tricks unnecessary? A more detailed explanation of this point would strengthen the theoretical foundation of your approach. - You demonstrate that K in Top-K selection can be treated as a hyperparameter in the proposed method. How sensitive is the model's performance to different K values? If the sensitivity is low, is extensive hyperparameter searching always necessary, or could a default value be recommended for most use cases? - The paper mentions "Our preliminary experiments" several times (e.g., lines 91, 473) without providing citations or section indicators. Could you clarify these references to improve the paper's clarity and reproducibility? Consider either expanding on these preliminary experiments in the main text or including them in an appendix. [1] Nallapati, Ramesh, et al. "Abstractive text summarization using sequence-to-sequence rnns and beyond." arXiv preprint arXiv:1602.06023 (2016). [2] Kwiatkowski, Tom, et al. "Natural questions: a benchmark for question answering research." Transactions of the Association for Computational Linguistics 7 (2019): 453-466.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

- In section 6 and appendix A.1, the authors appropriately present the limitations and societal impacts of their work. This acknowledgment of their model framework's constraints effectively strengthens the motivation for additional research topics in future work. The authors' awareness of these limitations demonstrates a commendable level of scientific rigor and transparency. However, the claim that performance can potentially reach 80-90% raises some questions. While this projection is intriguing, it would be more convincing if supported by more robust quantitative assessments. The addition of more detailed quantitative evaluations could provide valuable insights and significantly aid future research efforts. For instance, a breakdown of current performance bottlenecks and a roadmap for potential optimizations would offer clearer guidance for researchers looking to build upon this work. Furthermore, a more granular analysis of the trade-offs between model size, computational resources, and performance could enhance the paper's contribution to the field. Overall, while the authors have done a commendable job in addressing limitations and societal impacts, the inclusion of more quantitative metrics and projections would further solidify the paper's value as a foundation for future research in this area.

Summary

This work proposes a novel Multi-Head-Attention (MHA) mechanism which is more efficient than the standard MHA used in most transformers. The method---called SwitchHead---is relying on Mixture-of-Experts (MoEs) to save computation while retaining the same model capacity and performance. To achieve this, instead of naively selecting which head should be computed for each token, a small number of heads are **always** computed, and MoEs are used to modulate the value and output projections for each head. Using multiple transformer variants and datasets, they empirically show how SwichHead transformers reach a similar perplexity as the baseline, for the same number of parameters, while requiring fewer operations and memory. They also show how their approach can yield a significant speedup during training.

Strengths

The paper is well written and well motivated. While the components of the proposed method are not particularly novel, they are combined in a simple yet novel way. I find the experiments convincing. Matching the parameters and perplexity allows to clearly appreciate the gains in memory and MACs. The results of the MAC-matched experiments and the time comparisons further reinforce the potential impact of the proposed method.

Weaknesses

- As mentioned in the limitations section, the models are relatively small. - I could not find the sequence length used in your experiments, which prevents me from computing the number of tokens used during training. - I find it odd that dropout is used for the baseline but not for the switchAll models, on C4 I'm expecting dropout to slow down learning. - What is explaining the differences in the transformers' PPLs in table 4 compared to the PPLs in table 2?

Questions

See above.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Limitations have been discussed adequately.

Official Comment by Reviewer YarL

Thank you to the authors for their comprehensive response to the review comments. The efforts made to address the concerns and clarify the points discussed are highly appreciated. Given these responses, I am inclined to increase my score.