mlsanityData to 14 Oct 2026

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

arXiv 2605.28384 ↗First posted 27 May 2026cs.LGcited by 0

Not on the map. Why?

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

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

25 of 26 references.

References · 25

  1. 2017Attention Is All You Needcited by 60,361
  2. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,718
  3. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  4. 2020Longformer: The Long-Document Transformercited by 3,413
  5. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,073
  6. 2020Reformer: The Efficient Transformercited by 1,668
  7. 2020Rethinking Attention with Performerscited by 1,601
  8. 2019Generating Long Sequences with Sparse Transformerscited by 1,546
  9. 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,336
  10. 2023RWKV: Reinventing RNNs for the Transformer Eracited by 827
  11. 2023Retentive Network: A Successor to Transformer for Large Language Modelscited by 613
  12. 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 519
13 more
  1. 2020Efficient Content-Based Sparse Attention with Routing Transformerscited by 471
  2. 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 407
  3. 2025Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Freecited by 283
  4. 2024Mixture-of-Depths: Dynamically allocating compute in transformer-based language modelscited by 258
  5. 2023Zoology: Measuring and Improving Recall in Efficient Language Modelscited by 204
  6. 2025MoBA: Mixture of Block Attention for Long-Context LLMscited by 198
  7. 2021MoEfication: Transformer Feed-forward Layers are Mixtures of Expertscited by 152
  8. 2018Implicit Reparameterization Gradientscited by 85
  9. 2025The emergence of sparse attention: impact of data distribution and benefits of repetitioncited by 25
  10. 2025The Bayesian Geometry of Transformer Attentioncited by 13
  11. 2024Learning Linear Attention in Polynomial Timecited by 5
  12. 2025Bayes optimal learning of attention-indexed modelscited by 3
  13. 2026Variational Routing: A Scalable Bayesian Framework for Calibrated Mixture-of-Experts Transformerscited by 3

Built on by

Not cited yet.