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

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Bianchessi, Arthur S., Aguirre, Yasmin C., Barros, Rodrigo C., Kupssinskü, Lucas S.
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866910199677190144
author Bianchessi, Arthur S.
Aguirre, Yasmin C.
Barros, Rodrigo C.
Kupssinskü, Lucas S.
author_facet Bianchessi, Arthur S.
Aguirre, Yasmin C.
Barros, Rodrigo C.
Kupssinskü, Lucas S.
contents 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\times$ 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_22842
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Bayesian Attention Mechanism: A Probabilistic Framework for Positional Encoding and Context Length Extrapolation
Bianchessi, Arthur S.
Aguirre, Yasmin C.
Barros, Rodrigo C.
Kupssinskü, Lucas S.
Computation and Language
Machine Learning
I.2.6; I.2.7
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\times$ 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.
title Bayesian Attention Mechanism: A Probabilistic Framework for Positional Encoding and Context Length Extrapolation
topic Computation and Language
Machine Learning
I.2.6; I.2.7
url https://arxiv.org/abs/2505.22842