SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

We would like to thank the reviewer for their insightful review and for positive comments on the clarity and methodology of the paper. Please find our responses as follows:

similar techniques for transforming attention heads into MoE have been proposed in … authors can consider mentioning these works [2,3,4 inter alia]

We would like to thank the reviewer for pointing this out. We will include the discussion of these papers in the final version.

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. …

We agree that at this spot, this description is not very clear (this becomes hopefully clearer later, in light of our description starting at L97 until Eq. 10, that introduces MoE everywhere, in both source and destination sides). The reviewer is right to point out that everything is a function of the input x, including the output of the routing functions. In self-attention, all of key, value, query vectors are indeed computed from the input x for each position. But once we have these vectors, from the viewpoint of attention computation, the current query defines the destination, and keys and values define the source. In the MoE version of attention, we allocate a routing function for each of key/value/query projection; these routing functions belong to the source or destination side accordingly. Now if we compare Eq 10 and Eq 6, one can notice that the routing function in Eq. 6 effectively corresponds to what we define as the destination-side routing in Eq 10. We will improve the clarity of this passage in the final version. Thank you for pointing this out.

the explanation is somewhat confusing, why is it the case that in the worst case all possible source projections have to be computed?

Considering the ‘row’ of an attention matrix as ‘destination’ and ‘column’ as 'source’, routing on the “destination side” means that we only select/use K experts to compute a row. However, different rows can still decide to use different selections of these K experts. This means that in practice, there are more than K active experts per column. In the worst case, i.e., if all rows decide to use different K experts, this requires K times the number of row active experts, or simply the total number of experts (if such a product exceeds the total number of experts, which is typically the case as there are more rows than the number of experts). This requires all possible key and value projections to be computed even if only a subset is used in each row. Again, we agree with the reviewer that this is only implicit and confusing in the current description. We will improve this in the final version. Thank you very much for pointing this out.

does the 4 in the calculation of the total MACs for the projections already include the output projection?

Yes, it does. We will clarify this in the final version.

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].

We would like to thank the reviewer for pointing this out. We will include a discussion of [1] in the final version.

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.

We would like to thank the reviewer for pointing this out. The main difference is that MoA fixes the number of K and V projections to 1, thus the efficiency limitations we describe do not apply to it. However, this limits its expressibility. Despite having a single key and value projections, it computes K different attention maps.

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"?

The parameter-matched setting is crucial to evaluate the model’s expressiveness in the LLM tasks where the number of parameters has a high impact on the model performance. We consider this setting to be particularly important to evaluate the true expressiveness of MoEs compared to their dense counterparts.

Please note that we also provide additional results in a MAC-matched setting in Tab 4.

does not exactly fit into the "Limitations" section, since, according to the authors (ll. 266-269) FlashAttention is orthogonal to SwitchHead

We agree with the reviewer that this paragraph does not fit here. We’ll move this to Appendix A. Thank you for pointing this out.

the visualization in Fig. 2 is only for 1 layer, and not for all layers as stated in the text We would like to thank the reviewer for pointing this out. We will fix this in the final version.

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?

The effects of σ-MoE and SwitchHead are orthogonal to each other, as they affect independent parts of the network. In the experimental setting of the paper, they are both set up to result in the most significant speedup with a marginal perplexity gain. Thus, by construction, combining them would never result in an increased perplexity. Speed-wise, with the current implementation, σ-MoE does not provide wall-clock gains on this model scale, in contrast to SwitchHead which provides a significant wall-clock speedup (see Tab 5). It saves, however, around 1.5Gb of additional memory compared to SwitchHead alone in the 262M param setting of Tab 5. With a higher d_model, the savings are significantly more substantial.

We believe our response above resolves all the concerns that the reviewer has raised. These questions are extremely valuable for us and will enable us to improve our paper. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.

We would like to thank the reviewer for their insightful review and for positive comments on the clarity of the paper. Please find our responses as follows:

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.

We would like to emphasize that the differences between MoA and SwitcHead are more than just the activation function. SwitchHead also uses multiple heads, while MoA uses a single Q and K projection. Moreover, MoA does the weighted average after the attention matrix computation, which makes it slower. Section 3.2 lists all these differences.

The use of sigmoid activation function in MoE was introduced by σ-MoE [1] and intuitively motivated by similarity to an approximate feedforward layer. Moreover, the authors of [2] provide a detailed theoretical analysis of the sigmoid activation function for MoEs, and they show that it converges faster.

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.

The parameter-matched setting is crucial to evaluate the model’s expressiveness in the LLM tasks where the number of parameters has a high impact on the model performance. We consider this setting to be particularly important to evaluate the true expressiveness of MoEs compared to their dense counterparts.

While the compute-matched setup has values when considering certain practical settings, it gives an “unfair” advantage to MoEs in terms of comparison, as we can easily add extra parameters to an MoE without significantly increasing compute requirements. Here we wanted to show that our SwitchHead is capable, even without considering such an advantage, by evaluating its pure expressiveness in the more challenging parameter-matched setting.

Please note that we also provide additional results in a MAC-matched setting in Tab 4.

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.

We would thank the reviewer for pointing this out. It is the total number of parameters. We will clarify this in the final version.

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.

Thank you very much for pointing this out. Shared selection refers to the case where we tie the V and O projections' selections to be identical. This marginally reduces compute and memory requirements at a slight performance cost. We agree this was confusing, we will provide the details in the final version.

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.

This is an excellent suggestion. Thank you for pointing it out. For measuring the resource usage of MoA, we chose the fastest MoA model that can match the performance of the dense baseline, or simply the best MoA model when no MoA model can match the baseline performance. This resulted in choosing MoA with H=4 for the 47M model and MoA with H=8 for the 262M parameter model. Please find the updated Table 5 below. It can be seen that SwitchHead outperforms MoA in both memory usage and runtime in both cases.

We believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.

[1] Csordás et al. EMNLP 2023 (Findings). Approximating Two-Layer Feedforward Networks for Efficient Transformers.

[2] Nguyen et al. Arxiv 2024. Sigmoid Gating is More Sample Efficient than Softmax Gating in Mixture of Experts.

We would like to thank the reviewer for their insightful review and for positive comments on the clarity and methodology of the paper. Please find our responses as follows:

I could not find the sequence length used in your experiments, which prevents me from computing the number of tokens used during training.

The sequence length is T in Tab. 6 in the appendix.

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.

The dropout is used in the switchAll models but only in the feedforward blocks. We do not use dropout in the attention components. In some sense, a dropout-like behavior is naturally provided by the expert selection mechanism: in the early stages of the training, the expert selection is “random” (not trained yet), and in all stages, only a small percentage of the total number of experts is selected (the others can be considered as "dropped out"). Moreover, the baseline only uses dropout on the projections, and neither model uses dropout on the attention scores.

What is explaining the differences in the transformers' PPLs in table 4 compared to the PPLs in table 2?

They correspond to runs with different seeds, for historical reasons. We will update Tab. 2 in the next version to avoid this confusion.

We believe our response above resolves the concerns that the reviewer has raised. Thank you very much.

We would like to thank the reviewer for their insightful review. Please find our responses as follows:

… 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]?

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]?

We agree with the reviewers that naturally, more experiments on other tasks will strengthen our work. That said, with our modest compute resources, it would be hard for us to perform all of these experiments. We restricted ourselves to LM tasks since they are of central importance today, and many of the other tasks, such as QA, can be casted as LM. We would like to note that we already conducted experiments on a list of LM tasks (enwik8, WT 103, C4, peS2o) on two different scales (47M, 262M) with two different positional encodings (Transformer, RoPE) and evaluated them zero-shot on Lambada, CBT, and BLiMP. Currently, we are not advocating using our method for tasks other than self-attention in LM, although we believe it would have similar properties on those as well.

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.

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?

While we agree with the reviewer that comparing attention maps is an interesting topic, we are unaware of any quantitative analysis that can be indicative of a “quality of an attention map”, except the downstream performance of the model, which we measure on a wide variety of tasks and datasets. For example, it is hard to draw conclusions from raw entropy measurements as we could argue on both sides: high entropy attention maps integrate more information and distribute gradients to more tokens, potentially providing a better learning signal. On the contrary, they also “blur together” too much information, making it hard to filter out irrelevant tokens and sharply focus on the relevant ones. High entropy attention maps might be a mere byproduct of some heads being unused by the model (by setting its contribution low through V and O projections, which do not show up in the attention maps). We would also like to note that on algorithmic tasks, such as ListOps shown in Fig. 2, typically low-entropy, sharp, and simple attention maps work better, as these require exact focus on specific numbers and operations. Therefore, our current conclusion is to rely on the downstream performance of the model, assuming that models, including their attention maps, have to have good behavior to achieve good performance.

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?

This is a very interesting open question. As this is an empirical observation in our experiments, we have no theoretical explanations. In fact, the situation is the same for the general theoretical studies on MoE; the current literature lacks clear explanations on when exactly regularization is necessary in MoE models. Also, the direct comparison between sigma-MoE and Switchhead is not straightforward as one is selecting slices of two layers in an MLP simultaneously, and the other one selecting individual projections in the attention.

... 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?

… a more granular analysis of the trade-offs between model size, computational resources, and performance could enhance the paper's contribution to the field.

In paragraph 2 of Sec 3 (L124-131), we provide the details of our algorithm for selecting K. K=2 usually works well enough. Generally, increasing K always helps, but it makes the model slower and uses more memory since more projections have to be computed. Thus, the choice of K is a tradeoff. We chose the minimal K that matches or slightly outperforms the equivalent dense model. Increasing d_head and the number of heads have a similar effect.

Generally, we do agree with the reviewer that a more extensive hyperparameter analysis would strengthen the paper, but our compute resources are modest. We focused on the ablation studies that are maximally informative about our new model, such as verifying which of the projections are necessary to use (see Tab. 6 in the Appendix).

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?

We refer to “preliminary experiments” as experiments not included in the paper, but done before the final experimental protocol is decided (to the best of our knowledge, this is a rather common terminology). They are typically conducted in order to determine the final model/protocol to be studied and presented in the paper, i.e., the final setting on which we spend our compute resources. These are not necessarily done in the exact same setting as the experiments reported in the paper and, thus, are not directly comparable. As we mention in L512-513 of the appendix, we have done an order of magnitude more of these experiments than the ones included in the paper.

We hope our response above brings clarifications to all the concerns that the reviewer has raised. Thank you very much.