SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention
Authors: Róbert Csordás
Piotr Piękos
Kazuki Irie
Jürgen Schmidhuber
Venue: NeurIPS 2024 poster (2024)
Keywords: MoE
mixture of experts
attention
transformers
Primary area: deep_learning_architectures
TL;DR: We reduce the resource usage of the attention layer with Mixture of Experts.

Abstract: Despite many recent works on Mixture of Experts (MoEs) for resource-efficient Transformer language models, existing methods mostly focus on MoEs for feedforward layers. Previous attempts at extending MoE to the self-attention layer fail to match the performance of the parameter-matched baseline. Our novel SwitchHead is an effective MoE method for the attention layer that successfully reduces both the compute and memory requirements, achieving wall-clock speedup, while matching the language modeling performance of the baseline Transformer. Our novel MoE mechanism allows SwitchHead to compute up to 8 times fewer attention matrices than the standard Transformer. SwitchHead can also be combined with MoE feedforward layers, resulting in fully-MoE "SwitchAll" Transformers. For our 262M parameter model trained on C4, SwitchHead matches the perplexity of standard models with only 44% compute and 27% memory usage. Zero-shot experiments on downstream tasks confirm the performance of SwitchHead, e.g., achieving more than 3.5% absolute improvements on BliMP compared to the baseline with an equal compute resource.

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention
Róbert Csordás
Piotr Piękos² Kazuki Irie Jürgen Schmidhuber
Stanford University. Stanford, CA, USA
Al Initiative, KAUST, Thuwal, Saudi Arabia
*Center for Brain Science, Harvard University. Cambridge, MA, USA
"The Swiss AI Lab IDSIA. USI & SUPSI, Lugano, Switzerland rcsordas@stanf ord. edu, piotr p piekos@kaust.edvv sa, sa,
kirieffas harvard edu, juergen@idsia.ch
Abstract
Despite many recent works on Mixture Experts (MoEs) for resource-efficient Transformer language models, existing methods mostly focus on MoEs for feedfor- ward layers. Previous attempts extending MoE the self-attention layer fail match the performance of the parameter-matcher baseline. Our novel SwitchHead is an effective MoE method for the attention layer that successfully reduces both the compute and memory requirements, achieving wall-clock speedup, while matching the language modeling performance of the baseline Transformer. Our novel MoE mechanism allows SwitchHead to compute up 8 times fewer attention matrices than the standard Transformer. SwitchHead can also combined with MoE feed- forward layers, resulting in fully-MoF: "SwitchAll" Transformers. For our 262M parameter model trained on C4, SwitchHead matches the perplexity standard models with only 44% compute and 27% memory usage. Zero-shot experiments on downstream tasks confirm the performance of SwitchHead, e.g., achieving more than 3.5% absolute improvements BliMP compared to the baseline with equal compute resource."
Introduction
X
weight matrix
SEL
autelection logic
A
Wv
SEL
Wg
y
Figure 1: A schematic representation of SwitchHead. It consists of a few independent heads, each with multiple experts for value and output projections. Each head has single attention matrix.
Large language models (LLMs) have shown remarkable capabilities [1, 2, 3, 4] and great versatility [5]. However, training large Transformers [6, 7] requires considerable amount of computing power and memory, which is not accessible to most researchers, academic institutions, and even companics.
Work done at IDSIA.
Our code is public:
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Even running them in inference mode poode-typically much less resource- intensive intensive-requieees requires significant engineering effort [8]. Accelerating Transformers remains an important research question.
In this context, Mixture of Experts (MoE) layers [9. 10, 11] have become popular to efficiently scale up Transformers a large number of parameters [12, 13, 14, 15, 16, 17]. However, most of these works mainly focus on applying MoE to the 2-layer feedforward blocks [6], i.c., the multi-layer perceptron (MLP) components of the Transformer, while keeping the self-attention layers unchanged. Given that attention also accounts for a considerable amount of compute and memory usage in Transformers (especially for long context sizes), using MoE attention has potential to further improve resource efficiency in Transformers. While MoE-based attention remains underexplored in general, there are existing works on MoE approaches for attention [18, 19]. However, in practice. previously proposed methods typically require a lot of engineering tricks for successful training, and most importantly, only achieve modest reduction in computing and memory requirements in the end (as we also confirm in our experiments).
Here, we present novel MoE-based attention method, SwitchHead, whose mechanism allows to reduce the number of attention matrices that need to be computed and stored. Following a-MoE [17], our method uses non-competitive selection activation function (sigmoid), and does not require regularization or extra tricks for stable training. Importantly, we show that it is possible to compute the MoE projections outside of the attention core, which enables a significant reduction in the number of computed attention maps, resulting in significant resource savings. Our thorough investigation shows that it enough to choose the value and output projections from a pool of experts and share keys and queries between them.
We evaluate our method C4 [20], Enwik8 [21], peS2o [22] and Wikitext 103 [23], with two model sizes (47M and 262M). Additionally, we measure the zero-shot performance of our main models on Lambada [24], BLiMP [25], and Children's 's Books Test [26] datasets. Our experiments demonstrate that SwitchHead can achieve performance comparable to parameter-matched baselines with just fraction of the compute and memory budget. In addition, we introduce "SwitchAll", fully MoE-based Transformer model, that combines o-MoE- MLP layer with our SwitchHead attention, often outperforming dense baselines with the same parameter budgets.
Finally, we analyze the attention maps of our SwitchHead. We find that the attention maps taken over all heads qualitatively similar to the dense baselines, indicating significant reduction in redundancy without a loss of expressivity. In addition, expert selections are often interpretable.
2 Method
2.1
Background
The standard multi-head self-attention (MHA) layer [6] consists of four major steps: (1) compute key, query, and value projections. (2) compute the attention matrix, (3) use the attention matrix project the values, and (4) map the projected values to the output. Let h. T, Theedss dloodel dhead denote positive integers. Let x E ET ddunn denote an input the MHA layer with Theads heads, 1 be the sequence length, and demodet denote the size of the hidden representations of the model. Whx (K,VQQ E IR dhenatt X dhost
are the projection matrices for head h E (1,..., Pheads) Then K x K Q x: W and y x W (thus Kh, Q², V E IR xxdeel) are the keys, queries, and values, respectively. The attention matrix for the head h, A E RxT, and the output y E dxd'da are calculated as follows:
A = softmax
Q"K"'
Chead
= (A'V
V
|..| AA Thedd Wasd Wo
(2)
where denotes concatenation the last dimension, the softmax(-) is also over the last dimension, and Wo I Rhes dheal cheal X Xhuakd However, an alternative formulation reflects the role of Wo better. Let
us divide Wo along the first dimension into submatrices for each head, W E &Rdaes Xdmaas such that
Wo (W) I W Whsatt In this case, the output (Eq. 2) can be equivalently written as:
Σ A h
(3)

---

From this, it can be seen that all computations are local to each head. Computing the attention matrix A and the readout V requires compute in order of (Tneatts T²) MACs (multiplication- accumulation operation"). During training, it requires the storage of ( (2) for the attention
matrices and (Theeuss T dheaa)) for storing the sub-results of the projections. Given a sufficiently long sequence, computing the attention matrix and projecting the values will dominate the compute requirements due to the quadratic dependence on the sequence length T.
2.2 From Dense to SwitchHead Attention Layer
Our goal is to obtain resource reductions while maintaining the fundamental properties of attention and retaining a fully expressive attention matrix. For that, we start from the following observation: modern LLMs use tens of heads [2. 27]. Are so many of them all necessary? As we show later in Sec. 3, indeed, naively reducing the number of heads (while keeping the same number of parameters by increasing the head dimension) results in performance loss. Explaining the reason for the need for many heads is beyond the scope of this paper. Nevertheless, here are some hypotheses: (I) they provide multiple inputs for the operations that the network performs in each step, (2) they are specialized and provide inputs only for specific operations (in this case, each operation would use a different subset of heads), (3) they may provide diverse outputs due to different initializations, some being more successful than others, thus enabling better learning. Among these, (2) and (3) may offer an opportunity for resource savings: if not all heads are needed at the same time, it might be possible to switch among them depending on the context.
One naive method to achieve this is to use gating signal using linear projection Ws E RR Chand Rhaaa and use the heads with the highest score, by replacing Eq. with Eq. 6:
a sx)
E = arg topk 8, E),E C (1,..., Theaads)
(5)
y(t. c) = Σ [Al
hEE
where y|t,c] R denotes indexing the specific element of the output matrix y RTxuuai for timestep l. and channel c, and k is the number of active experts. Following the a-MoB method [17], use non-competitive selection function (sigmoid in Eq. 4). Now, let define the source side of attention as the keys and values and the destination side as the queries and output. Intuitively, the above method corresponds to choosing subset of attention heads based on the destination side alone". Our preliminary experiments confirmed that this method is indeed feasible for language modeling on WikiText-103. However, it difficult achieve acceleration and memory savings with this method. To see why, notice that the entries of the attention matrix A h depend on pairs of tokens, one for the source and one for the destination side, but the choice is made only based on the destination side. Thus, in the worst case, for each destination, a different source might be chosen, in which case all possible source projections have to be computed for the keys and values, which we would like to avoid.
Alternatively, we propose to improve the method above by introducing conditional computations for the source and destination projections independently of each other. That is, we parameterize each of key, query, value, output projection by an independent MoE. This avoids conditional computations that involve the attention matrix itself. Our solution implements this using Mixtures of Experts (MoEs). The concepts of "heads" are no longer well defined in the conventional sense: we redefine a head as an instance of a computed attention matrix. We call the total number of them Theads. For each head h, we define a separate list of E experts. The total number of experts is then Pheads E. Then, the projection matrices become Whe whes whe and O E R. x daseel where h denotes the head index and the specific expert. Then we compute the source-side expert selection as follows:
2 a (
8 arg topk(s k), E' C (I,..,
(8)
"The number of MACs is metric used in prior work [18], which is independent of both the specific hardware and implementation, unlike wall-clock time. For wall-clock-time measurements, see Sec. 3.7.
'To clarify, we allocate a routing function for each of key/value/query projections; these routing functions belong to the source or destination side accordingly. If we compare Eq. 10 and Eq. 6, one can notice that the routing function in Eq. 6 effectively corresponds what define as the destination side routing Eq. 10.
3

---

where W E R We compute the destination-side experts similarly: sD a(> D). EB arg top k), C E}, W Rd X. Then, the value projection V is
computed weighted sum of the selected experts:
Vh.
Σ BEE9
(9)
The key and query projections are computed similarly: K Eeeey s s'see<>W> and Q*
Exeey s The output projection also becomes an MoE:
Phoaak
y
Σ REED
(10)
As we' show, it not necessary to make all projections MoEs. In Section 3.1 we show that keeping single, head-speeifie copy of the query and key projections and reusing them for all experts is beneficial. We call this method SwitchHead.
Essentially, SwitchHead reduces the number of attention matrices that have to be computed (Theads) significantly, by using multiple experts per head. Note that our method does not depend on the specific implementation of the attention, allowing for easy experimentation and research. A schematic representation shown in Figure
Table 1: Performance of SwitchHead compared to different MoA variants. MoA can outperform the baseline, but only at price of using significantly more compute and memory. Also, SwitchHead outperforms the baseline dense Transformer. These results are on Wikitext 103. Table sorted by model perplexity.
#total params
Model
Tiseads
Perplexity
MACs
Mem (floats)
47M SwitchHead
12.27
170.4M
0.8M
Transformer
12.31
453.4M
3.5M
MoA
12.60
223.5M
1.3M
MoA
12.64
306.8M
1.9M
MoA
8
12.77
390.2M
2.6M
MoA
12.84
140.1M
0.7M
262M
MoA
9.50
2.9G
9.9M
SwitchHeaa
2.0G
2.9M
Transformer
16
9.66
5.4G
21.0M
MoA
12
9.68
4.1G
14.7M
MoA
4
9,69
1.7G
5.1M
MoA
9.87
1.1G
2.7M
3 Experiments
We conduct our experiments in purameter-mutched setting [17] which better reflects the task of language modeling (than the FLOPS-matched setting often used to evaluate MoEs), Our main experiments use Transformer XI., because we found them to consistently and significantly outperform RoPE-based basclines [28] for a fixed amount of compute. We provide the details of this analysis in Appendix A.4. The conclusions on the effectiveness of SwitchHead are consistent in both cases.
As an important specification, under this parameter-matcher setting, we always contigure Switchhead such that it matches the perplexity of the baseline dense Transformer, and we maximize its resource reductions. For this, we follow systematic procedure. First, we set Tipeads of the dense baseline. We start with setting Theaas reductions. the resulting model underperformss we increase k: 1: 4 underperforms as well, we set Thends 4 and A/ = 2. We always set dhead so that the total number of parameters of the resulting model matches the number of parameters the baseline. This reasonably simple procedure ensures good amount of resource savings, while avoiding doing expensive hyperparameter search.
be the same as Theads and A: 2, which provide the most resource

---

Note that all the perplexity gains seen in the main result tables are the byproduct of imperfect matching, and our goal is achieve reductions in resource requirements. unless noted otherwise (See Sec. 3.5). Detailed hyperparameters of all our models can be found in Sec. A.5 in the Appendix. We use and adopt the Triton kernel of o-MoE [17] for our purposes.
For all datasets except the character-level Enwik8 [21], a SentencePiece tokenizer [31] with vocabulary size of 8k tokens. For most of our experiments, we use Transformer XL [32] with the context size being twice the size of the active/current chunk, because we found it to be significantly more resource-eficieen than the standard setup. However, in order to show that our method is also competitive in the standard Transformer with RoPE positional ecodings, we also demonstrate our main findings in this setup (Appendix A.4).
we use sub- word units [29, 30] obtained with
All models are trained for 100k batches. Some of the datasets we consider (C4 [20], and peS2o [22]) much larger. In this case, we train on the first 10 T Nuatch tokens of the dataset.
3.1 Which Projections Require an MoE?
As discussed in Sec. 2.2, each linear projection (keys, values, queries, and output) can potentially be replaced independently by MoE. Here first check which projection benefits from such replacement. As we target the parcheter-matched setting, using MoE where it is not necessary can have a negative effect. Since experts the number of parameters available for the more useful parts of the model. Thus, we did a search over all possible combinations of MoE versus fixed projections with two active heads and compared them to the parameter-matched baseline. We find that the output projection is necessary to match the performance of the baseline (for detailed results refer to Tab. 6 in the appendix). Having MoE in the key and query projections turn out to be unnecessrry. Models without the output and value MoE underperform the dense baseline with Thhaads 2 heads.
significant part of the parameter budget, they can reduce
In sum, the perffrrrfon model is the one using MoE for value and output projections. We use this model variant in the rest of experiments in this paper.
3.2 Comparison with MoA
The method most related to ours is the so-called Mixture of Attention Heads, or MoA [18]. Unlike SwitchHead, MoA uses single key and value projection and chooses 7hhaads active query and output projections from a pool of E experts.
Moλ computes the attention map for each selected expert and computes their weighted average after the attention computation takes place. In contrast, SwitchHead calculates the weighted average of the K selected experts before and after attention computation. Because of this, in practice, the same perplexity is achieved with the required number of computed attention matrices (Theads) which is much lower for SwitchHead compared to MoA, allowing significant resource savings.
Also, unlike MoA. SwitchHead uses non-competitive activation function (sigmoid) [17]. We confirm that with this, our method performs well without any regularization, while MoA requires three different regularizers.
We compare our method with MoA in Table 1. It can be seen that while MoA can slightly outperform method in terms of perplexity, it can only do at the price of significantly more resource usage. Given a similar computation and memory budget, our method consistently outperforms MoA.
3.3 Performance on Different Datasets
We test our methods diverse set of language modeling datasets, including C4 [20], Enwik8 [21], peS2o [22], at two different scales: a 47M and 262M parameters. We chose this experimental setting taking into account our bompute-budget and confidence in our results which are consistent in across various configurations.
The results are shown in Table 2. We compare our models to two baselines: one with the same number of heads as the total number of experts (Theass F.) of the SwitchHead models, and the other has the same number of heads as the number of active attention matrices (Dhcads) as models. Our
5

---

Table 2: Performance of SwitchHead compared to baselines different datasets and model sizes. It can be seen that the predictive performance of our SwitchHead model is comparable to the baselines, and is always better than the baseline with equal number of heads. Perplexity is shown for Wikitext 103, C4 and peS2o datasets, and bits/character (bpc) for Enwik8. Models sorted by perplexity.
Dataset
#total params
Model
Thyads
ppl/bpe
MACs
Mem (floats)
C4
47M
Switchleead
22.53
203M
0.8M
Transformer
10
22.71
453M
3.5M
Transformer
2
23.71
453M
1.4M
262M
SwitchHead
16.23
2.4G
5.6M
Transformer
16
16.28
5,4G
21M
Transformer
17.09
5.4G
8.4M
Wikitext 103
47M
SwitchHead
12.31
170M
0.8M
Transformer
10
12.32
453M
3.5M
Transformer
12.73
453M
1.4M
262M
SwitchHead
9.77
2.0G
2.9M
Transformer
16
9.80
5.4G
21M
Transformer
10.09
5.4G
6.3M
peS2o
47M
Transformer
10
12.83
453M
3.5M
SwitchHead
12.84
203M
0.8M
Transformer
13.37
453M
1.4M
262M
Transformer
16
9.78
5.4G
21M
SwitchHead
9,86
2.4G
5.6M
Transformer
10.11
5.4G
8.4M
Enwik8
41M
Transformer
1.10
1.6G
10M
SwitchHead
1.10
709M
2.8M
Transformer
1.13
1.6G
4.2M
models closely match the performance of the full, many-head baseline with the fraction of memory and compute requirements (see Sec. 3.7 for more details).
In addition, we verify the performance of our models trained on the C4 dataset downstream tasks in zero-shot manner. We consider Lambada [24], BLiMP [25] and Children's Book Test (CBT) [26]. The results are shown in Table 4: our SwitchHead models consistently outperform or match the performance of the baseline dense Transformer models.
3.4 SwitchAll
The goal of achieving more resource-efficien Transformers includes reducing the resource require- ments of both the MLP and the attention layers. o-MoE [17] was recently proposed parameter- efficient MoE method for accelerating the MLP layers. However, it remains unclear whether can be efficiently combined with SwitchHead, can have some negative interaction effect if combined in "SwitchAll", where every layer is MoE-based.
To verify this, take the baseline architecture of Csordás et al. [17] without any hyperparameter change and replace the attention layer with SwitchHead. The hyperparameters for attention are directly taken from the experiments shown in Tab. 2. The results are shown in Tab. 3. The combined, fully-MoE model often outperforms the dense baselines for each dataset and model size considered, except in the case of the 262M parameter model on the C4 dataset.
3.5 MAC-Matched Setup
All our experiments so far were calibrated so that the predictive performance (perplexity) matches to the performance of the baseline Transformer, and we were aiming for maximum resource savings. However, it is also a valid question setup, where the compute requirements of our model are matched to those of the baseline. We achieve this by increasing dhend and Tineas until we have the same MAC requirements as the baseline. This results in model with more parameters. For the small Transformer XL, increase dhead from 76 to
what the performance of SwitchHead a MAC-matched

---

112 and Nbeads from to 3. For large XL, we increase Pheads from 4 to 6 and dhead from 112 to 168. For the small RoPE model, we change Theads from to 3 and dmnddl from 64 84, for big Teadd from 4 to 6 and maddel from 112 to 168. We show the results in Tab. 4: MAC-matched models outperform the others by large margin both in perplexity and zero-shot task performance.
3.6 Shared Selection
For further time savings, we share the expert selection between the source and destination side. Acceleration is achieved by reducing the number of sorting and top-k steps compared to the full SwitchHead. However, this results minor performance loss, which might be tolerated in some cases where the acceleration is more important. See Tab. 4 for more details.
3.7 Wall-Clock Time and Memory Usage Estimation
In all of our tables, we report the number multipy--accumulatc (MAC) operations following Zhang et al. [18]. The reason for this is that the actual wall-clock time is highly implementation and hardware-dependent. Nevertheless, we measured the runtime and total memory usage of our entire training pipeline (including the feedforward layer) to demonstrate that our current (suboptimal) implementation is already capable of providing wall-clock time acceleration. We show the results in Tab. 5. The measurements are taken identical hardware with the implementation (including for the attention core), the only difference being the MoE-based projections for the attention. It can be seen that for both scales, SwitchHead trains around 1.5 times faster, while using 61 %-67% as much memory as the baseline.
We also report the performance of MoA for reference in Table 5. For measuring the resource usage of MoA, we chose the fastest MoA model that can match the performance the dense baseline, simply the best MoA model when MoA model resulted in choosing MoA with = 4 for the 47M model and MoA with H 8 for the 262M parameter model. SwitchHead outperforms MoA on both scales, both in wall clock time and memory requirements. Note that these measurements also include the MLP layers, the optimizer, and the gradient synchronization in the case of multi-GPU training.
match the baseline performance. This
Table 3: Performance SwitchAll (SwitchHead + o-MoE [17]) on different datasets and model sizes. Our SwitchAll model close or better compared to the baselines. Models sorted by perplexity. Note: We show the parameter count of the dense model. The parameter count for the big SwitchAil model is 259M because of the imperfect parameter matching.
Dataset
#total params
Model
ppl 4
MACs
Mem (floats)
Wikitext 103
47M
SwitchAll Transformer
2
12.17
170M
0.8M
10
12.32
453M
3.5M
262M
Transformer
16
9.80
5,4G
21M
SwitchAll
4
9.81
2.4G
5.6M
C4
47M
SwitchAll
2
22.09
202M
0.8M
Transformer
10
22.63
453M
3.5M
262M
SwitchAll
4
16.45
2.4G
5.6M
Transformer
16
16.58
5.4G
21M
peS2o
47M
SwitchAll
2
12.56
202M
0.8M
Transformer
12.83
453M
3.5M
262M
Transformer
16 9.78
5.4G
21M
SwitchAll
9.86
2.4G
5.6M
4 Analysis
In order to see how the network uses the attention heads, trained a small, 6-layer, 8-head Transformer on ListOps [33, 34]. The reason for this choice is that small, algorithmic tasks tend be more interpretable compared to language modeling tasks. We also train parameter- mattmee,e-matchd,, 2-head

---

Table 4: Performance of SeadchHead trained C4 dataset, compared dense Transformer baseline with matched number of parameters.
Model
#total params
ppl 4
Lambada
BLiMP f CBT 1
SwitchHead
47M
22.53
20.4%
75.7%
Transformer
47M
22.
20.4%
73.6%
SwitchHead MAC-
63M
21.18
23.5%
77.1%
SwitchHead Shared selection
47M
22.81
20.0%
74.6%
SwitchHead
262M
16.23
29.4%
79.6%
83.3%
Transformer
262M
16.28
28.2%
76.1%
83.6%
SwitchHead MAC-matched
376M
15.43
30.2%
79.4%
84.2%
SwitchHead Shared selection
262M
16.49
28.6%
79.4%
82.7%
Table Real-world resource usage of our method. The numbers shown below are for training time for the whole pipeline, including the feedforward layers. It can be that SwitchHead in the current implementation reduces both the runtime and the memory usage by a factor of 1.4-1.5.
Size
Model
ms/iteration
Rel. iter. time
RAM/GPU Rel. Mem.
#GPUs
GPU type
Transformer
473ms/iter
1.0
20.5G
1.0
47M
SwitchHead
342ms/iter
0.72
13.5G
0.65
RTX 3090
MoA
41 12ms/iter
0.87
15.3G
0.75
Transformer
670m//tter
1.0
20.5G
1.0
262M
SwitchHead MoA
442ms/iter
0.65
12.5G
0.61
V100
85 851m//iter
1.27
16.4G
0.80
SwitchHead model. Both models achieve around 95% accuracy on held-out IID validation set, in contrast to the dense 2-head model, which saturates around 80%. Note that ListOps is classification task and does not use autoregressive masking.
We visualize the maximum of attention heads for each layer, both for the standard Transformer (Fig. 2a) and SwitchHead (Fig. 2b). The attention maps qualitatively similar. Due to different initialization and learning dynamics, thus the overlap between the two models would not be perfect. Complete attention map visualizations can be found in Fig. 4 and 3 in the appendix.
In addition, we anlyze individual attention heads for SwitchHead. We find that often possible to interpret the selection weights: synthetic tasks, the output experts specialize according different operations, while the input ones distinguish numbers and closed parentheses. The attention map itself appears to distribute information about contiguous chunks of numbers (see Fig. 5 in the appendix).
Attention maps of the language models are more difficult to interpret. However, we visualize the attention maps of the 47M parameter Transformer XL and the SwitchHead model from Tab. 2. We find them to be qualitatively similar. We also identified induction heads [35] in both models, some examples shown for SwitchHead in Fig. 6a and for Transformer in Fig. 6b in the appendix. Other typical vertical line-lined attention patterns are shown Fig. and 6d.
Related Work
The method most closely related to ours is MoA [18], which introduces MoE style attention. It defines each attention head as an expert but shares the key and value projections between them. Unlike in our Switchhead, each of the selected experts requires separate attention matrix, which significantly increases its memory usage. Due to the use competitive softmax-bast activation function in the selection network, it requires complex regularization to prevent expert collapse [17]. In the original formulation, the number of active heads is high. Our experiments also confirm that MoA needs many attention heads to match the performance of the dense baseline (see Sec. 3.2), and it is only possible to do so with a significantly higher resource budget than our method.
Nguyen et al. [36] analyze the attention matrices, and they conclude that they are usually low rank. Motivated by this, the authors construct few (e.g., 2) "global attention matrices", and they compute each local matrix for specific heads by weighted average of those. However, they average the logits,

---

0.8
0.8
00.6
0.6
0.4
0.4
0.2
0.2
(a) Transformer, Layer
(b) SwitchHead Layer 3
Figure 2: An attention map of the (a) standard Transformer and (b) SwitchHead. The maximum of all heads in the given layer are shown.
not the final matrix, so each individual head-specific matrix has to be computed. This means that in the best case, they only save half of the computation associated with the attention matrix because the readout (Eq. 3) is still needed. For the same reason, memory savings are also low.
Peng et al. [19] propose to reweight the contribution of each head by gating function. However, they only reduce the number of total attention heads by one, presumably compensate for the parameters used by the selection logic. Their goal was not to reduce resource usage but to have better predictive performance, which they achieve. They use a softmax-based competitive selection mechanism. To avoid collapse, the gating function is trained only in some steps.
More broadly, there have been several works on MoE to accelerate language models. Shazeer et al. [11] introduce sparsely-gated mixture of experts. Fedus et al. [37] introduce Mixture of Experts in Transformers. Lepikhin al. [13] train a MoE-based LLM, and Clark et [15] analyze the scaling laws of MoE models. Lewis et al. [12] introduce an alternative method for preventing collapse. How- ever, none of these methods focus on the important, parameter-matched setting. Csordás et [17] introduce the conpeomeeeee activation based MoE method, o-MoE, which was shown to be suc- cessful in such setting, but the authors only focused on accelerating the MLPs and the attention.
Multi-Query attention [38] uses single key and value projection that is shared between the heads while using multiple queries. Our findings show that such configuration is suboptimal: using multiple output and value projections is the most important choice in our model design.
Dao al. [39] provides hardware-awar CUDA implementation of the entire attention layer, which avoids storing the attention matrix. By saving memory bandwidth in this way, they achieve significant wall clock time speedup, despite that the attention matrix should be recomputed in the backward pass. This orthogonal to our method and they can combined for further acceleration.
Limitations
Our models are modest in size compared to the current state-of-art LLMs. However, training such models is estimated to cost millions of dollars, which we cannot afford. Instead, we aim show the versatility of our model by choosing a diverse set of datasets, including Enwik 8, Wikitext 103, C4 and peS2o, and different positional encodings, such as Seansffomrr----stteee relative positional encoding and RoPE. We also demonstrate the competitiveness of our models in zero-shot downstream tasks. We believe that the evidence we provided enough for research group with a larger amount of resources at their disposal to verify our findings in sfate-of-the-at model.

---

The Triton kernel that we used is currently around 60% of the speed of a single dense matrix multiplica- tion of the size of single expert with cuBLAS. Even this, we showed wall-clock time speedup. We es- timate that 80-90% should be achievable with more optimal kernel. Model- paral-la training requires the implementation of aad-balancing system that can dynamically move experts between GPUs.
Conclusion
On wide range of language modeling datasets with different model sizes, novel Mixture- of-Experts (MoE) based attention method called SwitchHead achieves performance of parameter- matched dense counterparts, with only fraction of the computational cost and memory usage. SwitchHead drastically reduces the number of attention matrices that have to be computed, by using MoE for the value and output projections. Our method is stable and does not need additional regular- ization to prevent degenerate solutions (a well-known practical issue in many existing MoE models). Our method can also be successfully combined with MoE MLP layers, to obtain "SwitchAll" where every layer of the Transformer is MoE-based, achieving a huge reduction in resource requirements.
Acknowledgements
This research was partially funded by ERC Advanced grant no: 742870, project AlgoRNN, and by Swiss National Science Foundation grant no: 100021_192356, project NEUSYM. We are thankful for hardware donations from NVIDIA and IBM. The resources used for this work were partially provided by Swiss National Supercomputing Centre (CSCS) projects d123 and s1205.
References
[1] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models unsupervised multitask learners. 2019.
[2] Tom B Brown et al. Language models are few-shot learners. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
[3] OpenAl. Chatgpt. https: htttpp:///peaai.... httppeaaai..comboooooooetteeeeee chatppt, 2022.
[4] OpenAl. GPT-4 technical report. Preprint 2rXiv:2303.08774 2023.
[5] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee. Yuanzhi Li, Scott M. Lundberg, Harsha Nori, Hamid Palangi, Marco Túlio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Early experiments with GPT-4. Preprint arXiv:2303.1 12712, 2023.
[6] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Advances in Neural Information Processing Systems (NIPS), pages 5998-69988999 6008, Long Beach, CA, USA, December 2017.
[7] Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to recurrent nets. Neural Computation, 4(1):331-139 1992.
[8] Georgi Gerganov. llama.cpp. https: / //:::://ittttt //tgtttttb.///.///////// com/ /gggggggnnooollaaaaa.mo cpp, 2023.
[9] John B. Hampshire II and Alexander H. Waibel. The meta-pi network: connectionist rapid adaptation for high-performance multi-speaker phoneme recognition. In Proc. IEEE Int. Conf. on Acoustics, Speech and Signal Processing (ICASSP), pages 165-168, Albuquerque, New Mexico, USA. April 1990.
[10] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Compututaion, 3(1):79-87, 1991.
[11] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expper layer. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
[12] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. In Proc. Int. Conf. on Machine Learning (ICMI., volume 139, pages 3265-6274, Virtual only, July 2021.
10

---

[13]
Dmitry Lepikhin, HyoukJoong Lee. Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with condi- tional computation and automatic sharding. In Int. Conf. on Learning Representations (ICIR), Virtual only, May 2021.
[14] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research
( 23)):2222-2255 2022.
[15] Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Pagannni, Jordan Hoffmann, Bogdan Damoc, Blake A. Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack W. Rae, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan. Unified scaling laws for routed language models. Preprint arXiv:2202.00 01 169, 2022.
[16] Zewen Chi, Li Dong. Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj. Xia Song. Xinn-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, Louisiana, USA, December 2022.
[17] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. Approximating two-layer feedforward networks for efficient transformers. In Findings the Association for Computational Linguisties: EMNLP 2023, November 2023.
[18] Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), pages 41 4150-4162, Abu Dhabi, United Arab Emirates, December 2022.
[19] Hao Peng, Roy Schwartz, Dianqi Li, and Noah A. Smith. A mixture of h heads is better than h heads. In Proc. Association for Computational Linguistics (ACL), pages only, July 2020.
Virtual
[20] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research (JMLR), 21:140:1-140:67 2020.
[21] Marcus Hutter. The human knowledge compression prize. http::///pize...tterree 2006.
[22] Luca Soldaini and Kyle Lo. peS2o (Pretraining Efficiently on S2ORC) Dataset. Technical report, Allen Institute for AI, 2023. attps://githut com/allenai/pes20
[23] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In Ini. Conf. on Learning Representutions (ICLR), Toulon, France, April 2017.
[24] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring broad discourse context. In Proc. Association for Compu- tational Linguistics (ACL), Berlin, Germany, August 2016.
[25] Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu Wang, and Samuel R. Bowman. Blimp: The benchmark of linguistic minimal pairs for english. Transactions of the Association for Computational Linguistics (TACL), 8:377-392, 2020.
[26] Felix Hill, Antoine Bordes, Sumit Chopra, and Jason Weston. The goldilocks principle: Reading children's books with explicit memory representations. In Int. Conf. on Learning Representations (ICLR), San Juan, Puerto Rico, May 2016.
[27] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models. Preprint arXiv:2302. 13971, 2023.
(28) Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced transformer with rotary position embedding. Preprint arXiv:2104.09864, 2021.
11

---

[29] Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. Proc. Association for Computational Linguistics (ACL), pages 1715-1725, Berlin, Germany, August 2016.
[30] Mike Schuster and Kaisuke Nakajima. Japanese and korean voice search. In Proc. IEEE Int. Conf. on Acoustics, Speech and Signal Processing (ICASSP), pages 5149-5152, Kyoto, Japan, March 2012.
[31] Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proc. Conf. Empirical Methods in Natural Language Processing (EMNLP), pages 66-71, Brussels, Belgium, October 2018.
[32] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoe Le, and Ruslan Salakhutdinov. Transformer-XL. Attentive language models beyond fixed-length context. In Proc. Association for Computational Linguistics (ACL), pages 2978-2988 Florence, Italy, 2019.
[33] Nikita Nangia and Samuel R. Bowman. ListOps: A diagnostic dataset for latent tree learning. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL-HLT), pages 92 22-99, New Orleans, USA, June 2018
[34] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. In Int. Conf. on Learning Representations (ICLR), Virtual only, April 2022.
[35] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Sriffeld-Doddie Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-contex: learning and induction heads. Transformer Circuits Thread, 2022. https:/raansfommerriccuusssppppp httpa22222i.oottttttaaaannnanaia weads/index.htm..
[36] Tan Nguyen, Tam Nguyen, Hai Do, Khai Nguyen, Vishwanath Saragadam, Minh Pham, Duy Khuong Nguyen. Nhat Ho, and Stanley Osher. Improving transformer with an ad- mixture of attention heads. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, November 2022.
[37] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Preprint arXiv:2/01.0396/ 2021.
[38] Noam Shazeer. Fast transformer decoding: One write-head all you need. Preprint 19 2019.
[39] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, Louisiana, USA, December 2022.
[40] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Int. Conf. on Learning Representations (ICLR), San Diego, CA, USA, May 2015.
12

---

A Appendix
A.1 Comment on Flash Attention
The resource reductions from Flash Attention might be, in many cases, larger than those from our method alone. However, Flash Attention depends on GPU-specific memory bandwidth/comput trade-offs, which might not be available on all hardware, especially edge devices. SwatchHead and FlashAttention can also be combined for further speedups. We demonstrated the viability of this setup in our RoPE experiments. Additionally, certain architectures, such shared- laaee transformers, might require a drastic increase in the number of heads, which Fltentntention alone might not be able to do.
A.2 Resource Usage of Different Methods
In this section, we discuss the compute and memory usage of different attention variants. We will define the compute in terms of the number of aacliplyy-ccumullte operations (MACs, also used by Zhang et al. [18]), which is arguably better defined than FLOPs (e.g., does one step of the matrix multiplication count as 1 FL.OP or 2? Do we include the softmax?). All calculations will be presented for a single attention layer for single sequence, and they are presented this way in all our tables. Both the memory and compute requirements scale linearly with both the batch size and the number
of layers.
Consider sequence of inputs of length T' with representation size dmadel- Let dnead be the width of the key, query and value projections used for the attention layer. For Transformer XL-style attention, let the size of the context be CT, where C' - 1 is the number of past chunks included in the context of the current attention step. We can divide the computation into two major parts: calculating the projections, which do not involve the attention map, and calculating the attention map and projecting the sequence of values using it.
First, consider the case of the standard Transformer XL. [32]. Here, from the input x E TR. descid we calculate the K", Q". V4 R ddhex using projection matrices of shape Rdddd Xchaed The output after the attention is projected in a similar manner (Eq. 3). Thus, the projections take a total of 4T'd chen MACs per head. For backpropagation, we have to store all the intermediate results. This takes T Tdnead numbers of Kah Q" and V*h Also, the projected values should be stored. They have an identical shape, therefore, the total memory used by projections is 4T dhene numbers per head. Now consider the resource usage related to the attention matrix. It involves calculating the product of Q² Q²K²T which takes CT2 MACs (multiplication by C' is needed because the shape of K" and V for Transformer XL CT' drea)). The projection of the values with the attention matrix A A is similar. For the memory usage, the attention needs CT2 numbers, but it needs to be stored both before and after the activation function. In addition, calculating the projection of the position encodings is necessary. This depends on the implementation. but in our case, it involves a matrix multiplication, and the total amount of computation 2cheaa and it needs Tch numbers of storage. Thus the resource requirements are:
A XXL MAC Theeads (4T'dheaddmodde + 22TTT + 2CT (headdmodd)))
(11)
XXL mem
Theads (4T + 2CT2 + 2CT dhead))
(12)
The resource usage of SwitchHead is different. First, the number of heads Thends is significantly reduced, but dhead is typically larger. Additionally, there are i experts active at the same time. Here, we only consider the case where the value and outputs are experts, but Q" and K are not (this version performs the best; see Sec. 3.1). Then, we have two projections that are identical with that of Transformer XL, and two MoE-based projections. These use T Themot MACs to calculate the projection and another Tkdsead to calculate their weighted average. With a smart kernel implementa- tion, memory usage is not affected by k. thus the formula remains the same as Eq. 12 (note, however, that Theads and daeud are very different in practice). The compute requirement can be calculated as:
A-Switchllead = Thheals (2T Cheehe + 27 2Tkhhhaa (dmoste + 1) 2CCT2 + 2CT dmenel) (13) MAC
Additionally, the expert selection logic needs minimal additional resources, which can be ignored. Note that the comparison between the MACs of the standard (Eq. 11) and SwitchHead (Eq. 13) depends on the exact values of the hyper-parameters. However, as we Ill see in See. 3, in our typical
13

---

Table Performance of SwitchHead with E experts and Thhaads heads. Different projections are either experts or fixed for the given head. Columns V, K, Q. and O show whether the given and Tineads are shown.
projection is an expert. Parameter-matched baseline with Theads Models sorted by perplexity. 47M parameters models on Wikitext 103.
Model
Theads
Q
Perplexity 4
Switchllead
2
12.27
Switchllead
2
N
N
Y
12.30
Transformer
10
12.31
SwitchHead
2
N
Y
N
Y
12.36
SwitchHead
2
12.37
SwitchHead
12.42
SwitchHead
2
Y
N
N
N
12.45
SwitcHHead
2
N
N
Y
Y
12.45
SwitcHHead
2
Y
N
12.51
SwitchHead
2
12.57
SwitchHeaa
2 N
Y Y
12.59
SwitchHead
2
Y
Y
Y
12.61
SwitchHead
2
Y
Y
N
12.69
Transformer
2
12.74
Switchllead
2
Switchleaa
2
N
Y
12.79
SwitchHead
2
N
Y
Y
N
12.90
configurations, SwitchHead provides good predictive performance with significantly lower Theads compared to the standard Transformer, resulting in reduced resource usage in the end.
The resource requirements of MoA [19] are very similar to those of Transformer XL, except that uses single shared key and value projection for each head.
M
(2nheads + chett + 2Thends CT + 2CT 2CT Cmedd
(14)
N MoAA mem
(2Mheads + 222 + 2mheads CT + 2CT dhead
(15)
A.3 The Importance of Different Projections
order to analyze which projections are the most important to be mixture-of-experts, we exhaustively tried all combinations. We analyze our 47M parameter models WikiText 103 dataset. We show the results in Tab. 6. We also include parameter-matched baseline with two heads, which serves a lower bound for the performance. We found that the value and output projections most important, and having key and query projections hurts the performance. This is possible because perform all our experiments in parameter-matcher setting. Allocating parameters to these projections uses the budget that can be otherwise spent on other parts of the network. In our preliminary experiments, found that, allowing the parameter budget increase, experts always help.
the
A.4
RoPE Positional Encodings
All our experiments the main paper have used a Transformer XL model. Thus, it remains unclear whether SwitchHead is specific to this model can be also used with other attention methods. As an alternative, we consider RoPE positional encodings [28] without the XL cache (thus, the attention matrices square). This is the standard setup used by modern language models, such as all versions of Llama [27]. We tested these models in Wikitext 103 and C4. The results are shown in Tab. 7, and zero-shot performance downstream tasks in Tab. 8. This shows that SwitchHead also performs well in the standard setup and not tied to Transformer XI..
A.5 Hyperparameters
We train all our models with Adam optimizer [40], with batch size of 64, a learning rate of 0.00025, and gradient clipping with a maximum norm of 6. Large models (> 200K parameters) use learning
14

---

Table Perplexity of SwitchHead compared to dense baseline, using RoPE positional encoding and no XL cache. Memory usage is specified in number of floats. Models sorted by perplexity.
Dataset
#total params
Model
Thheads
ppl +
MACs
Memory
Wikitext 103
45M
SwitchHead
12.75
285.6M
1.3M
Transformer
10
12.78
560.9M
6.1M
Transformer
12.96 2
560.9M
1.9M
244M
SwitchHead
10.00
4.2G
18.4M
Transformer
16
10.17
6,4G
37.7M
Transformer
10.26
6.4G
8.4M
C4
45M
SwitchHead
2
23.69
285.6M
1.3M
Transformer
10
23.79
560.9M
6.1M
244M
SwitehHead
4
16.41
4.2G
18.4M
Transformer
16
16.35
6.4G
37.7M
Table 8: Zero-shot task performance of SwitchHead using RoPE positional encodings and no XL cache, trained on C4 dataset, compared to dense Transformer baseline with matched number of parameters.
Model
#total params
ppl
Lambada
BLiMP f CBT
SwitchHead
45M
23.69
20.9%
77.3%
Transformer
45M
23.76
20.3%
73.8%
SwitchHeaa MAC-matched
54M
22.18
22.6%
77.4%
SwitchHead Shared selection
45M
23.63
20.3%
76.0%
SwitchHead
243M
16.41
30.5%
79.9%
83.8%
Transformer
243M
16.35
29.8%
76.1%
83.9%
SwitchHead MAC-matched
314M
15.63
30.5%
80.5%
84.6%
SwitcHHead Shared selection
243M
16.59
28.1%
79.1%
83.7%
rate warm-up of 4k steps. All models, except the SwitchAll model, dropout on the MLP layers, 0.1 for the small models and 0.2 for the large ones. Detailed hyperparameters are shown in the Tab. 9. o-MoE related hyperparameters for the SwitchAll models are identical to those of Csordás et al. [17]. For Transformer XL models, we always use single additional chunk of context, both in training and validation time. dhead and dg are derived in a systematic way, see Sec. for more details.
A.6 Note the Parameter Count of the SwitchAll
can be seen Tab. 3 that the parameter count of the SwitchAll models is often less than that of their dense counterparts. The reason is that normally compensate for the final difference in the number of parameters by increasing du (see Sec. for details of the parameter matching). However, that only be done in very coarse-grained way with o-MoR: the size of all experts must be increased at once, and the CUDA kernel supports only sizes of multiple of 4. Therefore, increasing the size of the experts would add too many parameters and the model would outgrow the baseline. For this reason, we simply keep the hyperparameters for Csordás et al. [17] and combine them with our SwitchHead configuration from Tab. 2.
A.7 Visalizing all Attention Heads
As discussed Sec. 4, we analyze the attention maps SwitchHead and compare them with the dense models. We show all the attention maps of the models trained ListOps in Fig. and Fig. We show individual heads of SwitchHead, including the expert selection scores in Fig. 5. Some selected attention maps of our 47M parameter models on Wikitext 103 are shown in Fig.
A.8 Compute Requirements
We report the compute used for our experiments, including the GPU type, count (the number of GPUs used per experiment, and not the total in the machine), and the runtime in "hhmm" format
15

---

Table 9: Hyperparameters used for our models.
Model
Dataset
Pheads
#params
dhed
dg
E
k
T
Taayees
SwitchHead
2
47M
76
2080
5
256
16
0.1
Transformer
C4
10
47M
41
2053
256
16
0.1
Transformer
47M
205
2053
256
16
0.1
SwitchHead
262M
112
4188
4
512
18
0.25
Transformer
16
262M
64
4110
512
18
0.25
Transformer
262M
256
4110
512
18
0.25
SwitchHead
47M
76
2080
5
256
16
Transformer
Wikitext 103
10
47M
41
2053
256
16
0.1
Transformer
2
47M
205
2053
256
16
0.1
SwitchHead
262M
132
4147
8
512
18
0.25
Transformer
Wikitext 103
16
262M
64
4110
512
18
0.25
Transformer
262M
512
4110
512
18
0.25
SwitchHeaa
47M
76
2080
5
256
16
0.1
Transformer
PeS2o
10
2053
256
16
Transformer
47M
205
2053
256
16
0.1
SwitchHead
262M
4188
4
512
18
0.25
Transformer
peS2o
16
262M
64
4110
512
18
0.25
Transformer
262M
256
4110
512
18
0.25
SwitchHead
2
41M
112
2088
4
512
12
0.25
Transformer
Enwik8
41M
64
2053
512
12
0.25
Transformer
41M
256
2053
512
12
0.25
Switchllead (RoPE)
45M
64
2092
5
512
16
0.1
Wikitext 103
Transformer (RoPE)
10
45M
2053
512
16
SwitchHead (RoPE) Transformer (RoPE)
243M
100
4136
4
1024
18
0.25
Wikitexi 103
16
244M
64
4110
1024
18
0.25
SwitchAll
Wikitext 103
47M
76
1648
5
256
16
0.25
SwitchAll
Wikitext 103
259M
112
4096
4
512
18
0.25
SwitchAll
C4
47M
76
1648
5
256
16
0.25
SwitchAll
C4
4
259M
112
4096
4
512
18
0.25
SwitchAll
peS2o
47M
76
1648
5
256
16
0.25
SwitchAll
peS2o
4
259M
112
4096
4
512
18
0.25
Tab. 10. We report the total number of CPUs (NCpu) and RAM because they are shared between concurrent runs. Note that most of the experiments were done prior to the much faster. Triton-basc kernel implementation. Because of this, the runtimes appear longer for SwitcHead compared the baseline. For timing benchmarks with our new kernel, see Tab. 5.
Note that we only report the resources used for the paper here. We estimate that the total of the failed experiments and preliminary runs is around 10 times higher than this.
16

---

01
(a) Layer
(b) Layer 2
(c) Layer 3
025
425
0.30
53350
5.030
535
0.0
0.5
533
2,000
0.3
203
50C
(d) Layer
(c) Layer 5
Layer 6
Figure 3: The maximum of all attention maps for Head model ListOps.
17

---

04
002
(a) Layer
(b) Layer 2
(c) Layer 3
03
0.40
0.5
029
020
0.5
-61
(d) Layer
(c) Layer 5
Layer 6
Figure 4: The maximum of all attention maps for standard Transformer model ListOps.
18

---

339
(a) Layer 1, head
(b) Layer 1, head
(c) Layer 2, head
-55
(d) Layer 2, head
(e) Layer 3, head
(f) Layer head
$.50
gg) Layer 4, head
(h) Layer 4, head
(i) Layer head
400
429
0455
045
( Layer 5, head
(k) Layer 6, head
()) Layer head 2
Figure 5: Details for individual heads of the SwitchHead model on ListOps. On the left side of each attention plot, the selection of the output projection expert is shown. Similarly, at the bottom, the selection the value projection selection is visible. In the selection maps, dark blue always corresponds to while white O. The adaptive scale shown the right of the attention map is for the map only.
19

---

(a) SwitchHead Layer 12. Induction head.
(b) Transformer XL Layer 10. Induction head.
(c) SwitchHead Layer 9. Stripe pattern.
(d) Transformer XL Layer 8. Stripe pattem..
Figure 6: Induction head copying the rare name "Homarus" in (a) SwitchHead and (b) Transformer XL baseline. The attention matrix is square because it is the first chunk of the sequence, without any extra context. Typical vertical line pattern in (c) SwitchHead and (b) Transformer XL. baseline.
20

---

Table 10: Training hardware information for the experiments reported in the paper
Model
#params
Dataset
G
GPU Type
NaPu
Ncrv
RAM
Duration
SwitchAll
259M
C4
V100-32GB-LS
40
503G
24:06
SwitchAll
259M
PeS2o
V100-32GB-ILS
40
503G
30:00
SwitchAll
259M
Wikitext 103
RTX 4090
24
251G
22:58
SwitchAll
47M
C4
2
RTX 3090
24
220G
22:14
SwitchAll
47M
PeS2o
RIX 3090
24
220G
22:49
SwitchAll
47M
Wikitext 103
RTX 3090
251G
6:03
SwitchHead
243M
Wikitext 103
V100-32GB
40
503G
147:09
SwitchHead
262M
C4
V100-32GB-LA
8
40
503G
26:38
SwitchHead
262M
PeS
V100-32GB-LS
40
503G
27:43
SwitchHead
262M
Wikitext 103
V100-32GB
40
503G
31:42
SwitchHeead
41M
Enwik8
V100-32GB
13:45
SwitchHeea
45M
Wikitext 103
RTX 3090
24
251G
17:28
SwitchHead
47M
C4
V100-32GB
40
503G
15:36
SwitchHead
47M PeS2o
V100-32GB
40
503G
16:17
SwitchHead
47M Wikitext 103
RTX 3090
251G
13:09
Transformer
262M
C4
V100-32GB
40
503G
11:55
Transformer
262M
C4
16
V100-32GB-LS
8
40
503G
20:21
Transformer
262M
P2S2o
4
V100-32GB
40
503G
17:08
Transformer
262M
S2o
16
V100-32GB-LS
40
503G
25:56
Transformer
262M
Wikitext 103
P100-16GB
12
62G
0:00
Transformer
262M
Wikitext 103
16
A100-80GB
2
64
503G
31:51
Transformer
41M
Enwik8
RTX 3090
24
220G
15:38
Transformer
41M
Enwik8
V100-32GB-LS
40
503G
16:04
Transformer
47M
C4
V100-32GB
40
503G
10:29
Transformer
47M
C4
10
V100-32CH 32GB
40
503G
16:57
Transformer
47M
PeS2o
V100-32GB
40
503G
11:07
Transformer
47M
S2o
10
V100-32GB
40
503G
17:55
Transformer
47M
Wikitext 103
V100-32GB
40
503G
10:06
Transformer
47M
Wikitext 103
10
'100-32GB
503G
18:51
Transformer (RoPE)
244M
Wikitext 103
16
RTX 3090
24
251G
30:30
Transformer (RoPE)
45M
Wikitext 103
10
V100-32GB
40
503G
15:30
21

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We summarized the motivation, method, and main findings in these sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in Sec. 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
22

---

Justification:
This is an empirical paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: show all the hyperparameter configurations Appendix A.5, and we provide the code for our experiments.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
23

---

Answer: [Yes]
Justification: We provide the code for our experiments. automatically downloads all the data that it needs.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: guwgge/des///// /Ct///ooodiessseeeeeeee ssionnoiicyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Our methodology is explained in Sec. 3, and the full table of hyperparameters is presented in Appendix A.5.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Our experiments involve large models that are very expensive to train, and we do not have sufficient compute resources run multiple seeds them.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the type of hardware used for our main experiments in Appendix A.8.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:/ /ttp:::///uupppeee cctpublic//tttoe /ptpbbii//ttccctsiieleieeees
Answer: [Yes]
Justification: We read the Ethics guidelines, and to the best of our knowledge, we are complying with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is a special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
We consider our paper to be a foundational research paper without direct consequences.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
25

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The models in this paper are small by modern standards and we do not release pre-trained weights.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Justification: Our code is under MIT license and the paper is CC-BY 4.0. To the best of our knowledge, we always credit the reused code we reuse any.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
26

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide the source code and instructions on how to run it.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not work with human subjects,
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
We do work with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
27

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "ULC8Ig0tND", "reviewer": "Reviewer_58yB", "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.", "questions": "See above.", "limitations": "Limitations have been discussed adequately.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and well motivated. \nWhile the components of the proposed method are not particularly novel, they are combined in a simple yet novel way. \nI 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. \n- I could not find the sequence length used in your experiments, which prevents me from computing the number of tokens used during training. \n- 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. \n- What is explaining the differences in the transformers' PPLs in table 4 compared to the PPLs in table 2?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> I could not find the sequence length used in your experiments, which prevents me from computing the number of tokens used during training.\n\nThe sequence length is T in Tab. 6 in the appendix.\n\n> 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.\n\nThe 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.\n\n> What is explaining the differences in the transformers' PPLs in table 4 compared to the PPLs in table 2?\n\nThey correspond to runs with different seeds, for historical reasons. We will update Tab. 2 in the next version to avoid this confusion.\n\nWe believe our response above resolves the concerns that the reviewer has raised. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q9GazvrpTP", "reviewer": "Reviewer_HTZS", "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.", "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.\n- 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?\n- 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.\n- 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.\n- 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?\n- 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.\n\n[1] Nallapati, Ramesh, et al. \"Abstractive text summarization using sequence-to-sequence rnns and beyond.\" arXiv preprint arXiv:1602.06023 (2016).\n\n[2] Kwiatkowski, Tom, et al. \"Natural questions: a benchmark for question answering research.\" Transactions of the Association for Computational Linguistics 7 (2019): 453-466.", "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.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed MoE method applied to self-attention layer reduces computational and memory requirements while preserving language modeling performance of dense baselines.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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.\n- The paper mentions \"preliminary experiments\" multiple times without proper citation or section indicators, which may confuse readers and reduce the reproducibility of the work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their insightful review. Please find our responses as follows:\n\n> … 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]? \n\n> 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]? \n\nWe 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.\n\n> 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.\n\n> 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?\n\nWhile 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).\nWe 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.\n\n> 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?\n\n\nThis 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.\n\n> ... 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?\n\n> … a more granular analysis of the trade-offs between model size, computational resources, and performance could enhance the paper's contribution to the field.\n\nIn 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.\n\nGenerally, 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).\n\n> 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?\n\nWe 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. \n\nWe hope our response above brings clarifications to all the concerns that the reviewer has raised. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EBei9KgTln", "reviewer": "Reviewer_YarL", "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.", "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.\n\n2. 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.\n\n3. 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.", "limitations": "The authors have discussed the limitations of their study.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> 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.\n\nWe 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.\n\nThe 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.\n\n> 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.\n\nThe 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.\n\nWhile 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.\n\nPlease note that we also provide additional results in a MAC-matched setting in Tab 4.\n\n> 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.\n\nWe would thank the reviewer for pointing this out. It is the total number of parameters. We will clarify this in the final version.\n\n> 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.\n\nThank 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.\n\n> 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.\n\nThis 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.\n\n| Size | Model | ms/iteration | Rel. iter. time | RAM/GPU | Rel. Mem. | #GPUs | GPU type |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| 47M | Transformer | 473ms/iter | 1.0 | 20.5G | 1.0 | 1 | RTX 3090 |\n| | SwitchHead | 342ms/iter | 0.72 | 13.5G | 0.65 |  | |\n| | MoA | 412ms/iter | 0.87 | 15.3G | 0.75 | ||\n262M | Transformer| 670ms/iter | 1.0 | 20.5G | 1.0 | 8 | V100|\n| | SwitchHead |  442ms/iter | 0.65 | 12.5G | 0.61 |  | |\n| | MoA |  851ms/iter | 1.27 | 16.4G | 0.80 |  | |\n\nWe 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.\n\n[1] Csordás et al. EMNLP 2023 (Findings). Approximating Two-Layer Feedforward Networks for Efficient Transformers.\n\n[2] Nguyen et al. Arxiv 2024. Sigmoid Gating is More Sample Efficient than Softmax Gating in Mixture of Experts.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UZOZTJRFLk", "reviewer": "Reviewer_Q7cF", "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.", "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]\n- 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.\n- 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?\n- l. 439: does the 4 in the calculation of the total MACs for the projections already include the output projection?\n- 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.\n- 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\"? \n- ll. 225 - 227: the visualization in Fig. 2 is only for 1 layer, and not for all layers as stated in the text\n- ll. 282 - 287: does not exactly fit into the \"Limitations\" section, since, according to the authors (ll. 266-269) FlashAttention is orthogonal to SwitchHead\n- 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?\n \nOverall solid paper, I am happy to revisit my score upon the rebuttal.\n\n[1] Shazeer, Noam. \"Fast transformer decoding: One write-head is all you need.\" arXiv preprint arXiv:1911.02150 (2019).\n\n[2] Page-Caccia, Lucas, et al. \"Multi-head adapter routing for cross-task generalization.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[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).\n\n[4] Ostapenko, Oleksiy, et al. \"Towards modular llms by building and reusing a library of loras.\" arXiv preprint arXiv:2405.11157 (2024).", "limitations": "Limitations are addressed in section 6.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall, this is a well-crafted paper that presents a straightforward yet effective method for efficiently integrating MoE layers into self-attention blocks.\n\nOriginality: 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.\n\nQuality: good. The claims are well supported by evidences. Experiments are well designed and seem to be sufficient overall.\n\nClarity: the paper is well organized. In terms of writing, some formulations require further clarifications (see weaknesses/questions)\n\nSignificance: 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> similar techniques for transforming attention heads into MoE have been proposed in … authors can consider mentioning these works [2,3,4 inter alia]\n\nWe would like to thank the reviewer for pointing this out. We will include the discussion of these papers in the final version.\n\n> 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. …\n\nWe 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.\n\n> the explanation is somewhat confusing, why is it the case that in the worst case all possible source projections have to be computed?\n\nConsidering 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.\n\n> does the 4 in the calculation of the total MACs for the projections already include the output projection?\n\nYes, it does. We will clarify this in the final version.\n\n> 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].\n\nWe would like to thank the reviewer for pointing this out. We will include a discussion of [1] in the final version.\n\n> 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.\n\nWe 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.\n\n> 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\"?\n\nThe 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.\n\nPlease note that we also provide additional results in a MAC-matched setting in Tab 4.\n\n> does not exactly fit into the \"Limitations\" section, since, according to the authors (ll. 266-269) FlashAttention is orthogonal to SwitchHead\n\nWe agree with the reviewer that this paragraph does not fit here. We’ll move this to Appendix A. Thank you for pointing this out.\n\n> the visualization in Fig. 2 is only for 1 layer, and not for all layers as stated in the text\nWe would like to thank the reviewer for pointing this out. We will fix this in the final version.\n\n> 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?\n\nThe 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.\n\n\nWe believe our response above resolves all the concerns that the reviewer has raised. \nThese questions are extremely valuable for us and will enable us to improve our paper.\nIf the reviewer finds our response useful, please consider increasing the score. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: SwitchHead demonstrates impressive performance improvements in language modeling tasks, showcasing the potential of MoE in self-attention layers. Despite some minor limitations, the reviewers generally agree on the technical soundness and significance of the work.
Given the overall positive reviews, I recommend accepting this paper. Addressing the weaknesses identified during the review process, such as expanding the evaluation to additional tasks and providing more quantitative analysis of attention maps, would further enhance the paper's contribution.
Decision: Accept (poster)