# Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 26 references that are in the map.
- First posted on arXiv 27 May 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [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)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (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)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free](https://mlsanity.com/paper/2505.06708/) (2025)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [Zoology: Measuring and Improving Recall in Efficient Language Models](https://mlsanity.com/paper/2312.04927/) (2023)
- [MoBA: Mixture of Block Attention for Long-Context LLMs](https://mlsanity.com/paper/2502.13189/) (2025)
- [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021)
- [Implicit Reparameterization Gradients](https://mlsanity.com/paper/1805.08498/) (2018)
- [The emergence of sparse attention: impact of data distribution and benefits of repetition](https://mlsanity.com/paper/2505.17863/) (2025)
- [The Bayesian Geometry of Transformer Attention](https://mlsanity.com/paper/2512.22471/) (2025)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)
- [Bayes optimal learning of attention-indexed models](https://mlsanity.com/paper/2506.01582/) (2025)
- [Variational Routing: A Scalable Bayesian Framework for Calibrated Mixture-of-Experts Transformers](https://mlsanity.com/paper/2603.09453/) (2026)

## Abstract (arXiv)

Standard transformer architectures apply a single attention mechanism uniformly across all tokens and sequence positions, irrespective of local context or computational budget. We propose Meta-Attention, a framework that dynamically routes each token to the most appropriate attention strategy -- full softmax attention, linear (kernel) attention, or sliding-window local attention -- via a Bayesian Meta-Controller. Unlike prior routing approaches that use deterministic or prior-free learned routing, the Meta-Controller treats per-token mechanism selection as posterior inference under a compute-aware Dirichlet prior: routing weights are the output of an amortised variational posterior q(alpha | x_t; phi) trained with an Evidence Lower Bound (ELBO) objective that jointly encodes task performance and attention-mechanism cost. This design produces principled routing uncertainty estimates that govern the soft-to-hard routing transition, mitigates routing collapse without ad hoc load-balancing losses, and yields better compute-performance trade-offs than deterministic or prior-free learned routing at negligible overhead. Phase 1 empirical results on a Tiny LM benchmark confirm core predictions: the Bayesian controller's learned routing distribution implies a projected normalised FLOP cost of 25.1% under hard routing, vs. 59.3% for the prior-free baseline (-34.2 pp), and reduces routing entropy from 55.8% to 43.3% (-12.5 pp), demonstrating that the Dirichlet prior prevents routing collapse while the non-Bayesian model defaults to full attention. We present the Bayesian architecture, ELBO training objective, and a Phase 1 PyTorch prototype validating forward-pass correctness, posterior diversity, and a controlled ablation against a prior-free baseline. Code available at: https://github.com/KFEAL/meta-attention