# Bayesian Attention Mechanism: A Probabilistic Framework for Positional Encoding and Context Length Extrapolation

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 24 references that are in the map.
- First posted on arXiv 28 May 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks](https://mlsanity.com/paper/2412.15204/) (2024)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Landmark Attention: Random-Access Infinite Context Length for Transformers](https://mlsanity.com/paper/2305.16300/) (2023)
- [Transformer Language Models without Positional Encodings Still Learn Positional Information](https://mlsanity.com/paper/2203.16634/) (2022)
- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [Advancing Transformer Architecture in Long-Context Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2311.12351/) (2023)
- [Can Perplexity Reflect Large Language Model's Ability in Long Text Understanding?](https://mlsanity.com/paper/2405.06105/) (2024)
- [Scalable-Softmax Is Superior for Attention](https://mlsanity.com/paper/2501.19399/) (2025)
- [Attention: Marginal Probability is All You Need?](https://mlsanity.com/paper/2304.04556/) (2023)
- [Computation Mechanism Behind LLM Position Generalization](https://mlsanity.com/paper/2503.13305/) (2025)

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

- [Robust Filter Attention: Self-Attention as Precision-Weighted State Estimation](https://mlsanity.com/paper/2509.04154/) (2025)
- [The Transformer as a Polar State Estimator](https://mlsanity.com/paper/2605.11007/) (2026)
- [A Probabilistic Interpretation of KV Cache Eviction](https://mlsanity.com/paper/2608.28293/) (2026)

## Abstract (arXiv)

Transformer-based language models rely on positional encoding (PE) to handle token order and support context length extrapolation. However, existing PE methods lack theoretical clarity and rely on limited evaluation metrics to substantiate their extrapolation claims. We propose the Bayesian Attention Mechanism (BAM), a theoretical framework that formulates positional encoding as a prior within a probabilistic model. BAM unifies existing methods (e.g., NoPE and ALiBi) and motivates a new Generalized Gaussian positional prior that substantially improves long-context generalization. Empirically, BAM enables accurate information retrieval at 500× the training context length, outperforming previous state-of-the-art context length generalization in long context retrieval accuracy while maintaining comparable perplexity and introducing minimal additional parameters.