# Mixture of Attention Heads: Selecting Attention Heads Per Token

arXiv: https://arxiv.org/abs/2210.05144 · HTML: https://mlsanity.com/paper/2210.05144/

## Key facts

- Cited by 60 papers in the map (data through 14 Oct 2026).
- Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts.
- Builds on 26 of its 36 references that are in the map.
- First posted on arXiv 11 Oct 2022.

## Placed in

- Mixture-of-experts / [Expert routing strategies](https://mlsanity.com/moe/expert-routing-strategies/)

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Pay Less Attention with Lightweight and Dynamic Convolutions](https://mlsanity.com/paper/1901.10430/) (2019)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Scaling Neural Machine Translation](https://mlsanity.com/paper/1806.00187/) (2018)
- [DeepNet: Scaling Transformers to 1,000 Layers](https://mlsanity.com/paper/2203.00555/) (2022)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [On the Representation Collapse of Sparse Mixture of Experts](https://mlsanity.com/paper/2204.09179/) (2022)
- [Do Attention Heads in BERT Track Syntactic Dependencies?](https://mlsanity.com/paper/1911.12246/) (2019)
- [Very Deep Transformers for Neural Machine Translation](https://mlsanity.com/paper/2008.07772/) (2020)
- [A Mixture of h-1 Heads is Better than h Heads](https://mlsanity.com/paper/2005.06537/) (2020)
- [Tricks for Training Sparse Translation Models](https://mlsanity.com/paper/2110.08246/) (2021)

## Built on by (most cited first, 30 of 60)

- [Boosting Continual Learning of Vision-Language Models via Mixture-of-Experts Adapters](https://mlsanity.com/paper/2403.11549/) (2024)
- [MoEUT: Mixture-of-Experts Universal Transformers](https://mlsanity.com/paper/2405.16039/) (2024)
- [ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing](https://mlsanity.com/paper/2412.14711/) (2024)
- [A Closer Look into Mixture-of-Experts in Large Language Models](https://mlsanity.com/paper/2406.18219/) (2024)
- [A Survey on Inference Optimization Techniques for Mixture of Experts Models](https://mlsanity.com/paper/2412.14219/) (2024)
- [MoH: Multi-Head Attention as Mixture-of-Head Attention](https://mlsanity.com/paper/2410.11842/) (2024)
- [Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models](https://mlsanity.com/paper/2404.05567/) (2024)
- [CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models](https://mlsanity.com/paper/2404.08763/) (2024)
- [SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention](https://mlsanity.com/paper/2312.07987/) (2023)
- [Sparse Universal Transformer](https://mlsanity.com/paper/2310.07096/) (2023)
- [Model Swarms: Collaborative Search to Adapt LLM Experts via Swarm Intelligence](https://mlsanity.com/paper/2410.11163/) (2024)
- [Layerwise Recurrent Router for Mixture-of-Experts](https://mlsanity.com/paper/2408.06793/) (2024)
- [ModuleFormer: Modularity Emerges from Mixture-of-Experts](https://mlsanity.com/paper/2306.04640/) (2023)
- [Unlocking Emergent Modularity in Large Language Models](https://mlsanity.com/paper/2310.10908/) (2023)
- [Selective Task Group Updates for Multi-Task Optimization](https://mlsanity.com/paper/2502.11986/) (2025)
- [BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts](https://mlsanity.com/paper/2408.08274/) (2024)
- [RingMoE: Mixture-of-Modality-Experts Multi-Modal Foundation Models for Universal Remote Sensing Image Interpretation](https://mlsanity.com/paper/2504.03166/) (2025)
- [Mod-Squad: Designing Mixture of Experts As Modular Multi-Task Learners](https://mlsanity.com/paper/2212.08066/) (2022)
- [Scattered Mixture-of-Experts Implementation](https://mlsanity.com/paper/2403.08245/) (2024)
- [Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routing](https://mlsanity.com/paper/2505.00315/) (2025)
- [MergeME: Model Merging Techniques for Homogeneous and Heterogeneous MoEs](https://mlsanity.com/paper/2502.00997/) (2025)
- [Experts Weights Averaging: A New General Training Scheme for Vision Transformers](https://mlsanity.com/paper/2308.06093/) (2023)
- [MoQAE: Mixed-Precision Quantization for Long-Context LLM Inference via Mixture of Quantization-Aware Experts](https://mlsanity.com/paper/2506.07533/) (2025)
- [Conditional computation in neural networks: principles and research trends](https://mlsanity.com/paper/2403.07965/) (2024)
- [Improving Transformers with Dynamically Composable Multi-Head Attention](https://mlsanity.com/paper/2405.08553/) (2024)
- [Learning to Skip the Middle Layers of Transformers](https://mlsanity.com/paper/2506.21103/) (2025)
- [Not All Models Suit Expert Offloading: On Local Routing Consistency of Mixture-of-Expert Models](https://mlsanity.com/paper/2505.16056/) (2025)
- [MISA: Mixture of Indexer Sparse Attention for Long-Context LLM Inference](https://mlsanity.com/paper/2605.07363/) (2026)
- [Adaptive Computation Modules: Granular Conditional Computation For Efficient Inference](https://mlsanity.com/paper/2312.10193/) (2023)
- [BTS: Harmonizing Specialized Experts into a Generalist LLM](https://mlsanity.com/paper/2502.00075/) (2025)

## Abstract (arXiv)

Mixture-of-Experts (MoE) networks have been proposed as an efficient way to scale up model capacity and implement conditional computing. However, the study of MoE components mostly focused on the feedforward layer in Transformer architecture. This paper proposes the Mixture of Attention Heads (MoA), a new architecture that combines multi-head attention with the MoE mechanism. MoA includes a set of attention heads that each has its own set of parameters. Given an input, a router dynamically selects a subset of k attention heads per token. This conditional computation schema allows MoA to achieve stronger performance than the standard multi-head attention layer. Furthermore, the sparsely gated MoA can easily scale up the number of attention heads and the number of parameters while preserving computational efficiency. In addition to the performance improvements, MoA also automatically differentiates heads' utilities, providing a new perspective to discuss the model's interpretability. We conducted experiments on several important tasks, including Machine Translation and Masked Language Modeling. Experiments have shown promising results on several tasks against strong baselines that involve large and very deep models.