Research map / Architectures and efficiency / Mixture of experts

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

NeurIPS 2024 · Architectures and efficiency › Mixture of experts · approach: Sparse MoE as regularization and attention · 8 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes SwitchHead, a mixture-of-experts attention layer that uses expert value and output projections (with shared keys/queries) to cut the number of computed attention matrices, and SwitchAll, which combines it with sigma-MoE feedforward layers, trained from scratch on language modeling.

Problem

Self-attention accounts for much of Transformer compute and memory, but prior MoE-attention methods (e.g., MoA) fail to match parameter-matched dense baselines or give only modest resource savings.

Approach

Make value and output projections independent MoEs with sigmoid non-competitive routing, keep a single head-specific key/query, so that far fewer attention matrices need computing; the MoE projections are applied outside the attention core, enabling up to 8x fewer attention matrices.

Key results

• On C4 at 262M parameters, SwitchHead matches dense perplexity (16.23 vs 16.28) with about 44% compute and 27% memory.
• On Wikitext 103 at 47M, SwitchHead reaches perplexity 12.27 vs 12.31 for the 10-head Transformer at 170M MACs vs 453M.
• Wall-clock training is about 1.5x faster with 61-67% of baseline memory in the reported setup.

Measured effects reported in the paper

Surprising findings

• Keeping key and query projections as fixed (non-MoE) is best; MoE on value and output projections is what matters, and multi-query-style shared K/V was found suboptimal.
• MoA needs many heads to match the dense baseline and its wall-clock cost can exceed the dense model (1.27x iteration time at 262M).

Limitations

• Models are modest in size (47M and 262M) relative to state-of-the-art LLMs.
• The Triton kernel runs at roughly 60% of cuBLAS dense matmul speed; model-parallel training needs load-balancing across GPUs.
• Perplexity gains in main tables are byproducts of imperfect parameter matching rather than the goal.

Experimental setup

Datasets: C4 Enwik8 peS2o Wikitext 103 ListOps

Benchmarks: Wikitext 103 C4 peS2o Enwik8 Lambada BLiMP Children's Book Test

Metrics: perplexity bits per character accuracy MACs memory (floats) wall-clock time per iteration

Baselines: Transformer XL dense baseline (matched heads) MoA Transformer with RoPE sigma-MoE

Builds on: sigma-MoE Mixture of Attention Heads (MoA) Transformer-XL RoPE

Closest papers by meaning

• UMoE: Unifying Attention and FFN with Shared Experts (NeurIPS 2025 · similarity 0.90)
• Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.89)
• MoEUT: Mixture-of-Experts Universal Transformers (NeurIPS 2024 · similarity 0.88)
• Sparse Layers are Critical to Scaling Looped Language Models (NeurIPS 2026 · similarity 0.88)
• Transformer Layers as Painters (AAAI 2025 · similarity 0.87)
• DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging (NeurIPS 2024 · similarity 0.87)