Bayesian Attention Mechanism: A Probabilistic Framework for Positional Encoding and Context Length Extrapolation
Authors: Arthur S. Bianchessi
Yasmin C. Aguirre
Rodrigo C. Barros
Lucas S. Kupssinskü
Venue: ICLR 2026 Poster (2026)
Keywords: Positional Encoding
Large Language Model
Transformer
Long Context
Attention Mechanism
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce BAM, a Bayesian theoretical framework for positional encoding (PE) that interprets existing methods as prior distributions within Bayesian attention, enabling significant improvements in transformers' long-context performance.

Abstract: 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 by more than $25\times$ in retrieval accuracy while maintaining comparable perplexity and introducing minimal additional parameters.

BAYESIAN ATTENTION MECHANISM:
A PROBABILISTIC FRAMEWORK FOR POSITIONAL ENCODING AND CONTEXT LENGTH EXTRAPOLATION
Arthur S. Bianchessi*
Yasmin C. Aguirre Rodrigo C. Barros Lucas S. Kupssinskii'
MALTA Machine Learning Theory and Applications Lab, PUCRS Porto Alegre, Brazil Kumumi Institute, Brazil
(arthur. biihhurs.iicccssi, yasmin, cardozo) Gedu.pucrs.br
boddrigo.barroe lucas. kupssinsku epucrs.br
rodrigo.barros@ www.bbbrrrroukuunnnn
ABSTRACT
Transformer- language models rely positional encoding (PE) to handle token order and support context length extrapolation. However, existing PE meth- ods lack theoretical clarity and rely limited evaluation metrics to substantiate their extrapolation claims. We propose the Bayesian Attention Mechanism (BAM), theoretical framework that formulates positional encoding as a prior within probabilistic model. BAM unifies existing methods (c.g.. NoPE and ALiBi) and motivates a new Generalized Gaussian positional prior that substantially improves long-context generalization. Empirically, BAM enables accurate information re- trieval at 500 the training context length, outperforming previous state-of- -or context length generalization in long context retrieval accuracy while maintaining comparable perplexity and introducing minimal additional parameters.
INTRODUCTION
Transformer-base neural models currently dominate language modeling (LM) due to their superior ability capture variable-length token dependencies via self-attention. Nevertheless, transformers inherently lack positional information, requiring the incorporation of Positional Encoding (PE). PE is vital, particularly for enabling LMs trained on shorter contexts to generalize to significantly longer sequences during inference- desirable capability known as context length extrapolation. However, the precise impact of PE on extrapolation remains poorly understood (Kazemnejad et 2023).
Several PE methods have been proposed to facilitate context length extrapolation, including Sinusoidal embeddings |Vaswanil (017)) RoPE et et al.] ALiBi (Press et al. (022), and even the omission of positional encoding entirely (NoPE) (Kazemnejad et al.) (023). Despite reported empirical successes extrapolation, two critical issues persist: (i) many existing PE techniques are empirically motivated with limited theoretical foundations, and thus their behavior is not well-understood (Liu ct al., 2024; Huang et al. 2023); ; (ii) cvaluation methods rely heavily on perplexity, which may inadequately reflect true extrapolation capability, as LMs can achieve low perplexity simply through localized attention patterns, as demonstrated in sliding-window evaluations (Huang et al. 2023).
To address these issues, introduce the Bayesian Attention Mechanism (BAM), theoretical framework that reframes self-attention as an expectation of values computed under a joint probabilistic model of content and position of tokens. Within BAM, PE naturally emerges as a prior distribution over token positions, clarifying the theoretical basis of existing techniques. Notably, we illustrate how NoPE and ALiBi correspond explicitly to Uniform and Laplace positional priors, respectively.
Leveraging this robust theoretical foundation, we propose new positional encoding strategy utilizing a Generalized Gaussian prior. Our approact introduces fewer than 1,000 additional parameters yet delivers substantially improved extrapolation performance, demonstrated clearly in retrieval-based
Code available http:///gtthha.o////ttttteaaaaaa

---

tasks and traditional perplexity evaluations. Thus, BAM serves both unified theoretical framework for analyzing PE schemes and as a practical method for enhancing long-context attention.
2 BAYESIAN ATTENTION MECHANISM
In this section, we motivate the perspective of framing attention as a Bayesian mechanism, supporting both token content and position information. We show that PE strategies can be seen as priors of Bayesian attention mechanism, hereby called BAM.
2.1 BAYESIAN ATTENTION AND THE JOINT PROBABILITY DISTRIBUTION Pij
Definition 1. For fixed query vector q R1 and key-value matrices K, V R Bayesian Attention Mechanism computes self-attention as an expectation over its values:
exp(score(q: kj))
self-attention qu K, V) exxxp(score( k.)) V
EyPoy = E
Definition 1 states that the attention mechanism can be expressed as expectation over values of the hth query, where Pij is the probability token j E [1, L] when attended by token This definition is consonant with the self-attention mechanism defined by Vaswani (2017) and with prior attempts to frame self-attention as an expectation over values (Singh & Buckley] 2023). where the scoring function is the scaled dot product and a (score(g, Kj) computes Pij Vj. The term Pij is usually called the attention weight, however we frame it as joint probability over content and positions.
Definition 2. In Bayesian Attention, Pıj is joint probability over the content of token j (fcont and its position relative to query qi (Spos).
Pij P(fcoot (o: k2) 1ppss (i, x p(gpos (i.j))
Definition 2 allows us to interpret Piy as a joint probability distribution dependent on both token content and position. Together, Definitions 1 and 2 frame positional encoding as probability distribution over the positions of the tokens within the context. Note that this definition, particularly in p(Scom (q; ký) 1pos (i, encompasses a dependency of content on position. This dependency is trivially modeled by a scalar Z when the scoring function is additive, as detailed below.
When framing PE as that probability distribution over tokens in context, probability distributions that explain positional encoding strategies such as NoPE (Kazemnejad et al.
derive parametrized
2023 and ALiBi (Press et al. 2022) and propose novel PE strategies with known behaviors.
Theorem 1. If the scoring function of the attention mechanism is additive, i.e., of the form Joom k,) + Opos( j), then Ptj is the product of the marginal probabilities over content and position, dependent normalizing scalar Z:r
Proof. By Definition 1, we have that self-attention (q: K, V) = (score(qq K)) = ΣvvpP' = E [V]. Following the assumption that the scoring function fcom(q ky) + Ipool (i, we have:
Pú = (score((;; kj))
exp (Fcom (qi, k,) + Spos (i,
Σ, exp (Jcoml (4im kz) Spec (2, z)) exp (Scone (qi kj)) exp (gpocc ()) Σ, exp (Scone (4t, kz) + Opos (i, =))
exp (fcom q, k2)) exp (Spas (i, j)) Ex(( exp (foont (q:, , kz) + Ipos (i, 2)))
exp (Seceaq k.))
EE, exp (gmaa =))
exp (foce k.))
exp (gpes x))
2

---

exp (Jam (q., k2))
exp (Sre (3, f3) (3,
Σx exp (Jomi knl)) Ex expp (Span
Es exp (fom (9h: kk)))
x exrp (3,
Di exp (fooe i k + Appa 21)
Pffcoal (qè. kj))
Pl (i,j))
Σg P(Scoua (qi, (4ikk)) plopos 2))
p(Jcom (qi, k))
pl Spos (1,
The distributions governing token content and position thus dependent by normalizing scalar factor We can further interpret Z from different perspectives, shedding light on the relationship between token content and position in self-attention (see Appendix K
This derivation shows us that adding positional information in the scoring function of the self-attention mechanism leads product of probabilities PE methods can actually be described by parametrized probability distributions, and that by defining particular distributions can force the model explicitly attend long context.
content and position. We now show that existing
2.2 POSITIONAL ENCODING AS PRIORS TO BAM
With Theorem 1, we can frame positional encoding as priors to BAM. In particular, we present lemmas that derive NoPE (Kazemnejad al. 2023) and ALiBi (Press et 2022 as specific prior distributions to Bayesian self-attention.
Lemma 1.
The causal mask in decoder models is special case of BAM prior where Causal Mask = p(gpos (i,j)) Uniform (,, over a given context Xi.....L
Lemma 2. ALiBi is special case of BAM prior where the token position distribution comprises both Uniform and Laplace distributions.
1 - - over context X1,....L
ALiBi p( (i, Uniform (1, (,, Laplace
0,
Lemma 3. ALiBi becomes local attention the relative length increases.
If Pij = softmax Kq. M.. + A..) then
lim Púj
Li-t -0c
Proofs. See Appendix
B.1
B.2 and
B.3
2.3 PE STRATEGY WITH A GENERALIZED GAUSSIAN AS PRIOR
Now change the distribution positions be Generalized Gaussian Distribution (GGD) and show its advantages over the existing PE methods. We call this new PE method GGD-BAM.
Let B [byylex where big 2 for i 1.... L and j 1...., L be matrix of non-linear biases that are added in the scoring function of the self-attention mechanism. This makes p(ppos(i, = GGD(H, a, B.j i), for B > and a > Self-attention is thus computed as:
softmax (q.KT. M.. + +B,))
When f = 0, 8 and we have an instance of ALiBi, i.c., Laplace prior.
3

---

Lemma 4. GGD-BAM becomes local attention the relative length increases, for any B > 0 and a 2 0.
If Pij = softmax (q:Kk) + M + B then
lim Pij
b- +
Proof. See Appendix B.4
Theorem 2. GGD-BAM is necessarily capable of seeing more context length than ALiBi:
GGD-BAM
Bij. -1 |
Ai -mlj - for some μ and a, and for 3
ALiBi
Proof. See Appendix B.5
2.4 RELAXING THE REQUIREMENT FOR 3
The requirement for 3 > 0 comes from the definition of GGD. We already showed that by making
< 3 1, GGD-BAM can see context beyond ALiBi. Now, we relax this requirement and allow < 0. This effectively increases the size of the tail of the distribution and makes GGD-BAM capable of ignoring local context and focusing on arbitrarily-lon context.
Theorem 3. GGD- BAM ignores local context for any B O and 2 0.
vB < 0, Va 2 0, if Pij softmax (q.Kg M. B.,) then
lim Pij 0,
|
Theorem GGD-BAM takes into account arbitrarily long context for any and 0.
vB 0, Vo 2 If P.
softmax (q, T + M. + B..) then lim
Ps # 0,
1/-1|0000
Proofs. See Appendix B.6 and B.7
2.5 INTUITIVE EXPLANATION OF BAM
To complement the formal derivations provided so far, include here intuitive visualization of the positional priors of BAM. Figure illustrates the probability distribution p(gpos (i, that modulates attention over token positions for fixed query position i. These curves reflect the prior belief of the model over the relevance of each position j when computing the attention for token i, as per the decomposition Pij x p(fcomt (qi kj)) X p(gpos(t, ().
Uniform Prior. This distribution corresponds to the absence of any positional inductive bias beyond the causal mask (NoPE). All positions within the causal window (i.e., j S i) are assigned equal probability, whereas all positions outside the causal window (j > i) are assigned probability zero.
Since Definition 2 states that Piy X p( (q:. k2)) x p(gpos (i, the Pıy of each token outside the causal window is zero. As proved in Theorem 1, this case reduces BAM to the NoPE baseline. The gray-shaded region in Figure represents the non-causal part of the sequence, i.e., positions i, which are masked out auto-regressive settings.
ALiBi as Laplace Prior. The ALiBi mechanism injects linearly increasing bias into the attention logits, which corresponds to Laplace prior over relative position lj - i|. The resulting prior has sharp peak near the query token and rapidly decays over distance. This can be seen in the green curve of Figure I which emphasizes local context. As demonstrated in Theorem 2, this behavior limits the attention window, making ALiBi sensitive to short-term dependencies but less effective capturing long-range interactions. As the relative distance between query and key increases, the probability tends to zero, making ALiBi local attention mechanism.
4

---

Positional Priors
query di position
Causal (Uniform)
Prior Probability Pposli.)
ALIBi (Laplace)
GGD-BAM (8 E (0. 1))
GGD-BAM < Q)
Masked tokens
this
is
long
context Token
for
the
LLM
<EOS>
Figure 1: Visual comparison of different positional priors (1 ) in BAM. Each curve represents the distribution over past token positions for fixed query qi fixed token position
GGD-BAM with 3 E (0, 1). The generalized Gaussian prior with fractional exponent produces heavier-tailed distribution than Laplace. As shown in the pink curve, GGD-BAM maintains significant probability mass across distant tokens. This reflects the theoretical result in Theorem 2, where BAM with 3 E (0, 1) exhibits slower decay rate and is able to attend longer-range dependencies without diminishing the influence of distant content tokens as fast as ALiBi. The effect the location parameter a over the GGD is to move the peak along the distribution. If μ > O, it moves the peak beyond the causal mask. Among the parameters of the GGD, offers a smaller impact on the generalization of BAM, and could be fixed in
with little harm perplexity and retrieval.
GGD-BAM with 0. Perhaps the counterintuitive is the of an inverted generalized Gaussian prior. The black curve illustrates a scenario where the prior probability is effectively zero in the vicinity of the query position and sharply concentrated the far end the context window. Note that 3 < ( is not valid parametrization of the GGD in strict probability sense. However, by relaxing the need for > O, we reach an intriguing theoretical result: the attention mechanism stops looking to local context and shifts to faraway tokens, allowing for arbitrarily-large context scenarios. Even though having attention heads with only negative would make the model blind to local context, attention heads with negative exponent to very long context windows. This is corroborated by our results (see Section B where the longest passkey retrieval was achieved by model in which some attention heads had negative B.
act retrieval heads capable of attending
Relaxing the requirement for 3 > 0 is important to increase context length extrapolation, though is desirable that all attention heads have 3 < 0. The parametrization / < O renders the model be unable to attend to local context, and this increases its perplexity language modeling. Therefore GGD-BAM is capable of not only encoding locality but also to explicitly suppress local content in favor of distant information. This behavior is beneficial in settings where important information appears in long-range context, such as causal reasoning and information retrieval.
Interpretability and Control. One of the central advantages of BAM is that these curves are not merely heuristics but correspond directly to explicit priors over token positions. This makes it possible to visualize, interpret, and even learn the attention pattern preferences of a model. Such view also offers a principled mechanism for extrapolation beyond training context lengths, by selecting priors that maintain probability mass over long sequences.
Scalable Softmax (SSMax). Standard Softmax- based attention mechanisms suffer from a phe- nomenon known attention fading, where the attention distribution becomes increasingly uniform as context length grows (Nakanishi, 2025). This occurs because the denominator in the Softmax computation increases with context size n, while the numerator for each token remains constant,

---

leading to vanishing attention peaks. To address this, Nakanishi (Nakanishi 2025)
proposes Scalable
Softmax (SSMax), which rescales the attention logits dynamically as function of sequence length: e logM)s;
2 5 Σラ e m where E IR is learnable scalar. Although SSMax is not
PE method, can be used address distinct and complementary challenge in long-context attention. the fading attention. This modification of the softmax function preserves the sharpness of the distribution regardless of input size, mitigating the flattening effect observed in softmax and improving long-contens generalization.
While BAM models positional structure as an explicit prior over token positions, it becomes more susceptible to fading aftention as it increases the context length that the LM can attend to. The two techniques are compatible and can be used together, enabling improved context length extrapolation both by better normalizing factor via SSMax and in the distribution over positions using BAM. This learnable normalizing factor in the softmax function also serves the purpose learning the normalizing dependence scalar Z that we introduced during the derivation of BAM.
The BAM matrix B. To use GGD-BAM, the only modification we need apply in the transformer is to add a relative position based matrix B to the attention score. Figure shows how the attention score is computed for the entire query matrix Q using GGD-BAM. As usual, the causal mask M masks tokens beyond qi by adding - 3o to those respective positions. Each value in B is computed according to the relative position j - i, a, 33 and μ. We fix = for most of our results.
B
q. kJ
k. ki q.kk
UP
00
00
-00
2
4
qaki quki qeke gak]
9
0
00
00
-1
0
1
2
3
1
qakk qs ggke] qakī qak]
+
00
-00
+
X
-2
14
a
q.k] qok] q.kk q,k[ 4skk
0
0
-00
-3
-2
-1
1
q.lc] q:kī q.k] q:ki qskk
0
0
3
-2
QK
Figure 2: Visual representation of the scoring function in GGD-BAM. The first matrix accounts for the content and the two others for the Uniform and GGD positional priors.
3 EMPIRICAL ANALYSIS
We perform empirical analysis evaluate the behavior of GGD-BAM in realistic long-context scenarios. We compare BAM and its variant coupled with Scaled Softmax (BAM SSMax) (Nakanishi 2025) against several widely used PE methods: Sinusoidal (Vaswani 2017). NoPE (Kazemnejad et al. 2023). RoPE (Su et al. 2024), Local RoPE (RoPE with sliding- window attention), and ALiBi (Press 2022 as well their versions coupled with Scaled Softmax.
ct al.
All models presented in this section contain approximately 120M parameters, including -25M for input embeddings and -95M for the transformer layers. Models were trained on the Web 10B dataset (Penedo et al. 2024) using the Mistral-7B v0.3 tokenizer (Jiang et al. (023). with training context length of 512 tokens. We evaluate model performance on two tasks: (1) language modeling
on long-context samples drawn from Fine 10B and Wikipedia (Foundation 2023); and (2) the
Passkey Retrieval task (Mohtashami & Jaggi 2023), which measures a model's ability to retrieve specific information from distant positions in the input sequence.
BAM introduces three learnable parameters, 9., 43, and 6ur for each attention head in each layer. This results in total overhead of 3 x Heads Layers parameters. In all experiments reported in this section, we train only 9,, and 0g. fixing 6,, = O. which amounts to just 384 additional parameters in a 120M parameter model. Further implementation and training details are provided in Appendix

---

3.1 PASSKEY RETRIEVAL
To assess the capability the LM to access and use long-range information, evaluate the models on the Passkey Retrieval task (Mohtashami & Jaggi 2023). This task measures whether language model can recall a specific five digit number called the "passkey" embedded somewhere within longer context window. In the passkey retrieval task, generate 20 sequences, each containing a passkey inserted at uniformly spaced positions: 0L. 1/
TO Tg .g....,L. L. In the end the sequence, we append the prompt <The passkey is :> and measure how accurate the model can predict the next five tokens. To avoid inaccuracies due to tokenization, each digit is considered a distinct token.
As shown in Figure 3 only models trained with BAM retained high accuracy when extrapolating to sequences up to 32,000 tokens. BAM SSMax maintained perfect accuracy across all tested lengths, demonstrating robust access to information throughout the full context window.
1.0
0.8
ALiBi SSMax
Accuracyl
0.6
BAM 55Max(ours)
NoPE SSMax
RoPE Local
0.4
RoPE SSMax
Sinusoidal SSMax
0.2
0.0
0
5000
10000
15000
20000
25000
30000
Sequence length
Figure 3: Passkey retrieval accuracy with distinct PE, BAM SSMax outperforms all PE methods maintaining perfect accuracy for context beyond 64x the training context length.
contrast, other evaluated PE methods such Sinusoidal, RoPE, and NoPE rapidly degrade to near-random accuracy beyond their training horizon. Even ALiBi, which showed good perplexity extrapolation (see Appendix D). struggles to maintain retrieval performance at very long context windows. Appendix H.4 shows similar trend to the evaluated PE methods without SSMax.
1.0
0.95
Position Sequence
20%
0.77
40%
0.77
0.6
60%
0.59
80%
0.86
0.2
0.0
<1K
51K
102K
153K
204K
256K
307K
358K
409K
460K
512K
Sequence Length
Figure 4: Passkey retrieval accuracy across context lengths and depths. The horizontal axis represent context length and the vertical axis represents the position of passkey the context. In the bottom row and last column, average accuracy across length and position, respectively.
7

---

In Figure 4 we show heatmap plot with passkey retrieval accuracy considering all possible depths of the passkey. We see that BAM is able to score perfectly in most lengths and depths while being trained only on length 512. BAM only degrades to 0% accuracy in 3/55 of the evaluated lengths and depths. Although accuracy is above 80% for x00x the training length, it appears that the model will eventually, however we did not have enough vram test for longer context.
degrade
The superior performance of BAM in this retrieval setting provides empirical support that its positional prior enables meaningful access distant content, rather than merely preserving surface-level fluency using local context. In Appendix D we assess the perplexity of GGD- in comparison to the
baselines in Wikitext (Foundation (023) and Fineweb (Penedo al. 2024 In Appendix |F] we assess GGD-BAM in the Needle in a Haystack (NIAH) subset of the RULER benchmark (Hsich)
al. 2024). as a means provide a thorough empirical analysis on long-context extrapolation.
3.2 ATTENTION WEIGHTS FOR 3 0. B0,Be (0,1) AND 3 1
To provide empirical evidence for the claims presented in Theorems and 4, we visualize attention weights from individual attention heads and their respective values of / in the Passkey Retrieval task.
To create this visualization, we craft a passkey retrieval prompt of length 841 (825 + 16). where the first 32 tokens are the task prompt-instruction to the model to remember the passkey; the following 25 tokens are the passkey itself; the next 768 tokens are filler text that the model should ignore; and the last 16 tokens the retrieval prompt answer that the model should complete with the passkey.
prompt
prompt task
prompt
prompt task
prompt
prompt task
task
task
task
passkey
passkey
passkey
passkey
passkey
passkey
filler
filler text
filler
filler
filler
filler text
text
text
text
text
prompt
prompt answer
prompt
prompt answer
prompt
prompt answer
answer
answer
answer
Queries
Keys
Queries
Keys
Queries
Keys
(a) 0g 0.205 (e.g. 8 8850
(b) 0g 0.456 (e.g. $ E ())
(c) 0g 1.031 (e.g. ~
Figure 5: Attention weights from GGD-BAM during the Passkey Retrieval task. When 3 < 0,
attention concentrates on distant keys (e.g., the passkey tokens), suppressing nearby content.
In Figure (a), with < O, the attention head effectively ignores local context and sharply focuses on distant tokens tokens-including including the passkey-as as predicted by our theoretical formulation. This behavior aligns with Theorems and 4. which show that for S O, the GGD-BAM prior suppresses attention weights near the query token and sustains probability mass across long distances. As result, these heads act as retrieval specialists, increasing the performance of the model for passkey retrieval.
In Figure (b), with @ (0, 1), attention exhibits long-tail -ail behavior, allocating attention mass more evenly across both nearby and distant positions. Compared to 3 ~ 1, this head decays more slowly, retaining mid- and long-range dependencies. This supports our claim in Theorem 2 that GGD-BAM with fractional 3 values can attend to longer contexts than ALiBi (8 ~ 1).
In Figure (c). with 3 1, the attention pattern is highly localized, closely resembling ALiBi's linear bias. Attention focuses on immediate neighbors, with a rapid decay over distance. This configuration is suitable for capturing local dependencies but is inadequate for long-range retrieval, as evidenced by its poor performance on the passkey task at large context lengths.
The visualizations in Figure show that BAM is highly interpretable. The negative fg distribution that was conjectured to improve long-context retrieval appeared after training, and during inference effectively caused higher attention weights towards the passkey that was further in the context.
8

---

3.3 83 AND & TRENDS AFTER TRAINING
We now analyze the trends of 0g and 0. after training. Our first analysis focus on the 120M model trained with distinct context sizes. Figure an interest trend where identify three linearly- separable clusters of parameters: the first cluster with 0g > O, where each attention head works similarly a Laplace distribution (ALiBi); the second cluster with -0.6 S 0g S which works as a retrieval head; the third cluster has fewer instances than the other two, with 0g < 0.6. We conjecture that this cluster works as a more aggressive retrieval head.
512
3
1024
0,
-1
03 -0.6
-0.6 6g <0
0p
-2
2
Figure 6: Trend of 0g x la regarding the 120M model trained on 512, 1024, and 2048 tokens.
When we analyze distinet model scales, the same three clusters emerge. Figure 7 shows the same clusters identified in Figure 6 providing evidence that these probability distribution over positions stable and transferable across many tasks.
1.18
432M
120M
0.
θμ -0.6
-0.6 6, <
-6
3
2
1
2
Figure Trend 0g x O,, on different model scales.
In Appendices D and H present additional empirical results for perplexity, for long-context per- formance on the Ruler Benchmark, as well as several additional ablation studies. We found that BAM has similar perplexity ALiBi and has no measurable impact inference time, while outperforming every baseline in context extrapolation across all evaluated tasks of the Ruler benchmark.
4 RELATED WORK
Several strategies for PE have been developed to allow Transformer- based LMs to encode token order. The most used PE methods are described below:

---

Sinusoidal Positional Encoding. Introduced by Vaswani et al. (Vaswani 2017). sinusoidal encod- ings inject fixed positional information into the model by adding position deseeden-dependentt vectors to token embeddings. In the self-attention mechanism, each q: k, and v, have both content and positional information. BAM, instead, explicitly models position as a separate probabilistic component and does not add positional information directly into qi kj, and V..
Rotary Position Embedding (RoPE). RoPE (Su et al. 2024) encodes absolute positions through complex-valuee rotations applied directly Q and K. RoPE encodes relative position by phase- shifting the token representations before the dot product. While RoPE introduces position at the dot-product level and preserves relative distance structure, it does not decouple content and position semantically. |Han & Jil (2025) show that RoPE asymptotically disentangles semantics and position information in additive components in self-attention logits. Thus, can see that RoPE asymptotically approximates BAM as the context length increases, though it lacks the flexibility and interpretability of assigning positional priors as established our formulation of BAM.
T5 Relative Bias. T5 (Raffel et al. 2020 avoids absolute encodings entirely and instead learns that is added to the attention logits. This PE method
bias b(i b(i-j) j) for each relative distance could be viewed in BAM as an empirical non-parametric distribution over the positions plapos ()).... However, both our theoretical grounding and experiments show that parametric distribution with fewer trainable parameters can achieve superior context length extrapolation.
No Positional Encoding (NoPE). Haviv et al. 2022) and Kazemnejad al. (023) examined transformers without explicit PE. Although the attention mechanism is purely content-driven, the authors were the first to highlight that the causal mask Myj (often not represented explicitly in the notation of other PE methods) in becoder-based transformers is sufficient to derive both absolute and relative PE. We showed NoPE to be special case of BAM under uniform positional prior.
Attention with Linear Biases (ALiBi). ALiBi (Press et al. 2022) injects a linear positional penalty into the attention scores. This formulation is interesting because allows context length extrapolation in language modeling without introducing learnable parameters to the LM. We show that ALiBi is a special case of BAM with a Laplacian prior and we test it as an initialization strategy in our ablation study (see Appendix(). We explain how ALiBi maintains low perplexity in longer context windows and why it fails to retrieve information as becomes local attention as the context length increases.
CONCLUSION
We introduced BAM, principled probabilistic framework that reconceptualizes positional encoding as a prior over token positions within attention. By framing the attention mechanism as a factorized joint distribution over content and position, BAM not only offers theoretical grounding for existing methods such as NoPE and ALiBi, but also motivates new families of positional priors. Our proposed Generalized Gaussian prior GGD-BAM significantly improves context length extrapolation in passkey retrieval task by more than 25x compared to other PEs while maintaining low perplexity.
Despite its simplicit--increasse parameters- (GGDBAM enables models to attend over significantly longer context windows without direct exposure during training. Experiments on Fine and Wikipedia show that BAM uniquely able to recover distant information even at 150,000-token sequences, where other methods collapse. Moreover, our theoretical results demonstrate that BAM can express attention patterns that emphasize distant context suppress locality, offering new axis of inductive bias design for Transformers.
the parameter count of the model in negligible 0.00032% trainable
Future work includes applying BAM to larger models, further exploring the interpretabilitty of learned positional priors, and extending the BAM framework to multi-modal input settings. Additionally, it remains an open question whether the extrapolation capabilities induced by BAM are preserved, or potentially enhanced, after instruction and preference muninnunng. Investigating BAM under these downstream adaptation regimes is crucial for understanding its robustness in real-world applications. We further discuss all limitations of this work in Appendix
10

---

ACKNOWLEDGMENTS
This study was financed in part by the Coordination for the Improvement of Higher Education Personnel (CAPES) Finance Code 001; by Conselho Nacional de Desenvolvimento Científico Tecnológico (CNPq)- Grant Number: 443072/2024-8; and by Fundação de Amparo Pesquisa do Estado do Rio Grande do Sul (FAPERGS) Grant Number: -3/2551-0000891-3.
This work was supported by Kunumi Institute. The authors thank the institution for financial support and commitment to advancing scientific research.
REFERENCES
Yushi Bai, Shangqing Tu, Jiajic Zhang, Hao Peng. Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong, et al. Longbench v2: Towards deeper understanding and reasoning on realistic long-context multitasss. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Long Papers), pp. 3639-3664, 3664,, 2025.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. 2018.
arXiv preprini
Wikimedia Foundation. Wikimedia downloads version 20231101.eco .nn, 2023. URL httph::t///mnm
wikimedia org
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Chi Han and Heng Ji. Computation mechanism behind Ilm position generalization. arXiv preprint 0rXiv:2503.13305, 2025.
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer Language Models without Positional Encodings Still Learn Positional Information. arXiv preprint arXiv:2203. 16634, December 2022. doi: 10.48 10.48550/arXiv.22003 16634. URL http:://axxi.crrrrte// abtrgg/bbb22233
16634
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt, Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03.4 03300, 2020.
Cheng-Ping Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.0666 2024.
Yutong Hu, Quzhe Huang, Mingxu Tao, Chen Zhang, and Yansong Feng, Can perplexity reflect large language model's ability in long text understanding? arXiv preprint arXiv:240S.06105, 2024.
Yunpeng Huang, Jingwei Xu, Junyu Lai, Zixu Jiang, Taolue Chen, Zenan Li, Yuan Yao, Xiaoxing Ma, Lijuan Yang, Hao Chen, al. Advancing transformer architecture in long-context large language models: A comprehensive survey. arXiv preprint arXiv:2311.1335 20233
Albert Q Jiang, A Sablayrolles, A Mensch, C Bamford, D Singh Chaplot, Ddl Casas, F Bressand, G Lengyel, G Lample, L. Saulnier, et al. Mistral 7b. arxiv. arXiv preprint arx0v:2310.06825, 10, 2023.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding length generalization in transformers. Advances in Neural Information Processing Systems, 36:24892-24928 2023.
Schoon Kim, Coleman Hooper, Thanakul Wattanawong, Minwoo Kang, Ruohan Yan, Hasan Gene, Grace Dinh, Qijing Huang, Kurt Keutzer, Michael W Mahoney, et al. Full stack optimization of transformer inference: survey. arXiv preprint arXiv:2302. 14017, 2023.

---

Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 02 2024. ISSN 2307-387X. doi: 10.1 Bacchacl_a_00o000 URL UKK/hhtttss//bb....r//0011161/62 ___00688
Amirkeivan Mohtashami and Martin Jaggi. Random-access infinite context length for transformers. Advances in Neural Information Processing Systems, 36:54567-54585, 2023.
Ken M Nakanishi. 2025.
Scalable-softmax is superior for attention. arXiv preprint arXiv:
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:308 11730811-3088 1-30849, 2024.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables imput length extrapolation. In International Conference on Learning Representations, 2022. URL
https://openrevier net wttp:/eeeeeeoemmeeeeee
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-61 2020.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
Ryan Singh and Christopher L Buckley. Attention: Marginal probability is all you need? arXiv preprint arXiv:2304004556 2023.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced
transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
A PRELIMINARIES AND NOTATION
We denote by x; the 7th token in sequence of length written as Xi,.., L. Each token is projected into query, key, and value via learned matrices: the query qe E Rixd is obtained from x; Wq q. and the full key and value matrices given by K XW E R and V XW, E Lxd respectively.
We define M as the standard causal mask used to enforce auto-regressive decoding constraints. We to denote slice from matrix. For instance, M. is the causal mask line vector for single qi-
In our formulation, we interpret attention weights as a joint probability distribution Ptj over two de- pendent components: content and position. The random variable foom (q;, ky) denotes the content sim- ilarity between the 7th query and the jth key, modeled by the conditional probability p(fceal (q: kj)). Similarly, the positional relation is modeled by the random variable 9po. (i. j), which represents the relative position of token j with respect token captured by the prior distribution (i.)).
a convention to denote the scoring function of the non-normalized logits in self-attention as adapt the notation used in ALiBi
score(q,, K) K M... To follow this convention, to present the linear biases in matrix form, so the scoring function becomes: score(q, K) =
q. M.. A. Finally, we introduce the BAM score(q,, K) q:k + M, + B,, by adding a matrix of non-lincar biases to the scoring function.
Another convention adopt is a ozz) softmax (z) E, exp(2,))
exp(x;)
12

---

Positional Priors
The prior over positions, p(opox (i.j)), is modeled using parametric distributions
over relative positions lj - z/. We consider the following parametric distributions:
Uniform distribution:
assigns equal probability mass to all valid positions before the
current token:
Pl (i,j)) Uniform(a, b, b) x IIaa b].
Laplace distribution: decays exponentially with distance, controlled by scale parameter
0:
Plopos (i,5)) = Laplace(μ= 0, o, exp
Generalized Gaussian distribution (GGD): introduces a shape parameter 3 > 0 and scale
a >> allowing flexible control over decay behavior:
Upe (i,j)) CC, a, x exp
This formulation generalizes several priors: when it recovers the Laplace distribution; when 3 it becomes Normal; larger values of B yield sharper, more localized priors; lower values yield large tailed distributions.
B PROOFS
B.I NoPE IS UNIFORM PRIOR TO BAM
Lemma 1: The causal mask in decoder models is special of BAM prior where
Causal Mask + p(gpos(( (i,j)) Uniform(L,4,j) over given context XI,...,L
Proof. The causal mask in decoder model changes the scores of every token X,+1,4+2.... L to inf:
0:
<
causal-self-attention = softmax(q:kt + M,,) where M.. = [-20;
otherwise
Since q: K has only content information and Mi has only positional information, can Theorem to rewrite it as
exp(q:k))
exp(M:.)
E, (exp(q.kk)))
E, (exp(M)))))
Only the softmax softmax(M,.) term depends on the token position, so is equivalent to p(gpos (i.j)):
p(gpos (i,j)) (exp((Mizz)) exp(M..) 1
ifj otherwise
0,
which is Uniform distribution over tokens X1...--
B.2 ALIB1 A PRIOR COMPRISING BOTH UNIFORM AND LAPLACE DISTRIBUTIONS
Lemma 2: ALiBi is special case of BAM prior where the token position distribution comprises both Uniform and Laplace distributions.
ALiBi + p(opos (i,j)) Uniform (1, i, Laplace 0, 1..15-1
- over context XI...,,
13

---

Proof. Let A = [0ıy]2xL where dzy =-mlj- =l, for i 1,..., and j 1..... L be the matrix with linear biases defined in ALiBi PE. The self-attention mechanism with ALiBi as PE is computed
softmax (q.K², Mi. +::)) V
where M. is the causal mask and -mli - il are the linear biases added by ALiBi. According to Theorem and Lemma rewrite the softmax as:
exp(q:k,)
exp(M (exp((Mix))
exp(A.))
(exp((a:k)))
Σ. (exp((Ai)))
ZalAB:
where ZALII is a scaling (normalizing) factor. Lemma further allows us to substitute the causal mask term with uniform distribution:
exp(q.k))
(xp(Ait) 1
Uniform (1,
Σ. (exp(q:kkk
Eg (exp(A..)) ZALIBI
The softmax (-mlj - il) term has no content information, just positional, so we further work on it:
exp(Ai))
exp
mlj - il)
Σ- (x(pp(( (exp(Aiz))
Σg exp
m il)
exp
mij il)
m77 2/2 Σ, exp
Laplace
(0, 1,
Σ, 2 exp
m7 il)
Laplace
0, j i
We can drop the scalar normalizing denominator Σg m exp - m 2 il) since ZALIBI accounts for the normalization of the whole expression. Back Definition we have p( (i, Uniform (1,4,j) Laplace (0, 1 i).
2
B.3 ALIBI IS LOCAL ATTENTION FOR LARGE J - LENGTHS
Lemma ALiBi becomes local attention the relative length lj il increases.
If Pij softmax (q.KT Mi. + Aii) then lim
Pij
1i-i -0c
Proof. We prove this lemma for a fixed query qu The scoring function ALiBi has three components: K.K. M... and Ai.. Let take the limit of the scoring function:
lim (q.K M.. + A,,) K. M.. 1
lim (A:.) 15-1|+000
We drop the limit the causal mask M,, the only effect of increasing the context size and the distance between and j in the causal mask is making the mask bigger in size, but it stills follows the same formation law with 0 to the left of the query and -o elsewhere.
lim (Aio) = lim (-mlj i|) = -00
1i-i
14

---

When plug -oo back into the scoring function that it becomes -oo, and consequently the softmax becomes 0.
lim +0
Pij
lim 1--0
softmax (q:K² + Mi + Aio)
B.4 GGD-BAM IS LOCAL ATTENTION FOR LARGE j - LENGTHS
Lemma 4: GGD-BAM becomes local attention the relative length increases, for any B > 0 and a > 0.
If Pij = softmax (q:K² + M. + B:,) then lim bjiil
Pij = 0
Proof. This proof is similar to Lemma 3. We prove this lemma for fixed query qu Let 0. The scoring function of GGD-BAM has three components q:k MM and Bi,, lets take the limit of the scoring function and see how it behaves:
lim (q.K3 M. + B:.)= M, 1
lim 15-1/0000
(B:.)
We drop the limit the causal mask M,. as the only effect of increasing the context size and the distance between and j in the causal mask making the mask bigger in size, but it stills follows the same formation law with 0 to the left of the query and oo elsewhere.
(-
(-13=
lim
(B:.)
lim 15-100000
11-00000
When we plug oo back into the scoring function we see that it becomes oo, and consequently the softmax becomes O.
lim Dj 0000 1000
lim
Pij
softmax (q.K + M. + B..) 0
B.5 GGD-BAM SEES MORE CONTEXT THAN ALIBI
Theorem GGD-BAM is able see more context length than ALiBi
B; 1
GGD-BAM
-ml < 1 for some μ and and for 3 <
ALiBi A;
Proof. For GGD-BAM to see more context than ALiBi, must be the case that the ratio between
Bi and Aij -| ; -m is less than 1 for some 3 < and for some a and μ. Let 0, a v/m
and B E (0, 1), then we have:
B
A
m
i
m
i
-m
i
8
< 1
15

---

Since the ratio between BAM and ALiBi can be less then 1, BAM shrinks at slower rate than ALiBi, effectively capturing longer contexts as | il increases.
B.6 GGD-BAM IGNORES LOCAL CONTEXT FOR B < AND a >
Theorem GGD-BAM ignores local context for any < and a > 0.
vB < O, Va > 0, If Pij softmax (q:K² + M, + B:,) then lim Pij 0,
1j-21-00
Proof. This proof is similar to Lemmas 3 and 4. We prove this lemma for fixed query q: Since i is fixed for a query, li - i - 0 implies that the token and are the same token. The scoring function of GGD-BAM has three components q:K*, M. and B;,, lets take the limit of the scoring function and see how it behaves:
(q:K + M. + B,.) = =q.K" + M, + lim (B:.) b-
lim
1-1|-40
We drop the limit in the causal mask M.. as the only effect of i and being the same token that My 0.
When the shape parameter is negative, then can perform the following manipulation.
(-13-01
(-
|8|
a
lim (B:.) = lim 1--11-00
lim 1--10>0
00
-
When plug -oo back into the scoring function that it becomes and consequently the softmax becomes O.
lim Pij = lim softmax (q:K7 + M, + B:,) = 1--1-0 15-i
Since the result of the softmax is zero, GGD-BAM does not attend to local context when B < 0 and
0.
B.7 GGD-BAM CAN HAVE ARBITRARILY LONG CONTEXT FOR B AND & > 0
Theorem GGD-BAM takes into account arbitrarily long context for any and a >> 0.
vB < 0, Vo > 0, If pt, = softmax (q.K3 T + M.. + B..) then
Pij + O, |
lim
Proof. This proof is similar to Lemmas 3 and 4 and Theorem 3. We prove this lemma for fixed query qı The scoring function of GGD-BAM has three components q M... and B... lets take the limit of the scoring function and how it behaves:
lim (q.K M,. + B..) q.K M,. 1-
lim (B:.) 15-10+00
We drop the limit in the causal mask M,. as the only effect of increasing the context size and the distance between and j in the causal mask making the mask bigger in size, but it stills follows the same formation law with 0 to the left of the query and co elsewhere.
When the shape parameter is negative, then we can perform the following manipulation.
(-
-
(-1-a) a |&8
(
lim (B:.)
lim 1000
lim +00
0
4000
16

---

When we plug 0 back into the scoring function that it becomes -o, and consequently the softmax becomes O.
lim Pij = lim 1-1
softmax (q:K² + M, + B:,)
15-il
= softmax (q:K² + M
Since the result the softmax not necessarily zero, Pij also is not necessarily 0, thus GGD-BAM can attend to arbitrarily long context when B < 0 and o > O.
C EXPERIMENTAL SETUP
Here we detail the experimental setup necessary to replicate our results.
Our models were based on Llama3 (Grattaioni al., 2024). We used RMSNorm Zhang & Sennrich 2019) and, on feedforward blocks, we used swiglu activation function (Shazeer) 2020p. We trained all our models on the Fineweb 10B dataset (Penedo et al., (024) utilizing the Mistral-7B v0.3 tokenizer (Jiang al.) 2023) for text processing. Training performed using context lengths of 512, 1024, and 2048 tokens to evaluate performance under different sequence lengths.
To maintain document separation within packed sequences, we employed an attention mask preventing self-attention between distinct documents (Grattafiori et al., (0024). The models were optimized using RAdam with decoupled weight decay set to 0.1 and an initial learning rate of 1 10 We applied cosine learning rate decay schedule, reducing the learning rate to a minimum of 0.1x x its initial value.
We trained LMs up 1.1 billion learnable parameters. Table details the configuration of each trained model, including embedding, trainable parameters, attention heads, and hidden layer size.
Table 1: Architecture details of the LM we evaluated in our study.
Attribute
120M
432M
1.1B
Parameters embedding
-25M
-50M
-67M
Parameters transformer
-955M
-38M
-IB
Parameters BAM (0a, 0g and 0₂)
576
1008
1440
Attention heads
16
24
32
Layers
12
14
15
Hidden size ff hidden size
768
1536
2048
768
2x 1536
4x2048
Learning Rate
x 10
5 x 10
3 x 10
Training was executed for approximately 19, 251 steps with global batch size 589, 824 tokens, leveraging up to 6 NVIDIA A6000 GPUs. For the 1.1B parameter models, we perform additional 256 training steps with a context length of 1024. This was necessary because bigger models tends to overfit in the trained context length (see Appendix H.5 for further details). The resulting models from these specific training configurations formed the basis for our subsequent performance evaluation.
In our implementation we define three learnable parameters, 0n, O2 and 0g for each attention head. [b,]]x where:
So B-
bij - (10-1) - (6 - - -m0 00
for i = 1, ... L and 1, L. So the GGD is parametrized in the following way: μ e
0₂1 and To evaluate our models in contexts longer than 10,000 tokens we implemented key-value caching (Kim
"a that ensures 0. 10 avoids division by zero when 3 < 0.
al. 2023
17

---

D PERPLEXITY EVALUATION
D.1 PERPLEXITY ON SHORT CONTEXT
We evaluate generalization for models with distinct PE methods by computing the perplexity for models trained on a context length of 512 in hold-out validation Fineweb with sentences up to 512 tokens and on Wikipedia with sentences up to 512 tokens. Results are shown in Table 2
Table 2: Perplexity of 120M models in Wikipedia and Fineweb data.
Fineweb Wikipedia
Sinusoidal
18. 18,9310
22.3573
ALiBi
18.5452
22.1771
NoPE
19.8609
23.7898
BAM
18.9281
22.2640
RoPE
18.3599
22.4428
Sinusoidal SSMax
19.1467
22.3150
NoPE SSMax
20.3899
23.7581
ALiBi SSMax
18.4967
22.2125
BAM SSMax
18.6897
22.1363
RoPE SSMax
20.1854
24.4507
We see that in context length seen during training, BAM outperforms almost all the compared baselines. The only exception is ALiBi, which outperforms BAM by less than 0.1 points in perplexity.
ALiBi outperforming BAM in this case is expected from probabilistic interpretation of PE. ALiBi is a Laplace distribution only focused on the local context whereas BAM has also the ability to use information in long context. We notice in experiments focused on long context (see Appendices F and Hy that this small gap in perplexity to ALiBi translates into huge long-context retrieval gain.
D.2 PERPLEXITY ON LONG CONTEXT
We evaluate context length generalization for models with distinct PE methods by computing the perplexity for models that were trained context length of 512 in hold-out validation set of sentences longer than 512. This a similar evaluation procedure as performed by Press et al. (2022).
Results shown in Figure 8 are consistent with those reported in the ALiBi paper (Press et al. 2022). confirming that Sinusoidal, RoPE, and NoPE fail to extrapolate to sequence lengths beyond those seen during training. For these models, perplexity increases sharply the context grows, indicating that the models are unable to maintain coherent predictions over long sequences. In contrast, only ALiBi, BAM, BAM combined with Scalable Softmax (SSMax), and RoPE local maintain a stable perplexity profile under context extrapolation. These models exhibit sub-linear or nearly flat perplexity growth as the input length increases to 32,000 tokens desenst-espite being trained only on context length of 512 tokens. This confirms our theoretical results where priors introduced by BAM provide robust generalization to unseen context lengths, on par with ALiBi's linear bias-based extrapolation.
We evaluate all the PE context length extrapolation regarding perplexity on Wikipedia dataset. Figure D shows the log-scaled perplexity and we see similar trend Fineweb 10B. BAM, RoPE and ALiBi are able to maintain low perplexity across all the evaluated context lengths. We see that SSMax has more impact in lowering BAMs perplexity when compared to other PEs.
However, this type of perplexity-basee evaluation has limitations (Hu et al., 2024). It does not measure whether the model attends to the full sequence or only relies on the most recent tokens. For instance, the RoPE Local variant implemented here applies local attention restricted to sliding window and still achieves competitive results with ALiBi. This indicates that models can make accurate next-token predictions without integrating information from earlier parts of the sequence. Hence, we claim that perplexity should not be taken as the sole measure of effective context length extrapolation. Indeed, information retrieval evaluation seems to be more suitable to assess extrapolation of trained lengths.
18

---

ALIBi
SSMax
BAM(ours]
BAM SSMax(ours)
105
NoPE S5Max
ROPE
RoPE Local
SSMax
Perperityt
10
Sinusoidal
pempextyy
10
Sinusoidal SSMax
10a
10a
5000 0
10000 15000
20000 25000 Sequence length
30000
5000
10000
15000 20000 25000 30000 length
(a) PE without SSMax
(b) PE with SSMax
Figure 8: Log-scaled perplexity computed up 64x the training context length of 512 tokens. BAM, RoPE Local and ALiBi are able to maintain the lowest perplexity on longer contexts.
106
10c
ALiBI
ALIBI SSMak
BAM(ours)
55Maxiours!
105
NoPE
10
NaPE SSMax
RoPE
RoPE Local
Perpexttyl
104
Sinusoidal
Perpexity(=
104
Sinusoidal SSMax
10a
10a
10a
10a
10
10a
5000
10000
15000
20000 25000 30000
5000
10000
15000
20000 25000
30000
Sequence length
Sequence length
(a) PE without SSMax
(b) PE with SSMax
Figure 9: Log-scaled perplexity on Wikipedia dataset, BAM, RoPE Local and ALiBi able to maintain the lowest perplexity on longer contexts.
E DOWNSTREAM EVALUATION
While perplexity provides a general notion regarding model capability in language modeling, does necessarily correlate with capability downstream tasks. Here we evaluate the large scale IB parameter models performance on downstream tasks from MMLU/Hendrycks et al. 2020). ARC-easy, and ARC-challenge benchmarks Clark et al. (2018). Table that BAM SSMax superior to RoPE in all the evaluated benchmarks.
Table 3: GGD-BAM RoPE Large-Scale 1B parameter models on MMLU and ARC Benchmarks
MMLU
ARC-Easy
ARC-Challenee
BAM SSMax 0.3716 0.3573
0.5770 0.5715
0.4132
RoPE SSMax
0.4123
F RULER BENCHMARK
To assess BGD-BAM capability to attend to long-context information, we assessed its performance in the NIAH subset of the Ruler benchmark (Hsieh al.) 2024). We deliberately chose only the NIAH
19

---

subset because it is designed to isolate the PE's ability to access information specific distance while ensuring the model maintains stable next-token prediction performance.
Performance other tasks, such as Question Answering and Variable Tracking, is known to correlate more with model size than with context extrapolation, and thus fall out of the scope of this paper.
F.1 GGD-BAM VS BASELINES (120M)
As seen in Table g GGD-BAM outperforms all the baselines on the three variations of the NIAH task. Specially in NIAH Single 1, our PE method outperforms all other PE methods by a substantial margin, being the only one able to perform retrieval above 6, tokens.
The second best performing PE is RoPE, followed by ALiBi. In the Single I version of Ruler NIAH, RoPE maintains 0.64 accuracy while ALiBi scores mere 0.16.
Regardless of the PE, in the Single 3 version of NIAH, where the Passkey appears in the form of UUID, only BAM achieves an accuracy above 0.8.
Table 4: Accuracy of 120M models on the NIAH subset of the Ruler Benchmark.
Task
PE
1K
1.5K
2K
3K
4K
6K
8K
10K
12K
NoPE SSMax
0.04
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
RoPE SSMax
1.00
1.00
1.00
1.00
0.64
0.00
0.00
0.00
0.00
ALiBi SSMax
0.96
0.26
0.22
0.18
0.16
0.
0.02
0.04
0.02
Single
RoPE Local
0.40
0.22
0.24
0.20
0.12
0.08
0.00
0.02
0.00
Sinusoidal SSMax
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
BAM SSMax (ours)
1.00
1.00
1.00
1.00
1.00
1.00
0.98
0.92
0.88
NoPE SSMax
0.70
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
RoPE SSMax
0.98
0.96
0.94
0.84
0.52
0.00
0.00
0.00
0.00
ALiBi SSMax
1.00)
0.46
0.12
0.10
0.06
0.02
0.00
0.00
0.00
Single
RoPE Local
0.84
0.36
0.10
0.14
0.10
0.00
0.00
0.00
0.00
Sinusoidal SSMax
0.70
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
BAM SSMax (ours)
1.00
1.00
1.00
0.88
0.24
0.06
0.02
0.00
0.00
NoPE SSMax
0.04
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
RoPE SSMax
0.28
0.30
0.10
0.06
0.00
0.00
0.00
0.00
0.00
ALiBi SSMax
0.24
0.04
0.00
0.00
0.00
0.00
0.00
0.00
0.00
Single
RoPE Local
0.10
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
Sinusoidal SSMax
0.08
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
BAM SSMax (ours)
0.84
0.68
0.42
0.08
0.00
0.00
0.00
0.00
0.00
F.2 LARGE SCALE GGD-BAM VS ROPE
Here we perform a large-scale experiment directly comparing RoPE and BAM with IB parameters in the Ruler benchmark. Results are show in Table 5 BAM achieves superior performance in comparison to RoPE every evaluated task in the Ruler benchmark, with a highlight of achieving almost perfect accuracy across tasks the Single subset.
Other tasks such as Multikey 2 and Multikey are harder for both models. This shows that only our pre-training may not be enough for models to perform such tasks. However, since our goal here is to assess how distinct PE behave on exactly the same training regime, we see that GGD-BAM clearly achieves better longer context in Single 1, and 3, MultiKey 1, MultiQuery and MultiValue.
G LONGBENCHV2
In Table 6. present the results of 1B parameter models the complete LongBenchV2 bench- mark Bai et al, (2025) limited in 131k tokens. We chose to report just the large scale 1B parameter models because smaller models perform close random guessing in these tasks. BAM outperforms
20

---

Table 5: GGD-BAM RoPE Ruler Benchmark Large-Scale IB parameter models.
Task
PE
1024
1536
2048
3072
4096
6144
8192
12288
16384
24576
32768
BAM SSMax
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1,00
1.00
1.00
1.00
Single
RoPE SSMax
1.00
0.88
0.68
0.40
0.30
0.10
0.02
0.00
0.00
0.00
0.00
BAM SSMax
1.00
1.00
1.00
0.98
1.00
0.88
0.82
0.46
0.18
0.06
0.02
Single
RoPE SSMax
1.00
0.82
0.62
0.32
0.18
0.04
0.00
0.00
0.00
0.00
0.00
Single
BAM SSMax
0.88
0.92
0.88
0.86
0.80
0.62
0.30
0.10
0.02
0.00
0.00
RoPE SSMax
0.76
0.30
0.16
0.04
0.00
0.00
0.00
0.00
0.00
BAM SSMax
0.84
0.86
0.94
0.92
0.86
0.76
0.66
0.56
0.24
0.10
0.06
MultiKey
RoPE SSMax
0.80
0.86
0.68
0.06
0.00
0.00
0.00
BAM SSMax
0.22
0.16
0.12
0.04
0.04
0.00
0.02
0.02
0.00
0.00
0.00
MultiKey
RoPE SSMax
0.26
0.14
0.04
0.00
0.00
0.00
0.00
BAM SSMax
0.12
0.04
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
MultiKey 3
RoPE SSMax
0.18
0.00
0.00
0.00
0.00
0.00
0.00
MultiQuery
BAM SSMax
0.88
0.88
0.85
0.82
0.71
0.68
0.62
0.44
0.23
0.10
0.02
RoPE SSMax
0.85
0.76
0.03
0.03
0.01
MultiValue
BAM SSMax
0,96
0.94
0.92
0.92
0.84
0.76
0.70
0.38
0.14
RoPE SSMax
0.96
1.00
1.00
0.78
0.70
0.42
0.32
0.10
0.04
0.08
0.00
RoPE in all evaluated tasks, with an overall score 5 points above RoPE. We opted not to show the Long Structured Data Understanding task because it has only four instances under 131k tokens.
Table 6: LongBenchy2 Benchmark: GGD-BAM vs RoPE, IB parameter models.
BAM SSMax
RoPE SSMax
Code Repository Understanding Long In-context Learning
41.7
25.0
36.4
30.3
Long-dialogue History Understanding
35.0
35,0
Multi-Document QA
26.5
25.3
Single-Documeg QA
26.5
18.8
Overall
28.6
24.2
H ABLATION STUDY
H.1 INITIALIZATION
In this section we study two different initialization strategies for the shape 03. scale 0,, and location 0μ parameters of GGD-BAM. The first initialization is Laplacian that replicates ALiBi, setting 02 different 0g for each layer and 6, = O. The second initialization start from Uniform distribution prior 83- 0, which is middle ground between ALiBis Laplacian and 03 < 0, with 9. O.
Figure 10 shows us that ALiBi initialization provides least extrapolation lengths. Our best results were achieved by using the initialization scheme of 03 and 0,, O. this initialization equivalent to assigning Uniform prior to all tokens in the context.
H.2 LEARNABLE PARAMETERS
The GGD prior in BAM is parameterized by a shape parameter 03, scale parameter O.. and location parameter or.. In this section, evaluate how training different subsets of these parameters affects the performance of GGD-BAM. Each contiguration introduces different number of additional trainable parameters: training only 0g adds 192 parameters; training both fg and 0., adds 384; and training all three parameters (0g, 9... and f,)) adds 576 parameters to the model.
that allowing all parameters 0g. Ga and 6₂ be learned during training lowers model capacity to extrapolate. The best result was achieved when both 03 and 6,, are learn during training, showing that both parameters are important for context length extrapolation.
In Figure II we
21

---

1.0
ALBİI Uniform
0.8
Accuracyl (+)
0.4
0.2
100000
200000
300000
400000
500000
Sequence length
Figure 10: Comparison of passkey retrieval accuracy of models trained on context length 512 training
just 0g and O2, with three distinct parameter initialization schemes.
1.0
Bg
Bp and 6, 8g and 6,,
0.8
Accuracyt (.)
0p. 9, and
0.4
0.2
100000
200000
300000 Sequence length
400000
500000
Figure 11: Passkey retrieval accuracy models training just: 03: 02 and 0,; and 0z. On, and 0,, All variations were trained context length 512. Training 0g and 02 while fixing Og 0 yields more extrapolation length.
If we compare results on Figure 11 other PE in Figure we that worst combination of training all parameters is superior to all other PE in long context passkey retrieval.
H.3 TRAINING CONTEXT LENGTH
Here we repeat the experiments for models trained on context length of 1,024 and 2,048, double and quadruple the original context length of 512. We show detailed results for training context length of 1,024 both on perplexity and passkey retrieval when compared other PE methods. And we also compare passkey retrieval accuracy between BAM SSMax on those three distinct context lengths.
Essentially, the trend of ALiBi, RoPE Local and BAM SSMax being the only PE methods that are able to maintain low perplexity longer context is maintained, this is possible to identify in Figure 12 As expected, NoPE and Sinusoidal PE are the first models to exponentially increase perplexity.
RoPE SSMax improved it extrapolation performance and was able to maintain perplexity par with ALiBi until around 7,500 tokens. is worth noting that to achieve such results with RoPE, we expanded the R, manipulation post-training that performed by Nakanishi (025) from 50x to 100x. Without this RoPE SSMax would perform on par with standard RoPE.
Regarding passkey retrieval accuracy, again all PE methods except BAM SSMax struggles to access long context information and maintain high accuracy, Figure [3 shows this trend. Although the Figure 13 shows context lengths up 32,000, evaluated BAM SSMax until 51 512,000 context
22

---

10²
107
BAM(ours)
105
10a
BAM SSMax(ours)
NoPE
Perpexity=
10P
ROPE
Perpexit(++)
10a
RoPE Local
10
RoPE SSMax
10
Sinuseidal
10a
10a
10a
102
10
10
5000
10000 15000 20000 25000 30000
5000
10000
25000 30000 Sequence length
15000
20000
Sequence length
(a) Perplexity on Wikipedia
(b) Perplexity on Fineweb 10B
Figure 12: Log-scaled perplexity computed up 32x the training context length of 1024 tokens.
BAM SSMax, RoPE Local and ALiBi are able to maintain the lowest perplexity on longer contexts.
1.0
0.8
ALiBi
BAM(ours)
Accuracyt
0.6
BAM 5SMax(ours)
NoPE
RoPE
0.4
RoPE Local
RoPE S5Max
Sinusoidal
0.2
0.0
0
5000
10000
15000
20000
25000
30000
Sequence length
Figure 13: Passkey retrieval accuracy of models with distinct PE the Passkey Retrieval task. BAM and BAM SSMax outperform all PE methods. Specifically BAM SSMax is capable of maintaining perfect accuracy for context 32x the training context length.
length and it maintained accuracy above 80% until 300,000 and did not drop to zero throughout the evaluation. Beyond 5 112,000 context length, we did not have enough vram to perform evaluation.
In Figure 14 for longer contexts appears to make the model more robust to long context generalization as the accuracy tend to drop slower. The model trained with context length of 048 tokens generalizes all the way to 512k tokens while maintaining accuracy above 90%. The model trained with context length of 512 achieves 40% accuracy on 512k, showing a correlation between trained context length and retrieval accuracy at 512k tokens. Nevertheless, our model trained on context length of 512 is competitive with all the others until 300, 000 where others achieve higher accuracy.
how training in longer context lengths affects GGD BAMM Generally, training
H.4 SCALABLE SOFTMAX
Here we test all PE methods without SSMax. Figure shows the same trend models with SSMax, BAM outperforms every other PE method. It is worth noting that SSMax improves context length
23

---

1.0
0.8
Accuracy (-)
0.4
0.2
2048
100000
200000
300000
400000
500000
Sequence length
Figure 14: Comparison of passkey retrieval accuracy of GGD-BAM models trained on context length 512, 1,024 and 2,048.
generalization almost all the assessed PE methods. This shows that fading attention indeed one of the problems in long context extrapolation, however BAM is superior other PE on both cases.
1.0
ALiBi
BAM(ours)
NoPE
0.8
RoPE
Accuracyt )+)
RoPE Local
0.6
Sinusoidal
0.4
0.2
0.0
5000
10000
15000
20000
25000
30000
Sequence length
Figure 15: Passkey retrieval accuracy with distinct PE. BAM outperform all PE methods.
By comparing the results from BAM and BAM SSMax. we note that SSMax plays role in maintain- ing the PE ability extrapolate longer lengths. This trend is also maintained when comparing RoPE and its SSMax version. The fact that Scalable Softmax improves context length generalization both for BAM and for RoPE shows that good PE necessary but not sufficient for context length extrapolation. The softmax function tends to zero for longer context windows, which is a problem both for language modeling and for retrieval. counterbalance this effect, scalable softmax applies rescaling factor to the logits (Nakanishi) 2025), fixing that limitation.
Scalable Softmax introduces rescaling factor x In(n), where n is the size of the input vector and is learnable parameter. Note that has similar effect the normalizing scaler Z obtained when framing PE as a Bayesian mechanism. The only effective difference that Z should be a function of both query and keys whereas is learnable parameter.
To understand the effect of the normalizing scalar in our models across distinct scales, we show in Figure [6 how this learaalle parameter is distributed after training. We see that although the distribution appears to have a heavier tail in smaller models, the shape of the distribution similar across model scales.
24

---

3.5
Model Size
1.28
432M
2.5
Normra Frsquaryy
1.5
0.5
0.5
1.0
1.5
2.0
2.5
3.0
5SMAX Paremeeerr
Figure 16: Scaling factor in scalable softmax in three distinct model scales.
1.18
4.0
432M
120M
432M Ext.
120M Ext.
Loss
3.0
2.0
5000
10000
15000
20000
25000
30000
Position
Figure 17: Loss of GGD-BAM before and after lightweight fine-tuning for context extension.
H.5 CONTEXT EXTENSION
We observed during our first experiments that smaller models generalized better to longer context. We found out that bigger models are prone to overfitting the trained context length.
We devised a lightweight fine-tuning with context length 1024 for 256 steps (in comparison 512 during the beginning of the training) to make bigger models generalize once again. In Figure 17] we show that after lightweight fine-tuning the loss our models become more stable context beyond the training length of 512 tokens. This shows that, even if bigger models can overfit the training context length, a lightweight fine-tume to-ing procedure can make them generalize to extended contexts.
After context extension, performed the PassKey Retrieval analysis and noticed that bigger models benefit more from the context extension than smaller ones. The results obtained by this lightweight fine-tuning procedure shown in Table 7
In Table we see that the 120M model generalizes up to 512x the trained context length with accuracy above 0.8. However, bigger models struggle in much shorter sequences. When analyzing each model to their context-extended counterparts, we see that all model-scales benefit from this procedure. The 1.IB parameter model, however, has most improvement of context extension.
25

---

Table 7: Context extension effect across different model sizes BAM SSMax on PassKey Retrieval.
Model
<IK
51K
102K
153K
204K
256K
307K
358K
409K
460K
512K
120M
1.0
1.0
0.9
0.9
0.9
0.8
0.7
0.7
0.8
0.6
0.4
120M Ext.
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
431M
1.0
0.5
0.1
0.3
0.1
0.1
0.1
0.2
0.1
0.0
0.2
431M Ext.
1.0
1.0
1,0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.0
1.IB
0.9
0.0
0.0
0.0
0.1
0.0
0.1
0.0
1.IB Ext.
1.0
1.0
1.0
1.0
1.0
0.9
1.0
0.8
0.9
0.7
0.7
H.6 MODEL SIZE
Here we compare how BAM SSMax compares to RoPE SSMax on large scale model size (1.1B) on the Passkey Retrieval setting. Models were trained with context length 512 and prompted to perform Passkey Retrieval with up to 512, 000 tokens. Figure [8] shows that BAM SSMax also dominates RoPE SSMax across all evaluated lengths, performing accurate Passkey retrieval in all contexts that we were able to assess using our available compute.
1.0
Sequence
20%
0.91
40%
60%
0.95
0.95
-.2
1.00
1.00
1.00
1.00
1.00
0.90
1.00
0.80
0.90
0.70
0.70
0.0
5iK
102K
Sequence Length
(a) BAM SSMax
0.09
in Sequenoe
20%
40%
0.09
0.6
60%
00%
0.4
0.09
1.00
51K
102K
153K
204K
256K
307K
409K
460K
Length
(b) RoPE SSMax
Figure 18: Passkey Retrieval accuracy on 1.1B models. the bottom row and on the last column, we average accuracy across length and position, respectively.
This ablation corroborates that BAM is capable of using information across longer contexts than RoPE, and that such a conclusion generalizes across bigger models (and not just in 120M settings).
H.7 PE IMPACT ON INFERENCE PERFORMANCE
Here access how distinct PE strategies impact model throughput during inference. To perform this experiment, we initialize BAM and all the baselines trainable weights of four distinct model sizes and run 100 samples with batch size and sequence length 512.
26

---

In 8 that for smaller models the cost of performing Scalable Softmax dominates the results. This is imperceptible in bigger models, since the only impact of Scalable Softmax appears to be in the standard deviation. BAM does affect model inference time in comparison to other PEs. When we account for the standard deviation, every model has equivalent inference time.
Table 8: Inference time (ms) and vram (GB) during the backward pass for distinct model sizes.
120M
430M
1.1B
Time
VRAM
Time
VRAM
Time
VRAM
Sinusoidal
32.84 + 0.40
1.734
50.98
0.37
4,933
117.26 + 0.42
12.329
Sinusoidal SSMax
32.82 1.13
1.736
53.39 0.28
4.933
120.91 + 0.55
12.324
RoPE
35.77 1.98
1.754
53.42 0.38
4.936
1
12.330
RoPE SSMax
36.04
1.742
5..65
4.934
123.34
12.325
ALiBi
0.73
1.738
5.22
4.933
119.34
12.329
ALiBi SSMax
3.099
1.736
54.22
4.934
122.09 0.33
12.325
BAM
0.53
1.739
4.936
120.31 + 0.38
12.329
BAM SSMax
33.01 0
1.737
59.50
4.936
133.14
12.319
NoPE
1.87
1.738
56.10
4.933
12.329
NoPE SSMax
38.49 1.30
1.736
58.36 0.70
4.933
124.95 0.60
12.324
I
LIMITATIONS
Despite the theoretical and empirical strengths of BAM and its instantiation with GGD, our study is subject several limitations, which we acknowledge and discuss below.
Model Scale and Generalization to Larger LMs. Our experiments were conducted Transformer models with up to 1.1 billion parameters due to limited compute availability. While our results show improvements context length extrapolation at this scale, it remains an open question whether these gains persist or even amplify in very large language models. However, we note that this evaluation regime consistent with prior work in the PE literature, including those introducing ALiBi (Press et al. 2022) and NoPE (Kazemnejad et al. 2023), which also validated their approaches using similar-scale models.
Dataset Scope and Representativeness. Our empirical evaluation of perplexity is currently re- stricted to two datasets: FineWeb 10B (Penedo et al.] (0024), and Wikipedia (Foundation, 2023). While these datasets provide coverage of both large-scale pretraining and structured text, this coverage not exhaustive. Broader evaluations across additional domains-suce as code, long-form scientific documents- would be valuable for assessing the robustness and generality of BAM-hased priors.
Coverage of Positional Encoding Methods. Due computational constraints and the complexity of reimplementing certain positional encoding strategies, our experiments do not encompass all meth- ods proposed in the literature. We did not evaluate the T5 relative position bias approach (Raffel et al.) 2020). which requires bucketing mechanisms and distinct architectural modifications. Nonetheless, believe that the set of baselines considered coverlnnooeingggg absolute positional encodings (Sinusoidal), rotary encodings (RoPE), relative linear biases (ALiBi)), and content-only baselines (NoPF)-providee representative and diverse comparison to assess the context extrapolation capabilities of BAM.
Generalization to Instruction/Preeereen Tuned LMs. We did not evaluate BAM in the context of instruction-tuned or preference-tuned models. LMS often undergo additional fine-tuning stage, such as supervised instruction following, reinforcement learning from human feedback, or direct preference optimization, which can significantly alter attention dynamics and generalization behavior. It remains open question whether the context extrapolation benefits introduced by BAM preserved, attenuated, or potentially enhanced in such settings. Assessing the how models with BAM PE perform after instruction-uue -uned architectures important direction for future work.
27

---

J BROADER IMPACTS
Improving context length extrapolation in Transformers has the potential to reduce the computational and environmental costs associated with pretraining large language models. Because attention scales quadratically with sequence length, training with long contexts is prohibitively expensive, GGD-BAM enables models generalize longer sequences without requiring direct exposure during training, potentially lowering the need for long-context pre-training.
This efficiency gain could contribute more sustainable and accessible language model development, particularly for institutions with limited compute resources. Furthermore, better long-range rangerange general- ization supports important applications such as legal and medical document processing, educational content understanding, and scientific analysis. However, these capabilities must be accompanied by careful evaluation ensure reliability and safety in high-stakes domains.
K ALTERNATIVE INTERPRETATIONS OF Z
K.1 STATISTICAL-PIIISIC INTERPRETATION Z
Let f; = fcom (g. kj). 9 9pos (i.
We introduce three normalization constants:
Zeont = Ee' Zpoe = Σ Es
Zioint = Σ btts
The usual Gibbs Boltzmmnnn partition function is
Zjointt Σ exp(f, + 9z),
and the corresponding free energy is Fjoint
In Zjoint.
The factorization normalizer Z that restores Σ, softmax(fy + +9)) after factorizing exp(f)
g3) = exp(f,) exp(g;)
2 Zeont Zpos
[soltmax(f3) x softmax(g;)]]
Zejoint
Its log, In 7 = In connt + In 7pos In Zioonnt is precisely the interaction free-energy between the content and position potentials. Finally, introducing an inverse temperature ? yields
Pig exp(7 (f, + 9;)),
so y could control how sharply the attention distribution peaks.
K.2
GEOMETRIC INTERPRETATION OF Z
Recall
Peont (j) softmax (fetmax(fcottaeee (ge, k;)),
(i.))s( (Spomax(9posii.)
and let
(p(fcomt (qi, k, k,))) Ppos (p(gpos (i.j))) j=1'
Peont
Both vectors lie in the probability simplex An- fx R xz 2 0, Σ3 2y 1}. Then
=
Σ P(Scom (4i,k))) (csi),p((ppor Pl = (peont, Ppas),
where (Pconts Ppus) E, (offamlg k;)), p(opus k,))),((ppos (i,j)).
-Dot-product as overlap. (Peont. Ppos) E [0, 1] measures how much the two distributions "agree"-itt is maximal when they coincide and minimal when they are disjoint.
28

---

Norms of probability vectors. Since lllll2 I lpI|ı 1 for any p E A the raw dot-product not a true cosine similarity unless one divides by PP ll2 ll ll2. We omit that division because need Ej p{fcmm (q:, (q), (i. exactly to quantify the
normalization gap of the product of two softmaxes.
Re-normalization identity. The product distribution Peont C Ppos sums to (Peont, Ppos) # 1. Inverting that sum,
Z = Σ P(fcom(qss kj)), p(gpos
precisely restores Ej (oooora P(pos (1, (1,
Cosine-similarity caveat. If one instead defined cos 0 Pe laa Ppos that extra normal-
(prom Ppos?
ization would destroy the simple re-normalization identity needed for attention.
Thus (Peont has clear geometric meaning: it compensates for the overlap (or misalignment) between the content-based and position-based probability vectors, ensuring their elementwise product yields a valid distribution.
K.3
INFORMATION CHEMMATTIO-TTOORTIIIC THEORETIC INTERPRETATION OF Z
Let
Pjoint (i) Pij softmax([Jom) (gi, kj) + Spoe (i,
and recall the marginals
Peont (j) softmax (q:: kgy)). Ppus (i) softmax (Spoce (i,
The mutual information between content and position under the joint distribution is
Pjoina
[cont; pos) = Σ Pioint (i) In
p(Scom (qi, k,))))poo((i))
Since
etsta;
ets e
Pjoint (j)
and p(fcom( k.kk)
Zjoine
Zoom
Zpos
one shows directly that
Zcont Zpcos
I(cont; pos) In
In Z.
Zjoint
Hence
Z (so(t(cott:
which implies:
Z = if and only if content and position are statistically independent (I = 0).
grows exponentially with the amount shared information between content and position.
Thus Z can be seen as the "informationncouplinng"" multiplier that normalizes the product of the two marginals into the true joint attention distribution.
29


Reviews:
[{"review_id": "5OPUt5n85x", "reviewer": "Reviewer_xnWw", "summary": "This paper introduces the Bayesian Attention Mechanism (BAM), a theoretical framework that reformulates self-attention as a probabilistic model where positional encoding (PE) acts as a prior over token positions. The authors show that existing PEs like NoPE and ALiBi correspond to specific priors (Uniform and Laplace) and propose a new Generalized Gaussian prior (GGD-BAM) to improve long-context extrapolation. They validate the approach empirically on language modeling, passkey retrieval, and the RULER benchmark.", "questions": "See above.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper provides a clear and rigorous probabilistic interpretation of positional encoding in self-attention, unifying multiple existing approaches under a single Bayesian view. The derivations are mathematically consistent and offer new insight into the role of priors in attention mechanisms.\n2. Framing positional encoding as a prior elegantly connects methods like NoPE and ALiBi that were previously viewed as unrelated heuristics. This contributes to a more principled understanding of extrapolation in transformers.\n3. The visualizations of positional priors and learned β values provide interpretability, showing how certain heads specialize in long-range retrieval.", "weaknesses": "1. The main experiments are restricted to models with limited parameters. It remains unclear whether the extrapolation benefits persist at the scale of modern large language models (7B–70B), where positional behaviors often change.\n2. The work focuses on retrieval and perplexity but does not assess downstream tasks like long context QA or summarization, where long-context comprehension is critical.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9wzZqyAVcv", "reviewer": "Reviewer_giCK", "summary": "This paper proposes Bayesian Attention Mechanism (BAM), a probabilistic formulation of self-attention where positional encoding (PE) is treated as an explicit prior over token positions.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Disclaimer: I don't know the field of positional encoding. I've alerted AC for my lack of knowledge in the domain.\n\nThe method can enable retrieval for context length far beyond training context length with small number of parameters added, which seems nice to me.", "weaknesses": "Disclaimer: I don't know the field of positional encoding. I've alerted AC for my lack of knowledge in the domain.\n\nI'm unsure about the experimental analysis section of the paper. I don't know if the benchmarks used in the paper includes most of the popular ones in the field of positional encoding. I don't know if the paper uses enough baselines to compare their method with.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wOddoLzkJW", "reviewer": "Reviewer_R1bm", "summary": "This paper introduces a new positional encoding design, BAM, which is a variant and improvement of the existing ALiBi. \nThe authors reformulate current positional encodings from a Bayesian perspective and provide a clear motivation for their proposed modification. Experiments in long-context perplexity and in-context copying tasks (PassKey Retrieval and NIAH) demonstrate BAM's effectiveness in length extrapolation for language models with up to 1B parameters.", "questions": "---\nQ1: Can the authors include additonal baselines to better demonstrate the effectiveness of BAM methods?\n\nSome relevant methods are missing and should be discussed in the experiments. For example, the RoPE variants (PI, NTK, YaRN, etc), which build on RoPE, should be included .   \nThey are the unique benefits of RoPE design, enabling length extraploation either at little to no cost.\n\nMoreover, discussing the latest design in relative positional encodings [1] would further strengthen the case for BAM's effectiveness.\n\n[1] HoPE: A Novel Positional Encoding Without Long-Term Decay for\nEnhanced Context Awareness and Extrapolation.\n\n---\nQ2: Can the authors continual pre-train these models with more tokens?\n\nWe appreciate the authors's efforts in conducting these experiments.\nHowever, the training corpus used for the pre-training of these toy models remains limited.\nGiven this, conclusions drawn from such a limited number of pre-training tokens are not entirely convincing. \nIt is well known that ALiBi can accelerate convergence during training. \nSo, will BAM still exhibit its advantages when pre-training with more tokens (at least 100B tokens)? \nHope the authors can clarify that.\n\n---\nQ3: More comprehensive evaluation of the long-context abilities.\n\nThe existing experiments (PassKey and NIAH) provide some preliminary evidence of BAM's in-context copying ability for long contexts. \nHowever, BAM's other abilities in long-context is still underexplored, such as many-shot ICL, integration, and reasoning?\n\nBuilding on the resolution of Q1 & Q2, could the authors test BAM on real-word tasks of long-context benchmarks (e.g., LongBench and L-Eval)?\n\n--- \nQ4: What's the generative abilities of these PEs with normal lengths.\n\nWe acknowledge the authors' experiment in Appendix D.1, where RoPE shows lower perplexity than BAM on Wikipedia.\nWe agree with the authors that perplexity alone does not accurately reflect these models' ability.\nSo, what is the performance of these models at different scales, using different PEs, on language modeling tasks such as ARC, Hellaswag, PiQA, etc.?\nCould the length-extrapolation ability of BAM emerge at the cost of its performance on shorter texts?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The perspective on improving ALiBi is both interesting and well-motivated.\n\n- The authors provide a theoretical analysis that explains their motivation and supports the proposed approach.\n\n- The authors conduct extensive experiments to demonstrate the effectiveness of their method, providing compelling evidence for its validity.", "weaknesses": "- Some experimental results are not convincing enough to support the authors’ contributions (See Questions).\n\n- The presentation of this paper could be improved. \nThe section on theoretical proofs somewhat detracts from the main focus of the work. \nWhile naming the new PE \"Bayesian Attention Mechanism\" is justifiable, it make sense but feels somewhat forced. \nSeveral important experimental conclusions are scattered throughout the Appendix, which makes it challenging for reviewers to follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5R8UrRvrhL", "reviewer": "Reviewer_b3U5", "summary": "This paper proposes the Bayesian Attention Mechanism (BAM), framing positional encoding as a Bayesian positional prior that unifies methods like NoPE and ALiBi and introduces GGD-BAM, achieving strong long-context retrieval up to ≈256K tokens with <1000 extra parameters and improved stability when combined with SSMax on models from 120M–1.1B trained on FineWeb10B/Wikipedia.", "questions": "1. Is there an empirical correlation between the normalization $Z$ and long-context extrapolation, and can $Z$ be optimized or learned to extend coverage?\n2. Does the optimal $β$ shift with larger training contexts (e.g., 2048), and can you provide experiments across different training-context regimes?\n3. Why do larger models (1.1B) tend to overfit the training context more, and can you analyze the representation or optimization dynamics behind this?\n4. Could a per-layer or learned $β$ schedule better balance local vs. long-range dependencies and reduce local-context suppression when $β<0$?\n\nOverall, while the idea is conceptually interesting and well-presented, the current empirical evidence is insufficient to justify acceptance at this venue.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Formalizes positional encoding as a Bayesian prior, providing a useful theoretical advance.\n- Demonstrates robust long-context retrieval (≈256K tokens) with minimal parameter overhead.\n- Presents ablations and visualizations that enhance interpretability.\n- Supplies clear experimental documentation and reproducibility details in the appendices.", "weaknesses": "1. Experiments stop at 1.1B parameters, leaving scaling behavior at 7B/13B+ unknown.\n2. Evaluation is narrow (FineWeb, Wikipedia, passkey retrieval, RULER–NIAH) and omits broader benchmarks and domain corpora.\n3. Training overhead analysis lacks detailed wall-clock and memory cost measurements for ultra-long fine-tuning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)