SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention | Summary

This article explains the key points of SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention.

• 2023-12-13 (arXiv), 38th Conference on Neural Information Processing Systems (NeurIPS 2024)
• Csordás, Róbert, Piękos, Piotr, Irie, Kazuki, Schmidhuber, Jürgen.
• Stanford University, Stanford, CA, USA, AI Initiative, KAUST, Thuwal, Saudi Arabia, Center for Brain Science, Harvard University, Cambridge, MA, USA, The Swiss AI Lab IDSIA, USI & SUPSI, Lugano, Switzerland
• Paper
• Github

Read this article in Korean

Summary

• Transformer language models spend substantial computation and memory on attention, the operation that combines information across tokens, while most methods that activate only selected model components target other layers. On C4, a language-modeling dataset, the 262M-parameter SwitchHead model achieves perplexity, a lower-is-better prediction metric, of 16.23 versus 16.28 for a dense Transformer with the same parameter count, using approximately 44% of its attention-layer computation and 27% of its attention-layer intermediate memory.
• SwitchHead applies Mixture of Experts (MoE), which selects a subset of learned transformations for each token, to value projections, which prepare information for attention to combine, and output projections, which map the combined information back to the model representation. Within each head, a parallel attention component, experts share key and query projections, which determine token-to-token attention weights; this allows multiple experts to use one attention matrix and reduces the number of matrices by up to 8 times in the reported comparisons. Training requires no additional regularization to keep expert selection stable.
• Across the complete training pipeline, SwitchHead reduces iteration time from 473 to 342 ms at 47M parameters and from 670 to 442 ms at 262M parameters on the respective matched hardware. The evidence covers several datasets and ways of encoding token position, but does not establish performance at the scale of current large language models or measure generation-time acceleration.

1 Introduction

Existing MoE approaches primarily accelerate feedforward blocks, the token-wise transformations within a Transformer, leaving attention costs largely unchanged. SwitchHead places expert selection outside the attention computation so that fewer attention matrices must be computed and stored. It also combines with σ-MoE, a feedforward MoE using independent sigmoid scores, which map each expert’s selection score into the 0–1 range, to form SwitchAll.

2 Method

The method separates routing, the selection of projection experts for each token, from the computation of token-to-token attention. Its final design retains ordinary key and query projections and activates selected experts only for values and outputs.

2.1 Background

Multi-head self-attention (MHA) computes several token-to-token weighting matrices in parallel: it first projects inputs into keys, queries, and values, then applies attention to values and projects the result to the output. The output can equivalently be written as a sum of head-specific contributions, making each head’s computation independent.

• For sequence length T, head count n_heads, and head width d_head, attention-matrix construction and value readout require O(n_heads d_head T²) multiply-accumulate operations (MACs), operations that multiply two numbers and add the result to an accumulator. The paper uses their count as a hardware-independent computation measure.
• Training stores O(n_heads T²) attention entries and O(n_heads T d_head) projection intermediates, so attention becomes increasingly costly as sequences grow.

2.2 From Dense to SwitchHead Attention Layer

Selecting complete heads using only the receiving token is difficult to accelerate because attention depends on both sending and receiving tokens; all source projections may still be required. SwitchHead instead selects source-side value experts and destination-side output experts independently, using sigmoid scores followed by top-k selection, which retains the k highest-scoring experts per token.

• Each head has E experts but computes only one attention matrix. Selected value projections are combined before attention, and selected output projections are combined after the attention readout.
• The final variant shares one head-specific key projection and one query projection across its experts. Sigmoid scores are independent rather than normalized across experts, and the reported training requires no additional regularization to prevent expert collapse, in which routing concentrates on too few experts.

The diagram places expert selection on the value and output paths around a shared attention matrix. Multiple projection choices therefore do not require a separate attention matrix for every expert.

3 Experiments

The main experiments match total parameter counts and configure SwitchHead to approximate the dense baseline’s perplexity while maximizing resource savings. They use Transformer XL, which incorporates an additional previous chunk of context, with total context twice the active chunk length; separate experiments use RoPE, rotary position embeddings that encode token position through rotations of query and key representations, without the XL cache of previous context.

• Configuration sets n_heads × E equal to the dense baseline’s head count, starts with two heads and two active experts, and increases the active expert count if needed. If four active experts remain insufficient, it switches to four heads with two active experts; head widths and remaining feedforward capacity are adjusted to match the parameter budget.
• All models train for 100k batches. Except for character-level Enwik8, data use an 8k-token SentencePiece vocabulary, obtained with a tokenizer that divides text into subword units.
• For C4 and peS2o, training uses the first 10⁵ × T × N_batch tokens, where T is the active chunk length and N_batch is the batch size. The authors treat small perplexity gains in the main resource-saving comparisons as a consequence of imperfect performance matching.

3.1 Which Projections Require an MoE?

An exhaustive comparison of expert versus fixed key, query, value, and output projections identifies value and output experts as the strongest configuration under the fixed parameter budget. At 47M parameters on Wikitext 103, this variant achieves perplexity 12.27, compared with 12.31 for the 10-head dense baseline and 12.57 when all four projections use experts.

• Output experts alone achieve 12.30, showing that output routing is particularly consequential in this ablation, an experiment that varies model components. Adding key or query experts consumes parameters that can otherwise support more useful components.

3.2 Comparison with MoA

Mixture of Attention Heads (MoA), the closest comparison method, shares key and value projections but computes a separate attention matrix for each selected query/output expert. SwitchHead combines experts around a smaller set of attention matrices and uses independent sigmoid routing instead of MoA’s competitive routing, which normalizes scores across experts and requires three regularizers.

• At 47M parameters on Wikitext 103, SwitchHead achieves perplexity 12.27 with 170.4M MACs and 0.8M stored floats; the four-head MoA variant achieves 12.60 with 223.5M MACs and 1.3M floats. These resource estimates describe one attention layer processing one sequence.
• At 262M parameters, eight-head MoA achieves slightly lower perplexity than SwitchHead, 9.50 versus 9.55, but uses 2.9G versus 2.0G MACs and 9.9M versus 2.9M stored floats. These are Table 1 comparisons, distinct from the Wikitext 103 results in Table 2.

At 47M parameters, SwitchHead has lower perplexity than all listed MoA variants while using fewer resources than the four-head MoA model. At 262M parameters, eight-head MoA slightly improves perplexity but requires more MACs and intermediate memory, showing the trade-off between predictive quality and resource use.

3.3 Performance on Different Datasets

Across C4, Wikitext 103, peS2o, and Enwik8, SwitchHead closely matches the many-head dense baseline and consistently improves on a parameter-matched dense model with the same reduced head count. For example, the 262M C4 models achieve perplexities of 16.23 for four-head SwitchHead, 16.28 for the 16-head baseline, and 17.09 for the four-head dense baseline.

• On Wikitext 103 at 262M parameters, two-head SwitchHead achieves 9.77 versus 9.80 for the 16-head baseline, using 2.0G versus 5.4G MACs and 2.9M versus 21M stored floats.
• Small performance losses remain: on peS2o at 262M parameters, perplexity is 9.86 versus 9.78. On Enwik8, the 41M-parameter models both achieve 1.10 bits per character, which measures average character-prediction uncertainty.
• Zero-shot evaluation, which tests pretrained models without task-specific training, includes Lambada for contextual word prediction, BLiMP for English linguistic minimal pairs, and Children’s Book Test (CBT) for prediction from book context. At 262M parameters, SwitchHead improves Lambada from 28.2% to 29.4% and BLiMP from 76.1% to 79.6%, while CBT decreases from 83.6% to 83.3%.

The reduced-head dense controls perform worse even when their parameter counts match, supporting the contribution of expert projections beyond head reduction alone. The 262M C4 comparison supplies the approximately 44% computation and 27% intermediate-memory figures, which apply to attention-layer accounting.

3.4 SwitchAll

SwitchAll combines SwitchHead attention with σ-MoE feedforward layers using the existing feedforward configuration and the attention settings from the main experiments. Table 3 supports compatibility, with lower perplexity on all three 47M comparisons but mixed results at the larger scale.

• At 47M parameters, SwitchAll achieves 12.17 versus 12.32 on Wikitext 103, 22.09 versus 22.63 on C4, and 12.56 versus 12.83 on peS2o.
• The larger SwitchAll models contain 259M parameters against the nominal 262M baselines. Table 3 reports 9.81 versus 9.80 on Wikitext 103, 16.45 versus 16.58 on C4, and 9.86 versus 9.78 on peS2o; these entries conflict with the surrounding prose’s identification of C4 as the exception to improvement.

All three small SwitchAll models improve perplexity over their listed dense baselines. At the larger scale, C4 improves while Wikitext 103 and peS2o are slightly worse; the actual larger SwitchAll parameter count is 259M.

3.5 MAC-Matched Setup

The MAC-matched setup spends SwitchHead’s computation savings on additional heads and wider projections, producing models with more parameters than the dense baseline. For Transformer XL, the small configuration increases head width from 76 to 112 and head count from two to three; the large configuration increases width from 112 to 168 and head count from four to six.

• On C4, the 63M model achieves perplexity 21.18 versus 22.71 for the 47M dense baseline, and the 376M model achieves 15.43 versus 16.28 for the 262M baseline.
• The 376M model reaches 30.2% Lambada, 79.4% BLiMP, and 84.2% CBT versus 28.2%, 76.1%, and 83.6% for the dense baseline. Its BLiMP score is slightly below the parameter-matched SwitchHead model’s 79.6%.

3.6 Shared Selection

Sharing expert selection between the source and destination sides reduces sorting and top-k operations but generally weakens predictive performance. At 262M parameters on C4, shared selection changes perplexity from 16.23 to 16.49, Lambada from 29.4% to 28.6%, BLiMP from 79.6% to 79.4%, and CBT from 83.3% to 82.7%.

The table separates parameter-matched, MAC-matched, and shared-selection configurations. The models with more parameters improve several results under matched computation, while the parameter-matched SwitchHead model shows a slight CBT decline.

3.7 Wall-Clock Time and Memory Usage Estimation

The complete training-pipeline benchmark uses identical hardware and attention-core implementations within each size comparison, changing the attention projections. It includes feedforward layers, the optimizer, and synchronization of gradients across devices, so its savings are smaller than the attention-layer analytical estimates.

• On one RTX 3090, the 47M model reduces iteration time from 473 to 342 ms and per-device memory from 20.5G to 13.5G. On eight V100 devices, the 262M model reduces iteration time from 670 to 442 ms and memory from 20.5G to 12.5G.
• The MoA comparison uses four active heads at 47M parameters and eight at 262M, selecting the fastest model that matches the dense baseline or the best-performing model when none matches it. MoA takes 412 ms and 15.3G at the small scale and 851 ms and 16.4G at the large scale; SwitchHead is faster and uses less memory in both comparisons.
• These results establish training acceleration for the measured implementations. MAC counts alone do not determine wall-clock speed on other hardware.

These complete-pipeline measurements demonstrate training savings beyond analytical operation counts. SwitchHead reduces iteration time and per-device memory on one RTX 3090 at 47M parameters and eight V100 devices at 262M parameters, and outperforms the listed MoA implementations on both measures.

4 Analysis

On ListOps, a synthetic task involving nested list operations, a six-layer, eight-head dense Transformer and a parameter-matched two-head SwitchHead both reach around 95% accuracy on held-out examples drawn from the training distribution, while a dense two-head model saturates around 80%. This classification experiment uses no autoregressive mask, the restriction that prevents tokens from attending to future tokens during language modeling.

• Layerwise attention maxima, formed by taking the largest weight across heads for each token pair, look qualitatively similar across the dense and SwitchHead models. Individual routing plots suggest output experts specialize by operation, while value experts distinguish numbers and closing parentheses.
• Language-model visualizations also show similar patterns, including induction heads that copy information from earlier matching context, illustrated here with the rare name “Homarus,” and vertical attention stripes. These examples provide qualitative evidence rather than a comprehensive measurement of equivalent internal behavior.

The layer-3 maps retain similar prominent attention structures despite SwitchHead using fewer heads. Each panel takes the maximum across heads, so the visualization supports qualitative similarity without establishing identical head-level representations.

5 Related Work

The related-work discussion distinguishes reducing attention matrices from reweighting heads or combining attention logits, the scores before attention normalization. Those alternatives can retain substantial matrix construction and readout costs, whereas SwitchHead shares an attention matrix across multiple projection experts.

• Multi-Query attention shares key and value projections across heads. SwitchHead’s ablations favor diversity in value and output projections under the tested parameter budgets.
• FlashAttention, a hardware-aware exact-attention implementation that avoids storing the full attention matrix, addresses memory access and is compatible with SwitchHead’s architectural changes. The appendix reports their combination in the RoPE experiments.

6 Limitations

The tested models are modest in size, and the paper does not verify its findings in state-of-the-art large language models. Its expert kernel, a specialized computation implemented in the Triton programming system, runs at around 60% of the speed of a corresponding single-expert dense matrix multiplication using cuBLAS, NVIDIA’s matrix-operation library; the proposed 80–90% target is an estimate.

• Model-parallel training, which distributes model components across devices, would require a load-balancing system capable of dynamically moving experts between devices.

7 Conclusion

The paper concludes that routing value and output projections allows SwitchHead to approach parameter-matched dense-model performance with fewer attention matrices and lower resource requirements, without additional regularization to prevent expert collapse. The experiments also support combining the method with feedforward MoE layers, although SwitchAll’s predictive gains vary by dataset and scale.

Appendix

• A.1 A Comment on Flash Attention: The authors acknowledge that FlashAttention alone may provide larger savings in many settings. They argue that the methods address complementary costs and report that combining them is viable in their RoPE experiments; possible advantages on edge devices, which may lack the relevant hardware trade-offs, remain conditional.
• A.2 Resource Usage of Different Methods: The resource formulas and the MAC/intermediate-memory columns describe one attention layer processing one sequence, with linear scaling in batch size and layer count. For context length CT, where C − 1 previous chunks are included, attention-matrix storage scales with n_heads CT² and attention-core computation with n_heads CT² d_head; SwitchHead reduces head count while adding selected-expert projection costs. Its memory formula assumes a kernel that avoids storing separate intermediates for every active expert, and the computation estimate neglects the small expert-selection overhead.
• A.3 The Importance of Different Projections: Table 6 compares all expert/fixed projection combinations for 47M-parameter Wikitext 103 models with two heads and E = 5 experts. Value/output experts achieve the best reported perplexity, 12.27, while key/query experts without value/output experts reach 12.90, supporting the final allocation of the fixed parameter budget.
• A.4 RoPE Positional Encodings: Without the XL cache, SwitchHead remains competitive on Wikitext 103 and C4. Table 7 reports Wikitext 103 perplexity of 12.75 versus 12.78 at the smaller scale and 10.00 versus 10.17 at the larger scale, while large-scale C4 perplexity is slightly worse, 16.41 versus 16.35. Table 8 reports BLiMP accuracy of 79.9% versus 76.1% for the larger parameter-matched comparison.
• A.5 Hyperparameters: Training uses Adam, a gradient-based optimizer, batch size 64, learning rate 0.00025, and gradient clipping, which limits gradient magnitude, with maximum norm κ specified in Table 9. The text assigns 4k-step learning-rate warm-up to “large models (> 200K parameters),” preserving the PDF’s stated threshold; feedforward dropout, which randomly suppresses activations during training, is 0.1 for small models and 0.2 for large models, except in SwitchAll. Transformer XL uses one additional context chunk during training and validation.
• A.6 A Note on the Parameter Count of the SwitchAll: Exact parameter matching is constrained because changing σ-MoE expert width changes all experts simultaneously, and the kernel supports widths only in multiples of 4. The authors therefore retain the existing feedforward MoE settings, leaving the larger SwitchAll models at 259M parameters.
• A.7 Visalizing all Attention Heads: Figures 3 and 4 show layerwise maximum attention maps on ListOps, Figure 5 adds per-head value/output selection scores, and Figure 6 presents selected Wikitext 103 attention patterns. The selection-score plots use a fixed 0–1 scale, whereas the attention plots have adaptive scales, so color intensity cannot be compared directly across attention panels.
• A.8 Compute Requirements: Table 10 lists training hardware, device counts, shared host resources, and durations in hh:mm. Most runs preceded the faster Triton kernel, so their durations should not be used as matched speed benchmarks. The authors estimate that failed and preliminary experiments cost around 10 times the resources reported for the paper’s included experiments.

Brief Thoughts

The projection ablation, reduced-head dense controls, and complete-pipeline speed measurements jointly support preserving projection diversity while computing fewer attention matrices. Practical assessment should distinguish analytical attention-layer savings from measured training savings and account for the untested larger scales and small regressions on peS2o and CBT. Table 3 shows mixed SwitchAll results at the larger scale, and Table 4 documents a slight CBT decline for parameter-matched SwitchHead. Both qualify the paper’s broader performance claims. Code is provided at https://github.com/robertcsordas/switchhead.