# Adaptive Head Budgeting for Efficient Multi-Head Attention

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 22 references that are in the map.
- First posted on arXiv 24 Apr 2026.

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

- [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)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [A large annotated corpus for learning natural language inference](https://mlsanity.com/paper/1508.05326/) (2015)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [DynaBERT: Dynamic BERT with Adaptive Width and Depth](https://mlsanity.com/paper/2004.04037/) (2020)
- [Compressing Large-Scale Transformer-Based Models: A Case Study on BERT](https://mlsanity.com/paper/2002.11985/) (2020)
- [Learning to Skip the Middle Layers of Transformers](https://mlsanity.com/paper/2506.21103/) (2025)
- [Pruning Attention Heads of Transformer Models Using A* Search: A Novel Approach to Compress Big NLP Architectures](https://mlsanity.com/paper/2110.15225/) (2021)
- [CATP: Cross-Attention Token Pruning for Accuracy Preserved Multimodal Model Inference](https://mlsanity.com/paper/2404.08567/) (2024)
- [Hybrid Dynamic Pruning: A Pathway to Efficient Transformer Inference](https://mlsanity.com/paper/2407.12893/) (2024)

## Built on by (most cited first, 1 of 1)

- [Budgeted Attention Allocation: Cost-Conditioned Compute Control for Efficient Transformers](https://mlsanity.com/paper/2605.05697/) (2026)

## Abstract (arXiv)

Multi-head attention enables Transformers to capture diverse representations, but all attention heads are typically activated for every input, regardless of task complexity. For coarse-grained tasks such as text classification, where relevant information is often global, this fixed allocation can introduce unnecessary computation. We propose BudgetFormer, a Transformer architecture that dynamically allocates attention heads on a per-input basis. The model learns both a head budget and a relevance distribution to select the most informative heads. To support effective head selection, we introduce a training strategy that balances exploration and exploitation. Experiments on text classification tasks show that BudgetFormer reduces FLOPs and memory usage while matching or surpassing the performance of standard multi-head attention. These results highlight adaptive head allocation as an effective approach to improving Transformer efficiency and performance.