TRA: Better Length Generalisation with Threshold Relative Attention
Authors: Mattia Opper
Roland Fernandez
Paul Smolensky
Jianfeng Gao
Venue: Accepted by TMLR ()

Abstract: Transformers struggle with length generalisation, displaying poor performance even on basic tasks. We test whether these limitations can be explained through two key failures of the self-attention mechanism. The first is the inability to fully remove irrelevant information. The second is tied to position, even if the dot product between a key and query is highly negative (i.e. an irrelevant key) learned positional biases may unintentionally up-weight such information - dangerous when distances become out of distribution. Put together, these two failure cases lead to compounding generalisation difficulties. We test whether they can be mitigated through the combination of a) selective sparsity - completely removing irrelevant keys from the attention softmax and b) contextualised relative distance - distance is only considered as between the query and the keys that matter. We show how refactoring the attention mechanism with these two mitigations in place can substantially improve generalisation capabilities of decoder only transformers.

Published in Transactions on Machine Learning Research (06/2025)
TRA: Better Length Generalisation with Threshold Relative Attention
Mattia Opper University of Edinburgh'
m.opperded..c.cuk
Roland Fernandez
rfemannd@miccrooofftccom
Microsoft Research
Paul Smolensky
pamo@miccrosoft.com
Microsoft Reseorch
Johns Hopkins Universily
Jianfeng Gao
iggo@miiroosoft.com com
Microsoft Research
Reviewed on OpenReview: https: openreview. net/ forum? id- BUCBNNMM
Abstract
Transformers struggle with length generalisation, displaying poor performance even on basic lasks. We test. whether these limitations can be explained through two key failures of the self-attention mechanism. The first is the inability fully remove irrelevant information. The second is tied to position, even if the dot product between key and query is highly negative (i.e. an irrelevant key) learned positional biases may unintentionally up-weight such information dangerous when distances become out of distribution. Put together, these two failure cases lead to compounding generalisation difficulties. We test whether they can be mitigated through the combination of a. selective sparsity completely removing irrelevant keys from the attention softmax and b) contextualised relative distance distance is only considered as between the query and the keys that matter. We show how refactoring the attention mechanism with these mitigations in place can substantially improve the generalisation capabilities of decoder only transformers
1 Introduction
Transformers Vaswani et al. 2017 have sparked a revolution in artificial intelligence. When pre-trained at scale, these models display impressive broad-spectrum capabilities, particularly the domains of natural language and code (Bubeck et al., 2023 Dubey al. 2024). This inchudes the ability to perform complex abstract symbolic processing through in-context learning (TCT.) (Brown al., 2020 Sinha et al., 2024 Smolensky et al. 2024h. On the other hand, those abilitics far from robust (Dziri et al. 2023 Mirzadeh et al.] 2024) and display dramatic failure modes such self-contradiction and hallucination (Liu et al., 2023a; Ji et al. 2023 Guerreiro et al., 2023). Similar failure modes are also observed in controlled synthetic settings where the correct solution should be trivial. Decoder-only transformers consistently display poor
generalisation to unseen lengths and dependency distances (Zhou et al. 2023 Liu al., 2023a
Addressing these issues can provide crucial benefits regarding both long context utilisation and length gen- cralisation more broadly. Firstly, current approaches generally involve training at certain sequence length and then adapting the positional encoding to longer context. Most commonly this entails additional fine- tuning for best effect, thereby incurring further compute cost penalty. Even then current SoTA foundation
Work completed while Mierosoft Research Pytorch code for core mechanism Appendix

---

Published in Transactions on Machine Learning Research (06/2025)
models still struggle with long context (Hsich et al. 2024). Ideally, we would want solution that simply works out the box. Put another way; if the transformer is in fact learning circuits Meerrill et al. 2022 Olsson eet al., 2022; Liu et al., 2023b), then their correct execution should be general purpose and robust, rather than contingent on specific quirks of the training data.
To understand how to get there, we should first look at the core capability that the attention mechanism was supposed enable in the first place, namely the ability to selectively retrieve and propagate information (Bahdanau et al., 2014 Luong et al. 2015 Vaswani al. 2017). We can break this capability down into three broad categories. (I) Purely positional: attending to ones local neighbourhood (i.e. contextualisation within sentence or phrase). (2) Purcly based on content: attending to the most relevant item irrespective of position. (3) Content followed by position: of the relevant items apply a positional bias (e.g. select the most recent). (3) is a crucial capability, it enables tracking the state an entity over time, and is a pre-requisite for robust reasoning. Our contention is that standard attention mechanisms can perform (1) and (2), but struggle with (3). This due to the way that semantic content and positional information combine. Position is generally treated independent of content. Which means that favouring one must. come at the expense of the other. This leaves the two types of information in conflict, fundamentally limiting self-attention.
Our solution involves a two step process. We first mask out irrelevant keys based on threshold applied to the raw attention weights, and then compute the relative distance but only as between the query and the keys that remain. We call this mechanism Threshold Relative Attention (TRA). TRA is a simple, modular addition that allows self-attention to learn general solutions to all three behaviours outlined above. We find that TRA's introduction leads to considerably improved length generalisation on synthetic benchmarks. As well robust, and even improved, language modeling perplexity for out-of-distribution sequence lengths, even up to an increase of 32x. Our findings contribute to the development of more robust attention mechanisms.
2 Model
2.1
Background: Relative Positional Encodings
Relative positional encodings DDai et al. 2019 Raffel et al. 2019) decompose the allention operation into two steps. First the standard QK" dot product is calculated, but this purely based on semantic relevance. Let's call that matrix S (semantic). This is then added with a learned positional bias matrix D (distance) to produce the final attention weight.
S QKT
(1)
Vdz
A = softmax (S + D)
(2)
Our contention is that a central cause of length generalisation issues stems from and D being treated as independent of each other. Tf D fit. the training distribution, then distribution shifts will necessarily begin to introduce errors. S more robust in this regard, because relevance isn't impacted by changing lengths. We want to make D more robust by making it contingent on S.
2.2 Contextual Relative Distance
We want to condition D on S. In other words we only want consider position over keys that matter. To achieve this we first introduce selective sparsity. We will allow certain keys to be fully removed from S, but do not force it to be one hot. Instcad we introduce threshold and climinate all keys which fall below it. Lets call that matrix S'. The simplest way to get S' is by applying ReLU.
QKT
S' ReLU
(3)
Vit.
2

---

Published in Transactions on Machine Learning Research (06/2025)
Next calculate boolean mask M that is true only for relevant positions:
Mrs { 1 ,
(4)
otherwise
We going to compute contextualised distance matrix, but only consider the relative distance to keys which make it through the threshold, this give us the contextualised relative distance. To do this consider the following example M, where again only keys which meet the threshold are assigned 1:
1
0
0
0
0
M
1
1
0
1 0
Our threshold mask M is m x n matrix, which we are going to use to get our contextualised relative distance matrix D, by applying the cumsum operation in the right to left direction (i.e. producing relative positions given causal masking). More formally we can say that for row index and column index j:
Do Ex Ma
(6)
For our example M from (5) this would yield the following output, noting that we omit the keys which do not meet the threshold in this example for clarity:
0
0
1
0
0
0
D
=
0
2
(7)
2 0
0
D now contains distances which are informed by context. For example, for the query in position two it views the key in position one as the closest possible item rather than its own key, which has been deemed irrelevant through the threshold operation. Similarly, for the query at position position three as the closest possible item and the key position as the next nearest, skipping over the irrelevant intervening keys. This enables operation (3) from the introduction ie. considering ordering but only over the relevant items. A capability which more traditional PE methods such as RPE (Dai et al., 2019), absolute (Vaswani et al. 2017) and RoPE Su et al. (0221) are unable to accomodate because they treat position as fined property rather than as adaptive to the context dependent needs of the query.
will view the key at
2.3 Parameterising D
Now we have arrived at contextualised relative distances, but still need some way of turning them into actual weight. In order to do utilise a parameterised forget gate. Lets denote the residual stream of size E for position i as with W, E Fx1 as the forget projection, b as a scalar bias, and o as the sigmoid non-linearity. Then the weight for each position is given as:
6, a (Wya, + b)
(8)
D/ = 8.
(9)
As each 8, is a sigmoided scalar value, the more times it multiplied with itself the smaller it gets. The closer &g is to the longer the decay takes, the closer to zero the shorter. This provides the model temporal
3

---

Published in Transactions on Machine Learning Research (06/2025)
memory. Enabling it to forget irrelevant past information if necessary, to treat all timesteps uniformly if temporal ordering is not a concern. We opt for a recency bias for several reasons. Firstly, it has been generally shown to be useful when modeling natural language (Shen et al. 2018 Zhu et al., 2021 Clark 2025 In particular ALiBi (Press et al., 2021 shows that a non-parametric recency bias significantly
cet al. improves length generalisation
translormers. However, though their non-paramelric version makes this constraint rigid. By choosing to parameterise instead we allow the model completely disregard position if needed, the benelits of which are demonstrated by Lin et al. (2025). Secondly, is unclear whether token occupying the ith position provides useful information beyond causal ordering. The first occurrence can be identified by proximity to the start of sequence token. While middle positions are better differentiated via their unique semantic context, a capability which has been identified and carefully studied in transformers by Ebrahimi et al. (2024). Finally, a temporally decaying memory is a powerful mechanism. Merrill el al. (2024) show that is better able to approximate sophisticated NC circuits given finite depth, compared to regular transformers. Morcover, is the mechanism which has empowerod the resurgence state space models (Peng et al. 2023 Orvieto et al. 2023: De al.] 2024), recent successful recursive networks (Opper & Siddharth 2024a). and has shown considerable promise with transformers (Csordás et al. 2021
Lin et al.
2025
2.4
The Final Attention Weights
We compute the final attention weights A':
Å = S' + log((D))
(10)
-lell
if Mij = 0 otherwise
(11)
Aij
A' softmax (A)
(12)
This gives us the full TRA mechanism. We log( log(D") for improved numerical stability and weight scale following Lin et al. (2025). Note keys which do not meet the threshold are completely removed from the final softmax by setting them to a large negative value (-1e11), as well as not counting towards distance as dicussed previously. This means that TRA exhibits both selective sparsity and allows semantic content to determine, and consequently synergise with, positional bias.
3 Experiments
3.1
Tasks
Our core hypothesis is that length generalisation failures occur in transformers because the standard attention mechanism puts semantic and positional information in conflict. TRA is designed to mitigate this, To test our hypothesis, we first turn to controlled synthetic tasks, followed by language modeling.
Flip-Flop Language Modelling: Introduced by Liu et al. (2023a) Flip-Flop language modelling is a algorithmic reasoning benchmark designed to test transformers ability to glitchlessly handle sequential de- pendencies. Sequences consist of symbol alphabet of three instructions: write (w), ignore (i) and read (r) each which is followed by a bit. To solve the task the model has to recall the bit that follows the nearest write instruction while ignoring all irrelevant information (e.g. in: w01101111101110001111 ing sequences will lead to failure. The model has to learn program to succeed. The ability to model a flip-flop is the foundation for all syntactic parsing and algorithmic reasoming capabilities (it necessitates the propagation of information according to abstract rules), and is pre-requisite to being able to maintain a state hicrarchically (Merrill et al. 2022 Liu et al. 2023b). The task consists of threc test sets: IID, and two out of distribution sets that vary the number of intervening ignore instructions. OOD sparse increases the number of ignore instructions and requires the ability to handle increased dependency distance. OOD dense decrcases the number of ignore instructions and consoquently probes whether the model retain focus in the presence of an increased number of attractors (i.c. write instructions). Input examples consist of strings of length 512 with the probability of an ignore instruction being generated set to 0.98 in the sparse set and 0.1
out: 1). Memoris-
4

---

Published in Transactions on Machine Learning Research (06/2025)
in the dense set. The criterion for success is perfection (i.e. 100% accuracy across all test sets). The authors find that transformers exhibit reasoning errors across wide range of architectural variants and that the issue is scale-tneariant, persisting even up to GPT-4.
Induction Heads: This is a test of associative memory, Given a string of unique symbols the model must learn the following reasoning pattern: a + b ... a ->? ? out: b, otherwise known as an induction head. This ability has been shown to emerge for specific heads and for specific input-output pairs after exposure to roughly two billion tokens of pre-training data (Olsson et al., 2022), and is hypothesised to be crucial mechanism for in-context learning. However, when tested in controlled synthetic settings, transformers' ability to learn the generalised pattern has been limited. [hou et al. (2023) show that models are only able to weakly generalise (from length R44 to 50) despite the fact that they should in theory easily be able to learn the circuit. Beyond basic positional reasoning the induction heads task is also a test of the attention mechanisms capacity for specificity. As length increases the number of irrelevant features grows and the model must be able to sharply discriminate so as select the correct successor.
Copy: A more difficult version of the induction heads task is copying non-unique strings. Here evaluate using a vocabulary of 10 symbols. The model must learn to absorb the history for each token in order to predict the correct successor. In out of distribution length settings difficulty is increased because the level of history required to distinguish between potential successors increases. Tt is also a task where our current models of transformer programs, which treat symbol manipulation as fully discrete (i.e. one-hot selection), cannot find a solution (Weiss et al.) 2021 Zhou et al. 2023), and therefore predict generalisation failure. Copying is a test of the models capacity to maintain memory. A given tokens capacity to memorise its past determines successful disambiguation of its successor.
For all lasks we use full sequence accuracy as the criterion (exact match). For the copy and induct tasks we train on sequences with input length 50 and evaluate on buckets of increased OOD lengths up to 300.
3.2 Baselines and Experimental Setup
For both TRA and the baselines we use the same core transformer backbone based on the Llama 3 variant of the architecture (Dubey et al. 2024); consisting on SwiGLU layers Shazeer (2020) coupled with RMSNorm
(Zhang & Sennrich 2019). We compare TRA with the following basclines:
No positional encoding (NoPE): Kazemnejad et al. 2023) claim that decoder only transformers are implicitly able to represent positional information and that explicitly encoding it is unnccessary. Our first baseline is therefore simply standard causal attention with no positional embedding.
Absolute positional encoding (APE): Positional information for each index is represented by a learned embedding which is added to the residual stream at layer zero This was the standard in the original variant (Vaswani et al. 2017) as well as GPT-2 [Raffrd el. al. 2019). Lengths not encountered during training will not have the corresponding embeddings trained.
Dai et al. 2019 and also adopted by T5 (Ralell et al. 2019). Relative Positional Encodings model position via a learnable additive bias term. Position is considered as to the relative distance of key to the query. OOD distances are all assigned the same learned value for maximum length.
Relative Positional Encodings (REL): Introduced by Transformer-Ne XL
Rotational Positional Encoding (RoPE): Introduced by Su et al. 2022 This approach models posi- tional information by applying a rotation to the keys proportional to the distance from the query. Recent work by [Barbero al. (2025) demonstrates that rather than encoding gradual distance based decay RoPE actually learns either fully positional attention (attend to predecessor or attend to first token) or full semantic attention. One coming at the expense of the other.
Label Encoding (Label): Introduced by Li & McClelland (2023) these seek to mitigate the OOD issues faced by APE by randomly selecting and then sorting indices from a range greater than sequence length. The authors demonstrate that this aides length generalisation under a limited OOD sctting (I25 -> 50), with the contention that this is due to the encodings learning to respect relative ordering rather than absolute value. The same mechanism is also proposed in concurrent work by Ruoss et al. 2023)
5

---

Published in Transactions on Machine Learning Research (06/2025)
Contextual Positional Encoding (CoPE): Perhaps conceptually the most related to TRA, CoPE (Golovneva. et al. 2024) utilises spikes in the S matrix in order to calculate positions. Tokens are as- signed fractional positions by interpolating betweeen a fixed set of absolute embeddings. Unlike TRA it does not completely remove irrelevant items from the softmax and uses an absolute rather than a relative scheme. This incurs further computational cost. because each query in each head needs to be matrix multiplied with a layer specific absolute positional embedding matrix. The authors attempt to mitigate this issue by limit- ing the number of absolute positions to a small number. TRA allows for similar behaviour but in a more computationally efficient manner, and allows removal of noise through its thresholding mechanism.
Forgetting Transformer (FoT): The recent Forgetting Transformer (FoT) (Lin et al., 2025) utilises a forget gate as its positional encoding. The crucial difference between it and TRA is that the forgetting transformer uses standard relative distance compared with TRA's contextualised alternative. This means that it can either enforce a recency bias or switch off any positional information completely. Crucially, what it cannot do is synergise the two (capability (3) from the introduction, which preciscly what TRA's contextualised relative distance enables. FoT therefore serves as the truest. test as to whether the issues we hypothesise regarding position are valid.
Differential Transformer (Diff): Introduced by Ye et al. (2024) the differential transformer utilises a twin heads approach where an auxiliary attention head looks to perform noise reduction by down-weighting attention to irrelevant keys. It tests to what extent simply the ability to remove noise from the attention dot product is beneficial- without any further considerations regarding position.
Hyper-parameters: For all models use the mini configuration from Ture et al. (201)) : four heads, four
layers and 256 dimensional embedding size. The MLP is set 2x embedding size for both the linear and gate hidden units. Our focus is on small models following prior work (Zhou et al. 2023). Furthermore, attention glitches have been shown to persist at scale (Liu et al. 2023a). and the true solution for these tasks should not require additional layers. We use a linear warm-up for the first 5% of steps coupled with cosine decay.
Other task and model specific pyper-parameters can be found in Appendix
3.3 Results
Table 1: Results represent the average across four random initialisations. Metric is full sequence accuracy (exact match). The numbers next to Copy and Induct represent input length ranges.
Model
NoPF
APF
REI.
RoPE
Tabel
FrT
Diff
CoPF
TRA
Flip Flop IID
100
100
0.0
0.0
100:
100 100=0.0
100
100 0.0
100
Flip Plop 0OD (Demee)
0.15 0.0
27.38 18.1
11.2:447.2
100 0.0
12.58 + 15.87
100
100::00 0.0
100 0.0
100 = 0.0
Flip Flop OOD (Sparee)
0.1
5.61
70.48 10.92
72.82 1.27
1 9.93
9..11:221
6.78
9.
100=0.0
Induct IID (0-50)
100 = 0.0
100 +: 0.0
100 0.0 0.0
100 0.0
99.96 +: 0.04
100 = 0.0
100 + 0.0
100 +: 0.0
100 = 0.0
Induct (50. 100)
13.42 1.66
0.0
11.87
45.36 43.99
15.00 2.83
4.43 0.0
20.12
100
Induct OOD (100-200)
0.0 0.0 0.0
0.0 0.0
0.0 +: 0.0
0.0 + 0.0
0.0 0.0
0.0 0.0
0.0=0.00 0.0
0.0 0.0
99.90 :: 0.0
Induct 0OD (200-300)
0.0+0.00 0.0
0.0 0.0
+ 0.0
0.0 0.0
B.0 0.0
0.0 0.0
0.0 0.0
0.0 0.0
99.33 + 0.0
Copy IID (0-50)
100 - 0.D
100 0.0
100 + 0.0
100- 0.0
100 +00++00 0.0
100 - 0.0
100 +
100 + 0.0
100 - 0.0
Capy 0OD (50-100)
12.43 = 1.66
0.0:0.0
3.1 2.5
1/44 4.79
26.97 : 6.11
97 17.65-2 : 4.7
1.61 0.29
86,47 : 21.17
100 = 0.0
Copy DOD (100-200)
0.0+0..0
0.0-00.0
0.0 + 0.0
0.0 + 0.0
0.0 0.0
66:04 + 40.49
0.0-0.0
40.80 + 31.27
99.87 + 0.14
Copy 0OD (200-300)
0.0
0.0 0.0
0.0 0.0
0.0:000
B.0 0.0
3.54 +: 4.35
0.0=0.0
2,08 :: 2.14
98.16 1.82
Results are in Table [ We draw the following conclusions:
Most PE methods struggle OOD: Absolute and relative positional embeddings consistently struggle out of distribution scttings. Positional information attached to specific indices. Best case they learn a bias in a particular direction (recency or the opposite), but even then this bias is input solely index dependent and only applicable for IID indices. Consequently, they struggle with OOD settings. RoPE also appears to learn a strategy which is fit to the training distribution and therefore struggles with out of distribution settings. This issue extends even to the differential transformer as it is based on RoPE. Consequently, despite having a more flexible selection mechanism it is still limited by its underlying positional encoding.
NoPE Generalises Wcakly: The NoPE hypothesis (Kazemnejad et al. 2023) states that decoder only transformers do not need positional embeddings because they are able approximate a counter. The theory is that they do so by selecting token which serves as anchor and sends some signal via the value message.
6

---

Published in Transactions on Machine Learning Research (06/2025)
Intervening tokens send a mull signal. The position of given token is then determined by the extent to which the anchor token's 's signal is diminished by the softmax. From Table 1 we can see that this strategy generalises better than absolute, relative and rotary PE, consistent with the author's findings. However, it still docs very poorly overall. NoPE fails at copy and induct indicating that when the anchor signal becomes too diluted due to increased length it fails to convey position accurately. The total failure of NoPE on the flip flop dense OOD set indicates issues when there are number of possible anchor tokens in close proximity to one another. Too many anchor signals lead to errors in the count. Results show that NoPE is insufficient for true generalisation.
Flexible PE Yields Improvements but is Insufficient: The strongest contenders of the basclines utilise more flexible positonal encodings. Label, which randomly selects from a larger set of positions, is the weakest among them, but performs better than simply using absolute positions. However, alone the approach is insufficient for strong generalisation. FoT performs well on the copy and flip flop tasks, but is not able to fully generalise to either. Similar performance patterns are obscrved with CoPE. However, both methods struggle to generalise to the induct lask which requires as its first step sharp attention to the predecessor token. This operation, also known token shift (Peng et al.) 2023). is purely positional and requires the token to completely remove attention from itself and then apply a sharp recency bias to capture the predecessor. TRA can use its threshold mechanism to perform this operation while CoPE and FoT lack this capability and consequently suffer.
TRA succeeds: it is able to robustly generalise to all tasks. Its performance on copy and induct demonstrate strong capacity for memory and associative recall. Its ability to solve flip-flops shows that it can combine positional and semantic information in complementary fashion. To the best of our knowledge it is the first transformer network to fully solve the flip-flop benchmark of Liu et al. (2023a). It achieves 100% accuracy across seeds, which till this point has been only been achievable by LSTM networks. However, both FoT and CoPE come very close to full generalisation To further tease apart performance differences on this core capability we generate an additional synthetic test in the same spirit as the original benchmark.
3.4 Flip-Flops++
Task: To further test the differences between contextualised positions and the bascline approaches we create a further synthetic test which we refer to as Flip-Flops++. The task is a follows: the input consists of instruction followed by sequence of random symbols combined with a trigger symbol. The set of instructions are as follows: return the symbol that occurs immediately after the first occurvence of the trigger, immediately after it, and immediately before or after the last occurence. For example, if we use the trigger symbol example could consist of: (instruction: before first, sequence: bexakleaztyab, out: x). The training set contains examples from each of the instructions and sequence lengths of 0-50. The test set consists of of sequences of lengths 51-500. This is a more challenging scenario than the original benchmark as dependency differences can be starker than in the original sparse set, and the model is required to behave differently depending on the instruction. Moreover, bocause we include having to retrieve from the first occurrence among the instructions it lets us investigate whother techniques which employ a recency bias are limited to only being able to allend to the most recent relevant item. For training we employ the same hyperparameters as before.
Table 2: Length generalisation results by instruction on Flip-Flop+++, taken over four random initialisations. Metric is full sequence accuracy (exact match).
Model
NoPg
APB
REL
RoPE
Label
EVT
Diif
CoPE
TRA
Match 64.50 19.36
18.37
95.78
5.1
09:27
96.50 à 3.91
18.08
92.01 4 2.07
0.0
95.64 =: 4.87
After Last Match
20 + 4.0
11.77 + 1.84
25.16
6.09
38.06
7.09
37.58
10.98
65.0
22.86
27.74
+ 1.82
89.32
10.21
99.84 0.28
Before Match
58.23
71.17
à 21.87
96.68
87.94
93.53
2.39
0.0
98.97 =
Before Last Match
13.45 1 0.69
10.12 1 2.21
27.06
9.02
33.86
15.01
31.31 = 12.97
76.12 4 13.72
23.25 1 4.28
91.3 20.3=00 6.32
100 1: 0.0
Results are shown in Table As before most PE methods struggle to generalise, though RoPE performs quite strongly with first match, echoing the findings of Barbero et al.) (2025). However, in this more challenging setting, only the contextualised schemes, CoPE and TRA, are able to robustly generalise. Though TRA performs better at last match while CoPF holds slight edge at first match, due to their use of different
7

---

Published in Transactions on Machine Learning Research (06/2025)
positional schemes. For natural language it is likely that the former is preferable as it better reflects causal structure (i.e. it allows understanding of an entity's current state). Notably FoT, which in the original Flip- Flop benchmark of Liu et al. (2023a) appeared to be strong contender, suffers substantial performance degradation in this more challenging sctting. This because, as we hypothesised in the introduction, is unable synergise positional and semantic information and consequently fails generalise to capability (3)
from carlier.
3.5
Why does contextualised relative distance matter?
In order to analysc the impact of contextualised relative distance we train two layer one head TRA and FoT models on the flip-flop language modeling task and compare the attention heatmaps for the final layer in Figure 1 TRA fully generalises in this limited setting, the only model to be able to do so out of all the baselines, and consistent with the optimal solution demonstrated theory by Liu et al. (2023a). On the other hand FoT fails, and the reason why clearly visible by contrasting the attention weights of the two.
(a) TRA final layer attention in FFLM task. Black means null attention weight.
26 Wo
mockiralie gic
has 402 don AK
(b) FoT final layer attention in FFLM task. Black means null attention weight.
Figure 1: Contrasting Attention Heatmaps between TRA and FoT. TRA synergises semantic and positional information while FoT must trade one for the other.
Figure la shows how TRA is able to combine both semantic content with relative positions. It first filters out irrelevant keys fully and is then able to apply strong positional bias to the remainder that truly matter, allowing it to cleanly select the most recent write instruction. Its recency bias Ta centre) is applied in a shifted manner only to what is relevant, which allows express a strong positional preference. FoT on the other hand [Tb), has no such mechanism and as a result must diminish on semantic content (contrast the strength of lb centre with Ia centre). This fact, coupled with its inability to completely remove irrelevant tokens, mcans that it is unable to handle long distance dependencies and consequently cannot learn semantic and positional preference such that they complement each other. We use FoT here as clean illustrative example, but note that the same failure mode must be true of all schemes that do not employ contextualised distance.
recency bias to be able focus
3.6 Language Modeling:
TRA enables improved generalisation in controlled synthetic settings, which we attribute its capacity for contextualised relative distance. What happens when we expose it more complex data like natural
8

---

Published in Transactions on Machine Learning Research (06/2025)
language? Investigating this question has two motivations. Firstly, contextualised relative distance requires selective sparsity to operate. This operation cuts the connection between columns and we must make sure that this does not harm learning in more complex domains by potentially limiting exploration. Secondly, allows us to investigate whether the generalisation patterns observed in syuthesis also occur in natural language. For our experiments we turn to the WikiText-103 benchmark (Merity et al. 2016). which consists of full Wikipedia articles comprising circa 100 million tokens, probing both scale and ability to handle long distance dependencies.
Hyperparameters: For language modelling increase model size to the medium configuration of Ture ct al. (2019); Opper et al. 2023a). This is cight layers, cight heads and a 512 dimensional embedding size. The MLP is set 2x embedding size for both the linear and gate hidden units as before. We use the GPT-2 tokenizer [Radford et al., 2019) which leads to a vocabulary size roughly 52k. Totaling circa 80 million parameters for both TRA and the baselines. We train using: window size 128, batch size 64, 100k steps. Our schoduling regime remains the same as before. Our evaluation consist of two parts. In distribution: we measure perplexity on the test set using the training window of size of 128. Out of distribution: increase the window size to 4096 and observe the extent to which perplexity changes with increased length.
Table 3: Test Perplexity WikiText-103 (4). Mean and standard deviation taken over four random initalisations. Models trained for 100k steps with context window of 128 lokens.
Model
NoPE
APE
Rel
RoPE
Label
FoT
CoPE
TRA
Perplexity
31.61 0.16
32.09
30.20
70
TRA
CoPE
NoPE
RoPE
60
Diff
Label
APE
55
Reletive
Training
Perplesity
40
25
500
1000
1500
2000
2500
3000
3500
4000
Token Position
Figure 2: WikiText-103 Test Perplexity (L) OOD sequences lengths. Models trained for 100k steps with window size of 128, and evaluated with window size to 4098. Results taken over four random seeds.
Results: In-distribution results shown in table 3 TRA remains competitive with the strongest base- lines. This means that performing model in-distribution is the differential transformer (Ye et al. 2024p. which contains more sophisticated mechanism for relevancy determination. Replacing the simple ReLU we use here with such a more sophisticated approach could prove a fruitful avemue for future work. The out of distribution case is shown figure / Most PE methods very rapidly begin to decline in quality, with similar degradation patterns emerging compared the synthetic settings. The three chief exceptions to this rule are TRA, CoPE and FoT. CoPE is able to maintain/slighhle improve perplexity for up to roughly 10x increase in
introduce selective sparsity without hurting learning. Notably the strongest

---

Published in Transactions on Machine Learning Research (06/2025)
context length, before performance begins to deteriorate. We believe this is because its use of fractional positions only temporarily alleviates generalisation difficulties, but ultimately leaves the core issue unad- dressed. Furthermore, increasing the number of positions in CoPE may delay degradation further, but any increase comes with a heavy computational burden and consequently makes CoPE difficult to scale as the sole mechanism for representing position. On the other hand, both FoT and TRA are capable of strong length gencralisation (up to 32x greater than training length). Moreover, they display a promising trend of improved perplexity with increased length which indicates that they are able to make use of increased context rather than simply being robust to it. Between the two, TRA demonstrates a consistent perplexity edge compared with FoT. We believe this is due to its improved capacity for handling out of distribution dependency lengths as demonstrated by the flip-flop language modeling results. In sum, we show that TRA can introduce selective sparsity to the allention weights without compromising performance, and that its improved generalisation capabilities in syuthetic settings also appear to be mirrored in the more complex setting of natural language. Finally, to provide an intuition for how TRA prefers to use its attention and forget gate we provide an analysis the attention heatmaps learned by our language models in Appendix C
4 Related Work:
Generalisation in Neural Networks: Neural networks, and particularly transformers, struggle with out of distribution generalisation (Hupkes et al. 2019
2023). Partially this is architecture dependent. Significant efforts have been made towards creating architectures that enable a more flexible information flow and consoquently have access to greater space of possible solutions (Fan et al.] 2020 Chowdhury & Caragea 2021 Opper et. al. 2023b Opper & Siddharth 2024b Soulos et al. 2024 DuSell & Chiang 2024p. However, these approaches tend to come at the cost of either flexibility or parallelisability the key properties which made transformers, and particularly their decoder only variants, take over the world.
Transformer Programs: A parallel effort has been made to understand what kind of finile depth solutions transformers should be able to represent (Weiss et al., 2021 Friedman et al., 2023 Smolensky et al., 2024). One motivation being the conjecture that transformers will learn to generalise iff they are able to represent a program which does so (Zhou et al.) 2023). However, empirical validation for this conjecture has thus far remained limited. Even when we expect the models to generalise they often only do so weakly, and rigorously bridging the gap between theory and practice remains a long standing challenge. A key motivation for TRA came from the tceoretical model presenled by SSmoennsy et al. (2024) which draws analogy between the transformer and the classical production system architecture from cognitive science. Their theoretical model has the attention mechanism first exactly match on specified conditions before applying a hard positional tie-breaker to perform one-hot selection. TRA arose out of an ongoing effort to create a learnable equivalent to their theoretical model.
Length Generalisation in Transformers: Several mechanisms have been posisted to improve transformer length generalisation. We partition this work into two categories. The first deals with trying to learn positional weights that are flexible w.r.t. the length distribution of the training data (Dubois et al. 2019:
Ruoss et al.] 2023; Li & McClelland 2023] Ray Chowdhury & Caragea 2023 Golovneva et al. 2024). Orthogonally, recent work has demonstrated that when extending to longer sequences the softmax tends towards uniformity and therefore introduces errors the attention signal (Ye et al. 2024: Veličković et al., 2024). TRA investigates whether these two streams of rescarch should combine. It attempts to reduce noise in the attention through selectively sparsity and learns general positional encodings with its forget gale.
Sparsifing Attention: Finally, substantial amount of effort has been made towards introducing sparsity into attention. One vein attempts to reduce the quadratic cost of the attention operation by applying fixed patterns (Child et al. 2019 Zaheer el al., 2021 Nawrot el al.) 2025), though doing so either requires restricting flexibility or is only applicable during inference time, Another has sought to allow sparsity within the softmax by allowing for zero probabilitics (Poters et al. 2019)), however, such algorithms often depend on computational complex operations and have delivered mixed results. Finally, Zhang et al. (2021) replace the softmax with ReLU, thereby introducing selective sparsity and demonstrating promising results on NMT.
10

---

Published in Transactions on Machine Learning Research (06/2025)
5 Conclusion
In this work we investigate whether fundamental issues with length generalisation occur due to conflict between positional and semantic information. As a probe for this hypothesis we present TRA, a simple modification to the attention mechanism designed to allow both types of information to operate in tandem. We find that the introduction of such a change significantly improves out of distribution generalisation. Morcover, the core innovation of TRA comes from marrying selectively sparsity with contextualised relative distance. We substantial opportunity for development of future architectures along this axis. Selective sparsity allows columns to focus what's important, while the forget gate enables memory over past time steps. While this latter mechanism does not enable full state tracking behaviour, temporal memory mechanisms have shown enhanced ability to approximate NC' solutions compared with standard transform- ers (Merrill et al., 2024). This means that combining a more sophisticated memory with a more flexible sparsification threshold should yield models which are able to approximate complex algorithmic reasoning more cleanly and with fewer layers. Pursuing research this direction has the promise of bridging the gap between our theoretical notions of what transformers should be able do and observed empirical realities, and in doing yicld models which capable of learning more robust, generalising attentive circuits.
11

---

Published in Transactions on Machine Learning Research (06/2025)
References
Damitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to aligu and translate. CoRR, abs/1409.0473, 2014. URL hem.appi.ssemmattcccccol.alaaa http:://ppp....aaaaaccccccoooerp
11212020
Federico Barbero, Alex Vitvitskyi, Christos Perivolaropoulos, Razvan Pascanu, and Petar Veličković. Round and round we go! what makes rotary positional encodings useful? In The Thirteenth International Conference on Learning Representations, 2025. URL
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Llya Sutskever, and Dario Amodei. Language models are few-shot learners. CoRR, 14165, 2020. URL https:://axxiy/ http://aawvvvr.aaaaaeeoe htrg/abbs20005. 14165.
Sébastien Bubeck, Varun Chandrasckaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, Harsha Nori, Hamid Palangi, Marco Tulio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Early experiments with gpt-4, 2023. URL https:
//arxiv. waaxivv.rg//aae atgg/bb//2033...../////// 12712
Rewon Child, Scott Gray, Alec Radford, and Llya Sutskever. Generating long sequences with sparse trans- formers, 2019. URL https://arxiv. orgo aat////s///...../// 10509
Jishnu Ray Chowdhury and Cornelia Caragea. Modeling hierarchical structures with continuous recursive neural networks. CoRR, abs/2106.06603. 2021. URL btttps://arxi..orr/bbb////// 06038
Christian Clark, Byung- Oh, and William Schuler. Linear recency bias during training improves trans- formers' fit to reading times. In Owen Rambow, Leo Wanner, Marianna Apidianaki, Hend AI-Khalifa, Barbara Di Eugenio, and Steven Schockaert (eds.), Proceedings of the 31st International Conference on Computational Länguistics, pp. 7735 7747, Abu Dhabi, UAE, January 2025. Association for Computa- tional Linguistics. URL https://acllatthollgg orgg/2/
org/2025. htt.://..rr.022255.....maaaaaaaa cowing-main.nnn 5177
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. CoRR, abs/2110.077332, 2021. URL https:
org/abs/2110.0rs 07733
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoe V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. CoRR, abs/1901.0225 2019.
URL httt:://axxi..or//aaaa htt//9/b///111111 028660
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. URL
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness, 2022. URL https://arxiv.. http://axxx..orrraa//////a/ 14135
Soham De, Samuel L. Smith, Amushan Fernando, Aleksandar Botev, George
Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, et al. The llama 3 herd of models, 2024. URL htttt:://aarttt http:://aarxiv..orabb/////a
21783
Yann Dubois, Gautier Dagan, Dieuwke Hupkes, and Elia Brumi. Location attention for extrapolation longer sequences. CoRR, 0ts//1111038877 2019. URL
12

---

Published in Transactions on Machine Learning Research (06/2025)
Brian DuSell and David Chiang. Stack attention: Improving the ability of transformers to model hierarchical patterns, 2024. URL https://arxiv.orr org/ httgg/bb/220 abs/aa10....oo/ 01749
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi Faith and fate: Limits of transformers on compositionality, 2023. URL
18654
MehammadReza Ebrahhii, Sunny Panchal, and Roland Memisevic. Your context is not array: Unveiling random access limitations in transformers, 2024. URL
05506
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Accessing higher- level reprosentations sequential transformers with foedback memory. CoRR, abs/ hbss2002.099402 2020. URL
httpss//aarziv..rr//a/////00 09402
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs, 2023. URL https:
waaxiv.org/bb//206.00rr 01128
Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Contextual position encoding: Learning to count what's important, 2024. URL
wttp://rrggabbb2.....g. 18719
Nuno M. Guerreiro, Elena Voita, and André F. T. Martins. Looking for needle in haystack: A com- prehensive study of hallucinations in neural machine translation, 2023. (ttpsttt::// /tttps://arxiv..oo www.aaiwwi.....aeeee....oeoooeee abs/
2208. 05309
Cheng-Ping Hsich, 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?, 2024. URL http:////axxiv.o/abbb//////
06654
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. The compositionality of neural networks: integrating symbolism and connectionism. CoRR, abs/1908.0883 2019. URL http::/arxiv..ot org/ abs/ 1908. 083511
Dieuwke Hupkes, Mario Giulianelli, Verna Dankers, Mikel Artetxe, Yanai Elazar, Tiago Pimentel, Christos Christodoollopullss, Karim Lasri, Naomi Saphra, Arabella Sinclair, Dennis Ulmer, Florian Schottmann, Khuyagbaatar Batsuren, Kaiser Sun, Koustuv Sinha, Leila Khalatbari, Maria Ryskina, Rita Frieske, Ryan Cotterell, and Zhijing Jin. A taxonomy and review of generalization research in nlp. Nature Machine Intelligence, 5(10):1 161 1174, October 2023. ISSN 2522-5839. doi: //o::///xx.oo..... wtt..//00//1111//2222221111 0204222266-02-00729--
URL
Ziwei Ji, Nayeon Lee, Rita Frieske, Tiezheng Yu, Dan Su, Yan Xu, Etsuko Ishii, Ye Jin Bang, Andrea Madotto, and Pascale Fung. Survey of hallucination in natural language gencration. ACM Comput. Suru., 55(12), March 2023. ISSN 0360-0300. doi: 10.1145/3571730. URL ftttps://doi.orr//11 10g//00.114/3571730
Amirhossein Kazemmejad, Inkit Padhi, Karthikeyan Natesan, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL htttps://opeereveee wttp:///oooooodddddrrreeeeee
Yuxuan Li and James McClelland. Representations and computations in transformers that support gener- alization on structured tasks. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL
Zhixuan Lin, Evgenii Nikishin, Xu He, and Aaron Courville. Forgetting transformer: Softmax attention with a forget gate. In The Thirteenth International Conference on Learning Representations, 2025. URL /ttps://openreeve
net rett/forummtd-qqqnayggrr
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing attention glitches with flip-llp language modeling, 2023a. URL https://arxiv.org//aa wwrg/bbs/200e 00946
13

---

Published in Transactions on Machine Learning Research (06/2025)
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Leurning Representations, 2023b.
URL https::/openrevvier http://rrrreiieeeeefoooooooddddd .wtt://oruummiii-e444yyyaaaeea
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. In Lluís Mårquez, Chris Callison-Burch, and Jian Su (eds.), Proceedings of the 2015 Conference on Empirical Melhods in Natural Language Processing, pp. 1412 1421, Lisbon, Portugal, September 2015. Association for Computational Linguistics. doi: 10. Saeelaanbbology.or//15--------
DD11-1166....11 1166. URL https:
Stephen Mority, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. CoRR, abs/
2016. URL http://aariiiv.orrrbb///0/// 07843
William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-8555 856, 2022. doi: 10.1162/ tacl 00493. URL attps://aclaanthoolls
www.orgg//20022 2ttr:/2222tacl-1..4 tacl-1.49/
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models, 2024. URL httpp::/aaxiv..ott org/
http://raiii...gg//////////// httgg//bb/2440.... 08819
Iman Mirzadeh, Keivan Alizadeh, Hooman Shahrokhi, Oncel Tuzel, Samy Bengio, and Mehrdad Farajtabar, Gsm-symbolic: Understanding the limitations of mathematical reasoning in large language models, 2024.
URL https: /tttpa//aaxxiv.rrt www.aaxwvvvv..aab..aaaaaaa wtwgg//aa.//....././g.// 05229
Piotr Nawrot, Robert Li, Renjie Huang, Sebastian Ruder, Kelly Marchisio, and Edoardo M. Ponti. The sparse frontier: Sparse attention trade-offs in transformer Ilms, 2025. URL attps://arxit org/abs/
2504.17760
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston,, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. URL hattps://arrae wwwpeiiivoorg//bbae trggabs/22055 11895
Mattia Opper and N Siddharth. Banyan: Improved representation lcarning with explicit structure. arXiv preprinl ar Xrv:244077 17771, 2024a.
Mattia Opper and N Siddharth. Self-strae at semeval-2024 task 1: Making self-structuring autoencoders learn more with less. arXie preprint orXiv:2404.01860, 2024b.
Mattia Opper, J Morrison, and N Siddharth. On the effect curriculum learning with developmental data for grammar acquisition. arXiv preprint arXiv:2311000190 2023a.
Mattia Opper, Victor Prokhorov, and N. Siddharth. Strae: Autoencoding for pro-trained emboddings using explicit structure, 2023b. URL https:://arxiw.org/bbbb/// 2ttbbb//33300000 05588
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023.
Bo Peng, Eric Alcaide, Quentin Authony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng Hc, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan S. Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, and Rui-Jie Zhu. Rwkv: Reinventing rmns for the transformer era, 2023.
Ben Peters, Vlad Niculae, and André F. T. Martins. Sparse sequencc--o--sequene models. CoRR, abs/1905.05502, 2019. URL
05702
14

---

Published in Transactions on Machine Learning Research (06/2025)
Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. CoRR, abs/2108.12409, 2021. URL https://arxxv..or//abbba/ 12409
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Πlya Sutskever. Language models are unsupervised multitask learners. 2019. URL
org//CorpusID:116022555 http:///waaaaaaooooooooooooooa/a
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 a unified text-to-text transformer.
CoRR, abs/1910.10666 2019. URL attp:://axaiv.orr//ass//910 wtt..orr/aabbs/11.10088...
Jishnu Ray Chowdhury and Cornclia Caragea. Monotonic location attention for length generalization. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabalo, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 28792 28808. PMLR, 23 29 Jul 2023. URL https:// proceedings mlr. wwwppp/eee//aappppppp/pp ppees//720./ray-chooour222 22mmm
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers, 2023.
URL https://arxiv... org/ a30/ 16843
Noam Shazeer. GLU variants improve transformer. CoRR, als als/2002.05502, 2020. URL https: /stps://arxiv.
Yikang Shen, Zhouhan Lin, Chin wei Huang, and Aaron Courville. Neural language modeling by jointly
learning synlax and lexicon. In International Conference on Learning Represenlations, 2018. URL https:
//openreview.net wtwpaarriiiennee//ouu///aaaae
Sania Sinha, Tanawan Premsri, and Parisa Kordjamshidi. A survey on compositional learning ai models:
Theoretical and experimental practices, 2024. URL
Paul Smolensky, Roland Fernandez, Zhenghao Herbert Zhou, Mattia Opper, and Jianfeng Gao. Mechanisms of symbol processing for in-context learning in transformer networks, 2024. URL https://arxiv.orr
abs/2410. 17498
Paul Soulos, Henry Conklin, Mattia Opper, Paul Smolensky, Jianfeng Gao, and Roland Fernandez. Com-
positional generalization across distributional shifts with sparse tree operations, 2024. URL https:
//arxiv.org/aart Dwtprii.orggabb/2441111 httgg/abs/242.11111 14076
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2104.09864, 2021. URL https::/arxxivoor///aa/a
Tulia Ture, Ming- Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: The
impact of student initialization on knowledge distillation. CoRR, abs/1908.00 1998.089622 2019. URL http:
7/axxiv. wra.ivi.org/aai hwwpg/abs/199.... 20001999.08862222
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, abs/1706, abs/1706.03362, 2017. URL http:
//axxiv.org hatg/abs/I1/ 1706.10660000 03762
Petar Veličković, Christos Perivolaropoulos, Federico Barbero, and Razvan Pascanu. softmax is not enough (for sharp aut-of-distrbbution). 2024. URL
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers, 2021. URL
aba/2106.065 06981
Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer, 2024. URL http:////axxvooo////////////
15

---

Published in Transactions on Machine Learning Research (06/2025)
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qilan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences, 2021. URL htttss//arxxiv.orr//////0000 wtt.:////abb//..1/00000
Biao Zhang and Rico Sennrich. Root mean square layer normalization. CoRR, abs 1911/1997777777 2019. URL
Biao Zhang, Ivan Titov, and Rico Sennrich. Sparse attention with linear units, 2021. URL https://arxix: org/ /tt/a/bb2214.001111
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization, 2023. URL
https://arxiv/ httpa/////ariivoooaaa/////// wrgg/ba////a/a/aoooaaaa// 16028
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. CoRR, abs/2107.02192, 2021. URL http:://arriv.oor//as/2000.00000
16

---

Published in Transactions on Machine Learning Research (06/2025)
A Implementation:
class TRA (nn Module):
def -_init__ (self, e_dim, n_heads, dropout-0.01). super (TRA, self) in_it______
self E =_diim
self e_dim // hehaads
self n_heads
self tokeys nn Linerr (e_dim, e_dim, bias=False)
self toqueries mn mm.Linear(e_dim e e_dim, Fias=
self tovalues nn Linear (e_ d__iin, e_dim, bias=Palse)
self delta_proj
Linear (e_dim, self.h)
11
self.c_out mn. Liner (inear e_dim)
# use qk norm but no scale to save params celf.qk_nors self beeeerrrbrbrebeeeeeee ("bias", torch. triu (torch oner diagonal-1) view 1,max_lem max_len) to bool))
13
(nsWormsself....
Meemntrise_afffe-Falle.
ones 1ma__eem max_ len),
self dropout nn. Dropout (dropout)
18
17
def
forward x):
18
b, 8, = x. shape # satch_size, seq len, e_dim transform q, k,
19
20
= self. toqueriess (x). vbe 8, self h, self e)) trw.s)).otseeeeeeeecc. (1, 2)
21
tokeys (x). view s, self h, self e) transpose 2)
= self. tovalues (x) vbe s, self self ee) transpose (1, 2)
23
qk norm (no scale -> can share) = self qk__nommq (q)
25
k self self.gk_norm(?. compute attn dot
27
s (q 0 krtranspose (-2, / torch sqrt (torch. tensor (self.e, dtype
float, device=x device))
28
= s masked fill (self. bias rerul(S))
float (' ('-1e11')
28
30
mask (S == 0)
31
compute positional weights
delta liogsiggoiddsell det (eltaaproo(x)) (x)) transpose (1, unsqueeze
contiguous () view (b, self s, 1)
D = (-mask) f) * delta D-flip(-1). cumsum (-1) (-1p
34
# dropout crucial to avoid loss spikes in language modelling
36
self dropout D)
[mask] -le11 mask: casual + threshold
38
A :=softmax(dim=-11))
39
protect against leaks where mask all from no-op attention
40
A = A. m.masked_ffit fill all unsqueeze (-1), 0)
41
compute output out A out out. transpose 2) contiguous () (be 8, self return self c_out (out)
42
43
Listing 1: TRA Implementation
17

---

Published in Transactions on Machine Learning Research (06/2025)
B Hyperparameters:
For synthetic tasks we train for one full pass through the data, for copy and induct this required circa 100k steps as our training set consisted of million examples, while for flip-flop language modeling this corresponded to 20k. We used batch size 128 for the former tasks and batch size 64 for FFL to accomodate larger sequence length (512).
For RoPE set the theta to 500k following Dubey et al 2024 For CoPE set npos max to 64 following the original authors (Golovneva et al. 2024 All models are trained using the AdamW optimiser with low dropout of 0.01 applied to the attention weights and feedforward hidden reprosentation.
C Attention Analysis: Language Modeling
Trained TRA models exhibit both completely content focused and completely positional attention heads, as well as ones where the trade-off somewhere in between. As general rule of thumb the more sparsely acti- vated the QK content weights, the weaker the degree of time decay enforced by the forget gate. Conversely, dense QK activations lead the model to focus on local window (see Figure 3 below). Attention patterns can be highly sparse indicating significant threshold pruning. Secondly, tokens exhibit considerable variance in time decay, impossible with fixed schemes such as ALiBi Press et al. (2021)
QK Weights
Positional Weight
Fnnal Atterdon Weights
109
102
200
200
200
303
300
200
100
200
250
300
350
50
102
200
302
50
230
220
230
302
Kay Tobena
Eay Noaaa
Kay Iaara
(a) Densely Activated Attention Head: Position Outweighs Content.
OK Weights
Posttional Weight
Final Attersion Weights
100
100
100
200
200
200
352
350
100
150
200
250
300
350
50
103
150
300
250
302
750
50
100
350
220
213
300
360
Troens
Cry Rosens
Kry Tolenn
(b) Mixed Activation Head: Content and Position Balanced.
Weights
Positional Weight
Attention Weights
100
100
100
200
200
200
252
200
303
359
350
350
50
100
150
200
250
300
350
50
102
150
220
250
302
200
50
100
250
220
230
300
350
Tokems
Day
Kay
(c) Sparsely Activated Attention Head: Content Dominates Position
Figure Attention Heatmaps from trained TRA language model. White signifies null attention.
18

---

Published in Transactions on Machine Learning Research (06/2025)
D Computational Efficiency and Potential Optimisations:
One potential disadvantage of TRA compared with schemes such RoPE, which remain the current defacto choice for positional embeddings, is that it introduces additional operations which can diminish runtime efficiency. To give an indication across two configurations:
Tiny: layers, heads, residual stream size 256, ffn 1024
Medium: layers, 8 heads, residual stream size 512, ffn 1024
On single Nvidia A40 card with window size 256 and batch size 64:
Tiny RoPE: 20k steps 58.73 minutes
Tiny TRA: 20k steps 65.88 minutes (circa seven minute increase)
Medium RoPE: 20k steps 156.34 minutes
Medium TRA: 20k steps 178.19 19 minules (circa 22 minute increase)
However, these differences without any effort towards runtime optimisation is significant scope for even greater efficiency. Though this is currently highly speculative, and hence left out of the main paper, there is no reason why the TRA algorithm could not be computed in a tiled manner and take advantage of the optimisations introduced by flash attention Dao et al. (022)): Dao (2023). While not trivial because the blocks we need to be processed in a causal manner to allow for online computation of the cumsum, the fact. that attention pabterns can be highly sparse and the forget gate can enable early stopping of row processing means that there should be considerable scope for greatly improving the efficiency of TRA.
end. We believe there
19


Reviews:
[{"review_id": "JynAZWMFBn", "reviewer": "Reviewer_YN4o", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- Introduces a novel positional encoding strategy that leverages contextually relevant positions informed by the QK dot product.\n- Demonstrates strong empirical performance across a range of benchmarks, outperforming several established positional encoding methods.\n- Provides clear and insightful analyses that help identify and interpret the relevant positional relationships captured by the model.\n\nWeaknesses\n- Section 2: The rationale behind the proposed framework is not sufficiently clear. Specifically, the design choices for computing contextual relative distances, their parameterization, and their integration into attention weights seem underexplained, given the multiple possible alternatives.\n- The formulation of the proposed method lacks contrast with existing approaches. The paper would benefit from a side-by-side comparison: presenting the formulation of a standard positional encoding method and directly juxtaposing it with the proposed approach to clarify what is novel and different.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper revisits positional encoding by incorporating richer contextual information derived from the query-key (QK) dot product. The approach begins by constructing a boolean mask that identifies relevant positions based on QK interactions. This mask is then used to compute a contextualized relative distance, which is parameterized and added to the attention weights prior to the softmax operation. The proposed method is evaluated against various positional encoding schemes—such as Absolute, Relative, and RoPE—across multiple benchmarks, consistently demonstrating improved performance. I particularly appreciate Figure 1, which effectively illustrates how the method integrates contextual position information into the final attention computation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would appreciate it if the authors could address the concerns outlined in the Weaknesses section—particularly by providing a more thorough justification of the proposed scheme and including formulation-level comparisons with existing positional encoding methods.\n\nI also have a few minor suggestions:\n- The transition from Equation (5) to (6) is unclear. Including more detailed explanations and intermediate steps would help readers better understand the underlying mechanism.\n- In Section 2.3, the variable $E$ is introduced without definition. Please clarify its meaning. Additionally, \"ith\" should be formatted as \"$i$th\" for consistency with mathematical notation.\n- The usage of \"1e11\" is somewhat opaque. While I understand it is intended as a large negative value to suppress attention weights (by pushing them towards zero after the softmax), a more explicit explanation would be helpful, especially for readers less familiar with this technique.\n\nOn a positive note, the experimental sections are clearly presented and easy to follow—I have no issues with their clarity or structure.", "broader_impact_concerns": "I have no major concerns regarding the broader impact of the proposed method. It presents a clear technical contribution as an enhancement to positional encoding, focusing on methodological improvements.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "uAy3kcLEUe", "reviewer": "Reviewer_DfmF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Weaknesses:\n- The paper doesn't analyze how sensitive the approach is to the threshold value or whether this introduces new hyperparameters that need tuning.\n- Completely masking out tokens based on a threshold could potentially lose valuable information in more complex settings.\n- The paper presents the results without a proper statistical analysis and some results appear to be from just 4 random initializations, which may not be enough for proper conclusions (for example: in table 1 they shows results on the Flip-Flop tasks and TRA achieves 100% accuracy while CoPE achieves 95.1% $\\pm$ 4.4%. The standard deviation suggests that in some runs, CoPE might have achieved close to 99.5%. Without formal statistical testing, it's harder to be confident about how reliable this difference is.\nAnd also, in the language modeling results, some perplexity differences between methods are quite small and without statistical significance testing, it's difficult to know if these small differences represent genuine advantages or just random variation)\n\n\nStrenghts:\n\n- I think the paper focuses on a core problem for the field of positional encodings and long length generalization and i really appreciate the approach the authors propose.\n\nPS. I pointed out more weaknesses then strenghts because i think you can work on those few points to make your paper even better, i do like the core approach and i think it has a value.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies one of the main limitations of transformer models: their poor generalization to sequence lengths not seen during training. The authors identify two key failures in the standard self-attention mechanism that contribute to this problem, the inability to properly remove irrelevant information from the attention distribution and the issue that positional biases can unintentionally up-weight irrelevant tokens when dealing with out-of-distribution sequence lengths\nTo solve these issues, the authors propose Threshold Relative Attention (TRA), which introduces two main modifications to the standard attention mechanism:\n- selective sparsity, to completely removing irrelevant keys from the attention softmax by applying a threshold to raw attention weights;\n- sontextualised relative distance, that means only calculating relative positional information between the query and the keys that pass the relevance threshold.\n\nThe authors implement this by first applying a ReLU to the QK^T dot product to eliminate irrelevant keys, then computing a contextual relative distance only among the remaining relevant tokens, and finally using a parameterized forget gate to create a temporal decay effect.\n\nThe approach is evaluated on synthetic tasks designed to test length generalization as well as on the WikiText-103 language modeling benchmark. Results show that TRA outperforms various positional encoding methods  in out-of-distribution settings, achieving almost perfect scores on synthetic tasks even when sequence lengths are up to 6x longer than those seen during training.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please work on the statistcal analysis of the experiments (as reported in weaknesses).\n- Plase also add an analysis on the threshold value sensitivity and if that introduces new hyperparameters.\n- I'd add connections to other sparse attention approaches and how TRA differs from those, and especially i'd underline how TRA is different from the paper \"Lin, Zhixuan, et al. \"Forgetting Transformer: Softmax Attention with a Forget Gate.\" arXiv preprint arXiv:2503.02130 (2025)\" that is very recent, but i can see a possible overlap, so i'd clarify properly on that to underline the value of TRA.\n- Finally, I'd add an analysis of training curves and convergence speed compared to baselines.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "080ZrMGeMS", "reviewer": "Reviewer_Zo6s", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The work is well positioned and highly relevant to recent literature \n-  The proposed method is unique and its strengths are clear.\n- There are a number of interesting synthetic experiments that showcase the strength of the method. \n- The attention maps in Figure 1 are an interesting addition and show the interpretability of the technique.\n\n\nWeaknesses\n\n- There is no discussion of the run-time of the method compared to RoPE. RoPE is particularly useful because it can be very easily optimised due to its 2x2 block-diagonal nature. It seems to me that these operations might slow down the forward and backward passes as they cannot be perhaps fused as easily. I believe it would be worth dedicating some space in the appendix to discuss this as a potential limitation (if it is indeed slower).\n- The way TRA is presented in Section 2 is slightly confusing as there are a lot of different symbols. It took me a while to parse it. For instance I don't believe there is a need to introduce the intermediate $\\hat{D}$ as one could simply write M directly. \n- For additional clarity, it would be useful to place in the appendix the pytorch/Jax code or pseudo-code for the implementation of TRA for completeness.\n- It would be interesting to see some results showing how for example there are configurations of TRA that can provably generalise over certain (synthetic) tasks while for instance RoPE cannot. At the moment the paper provides no theoretical insights which is fine, but this would help strengthen the work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper tackles the problem of length generalisation in attention mechanisms. In particular, they aim to address two known issues: 1) attention weights can never be 0 and noise therefore accumulates and 2) that positional encodings can \"fall out of distribution\" if they are calibrated on shorter context, thus harming length generalisation.\n\nTo tackle this problem, they propose a new attention mechanisms they call Threshold Relative Attention (TRA). TRA introduces a form of sparsity through the application of a ReLU function and a subsequent masking mechanism. It introduces a positional bias then by performing an interesting cumulative-sum type of operation which aims to behave like a learnable bias. \n\nThey evaluate this method on a number of synthetic benchmarks and on language-modeling.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Formatting issues with Eq.5 and Eq. 6. It seems like there are missing brackets over the matrices.\n\n- Add algorithm/pseudo-code of TRA (see weaknesses) in the Appendix.\n\n- Add runtimes of TRA, comparing to RoPE (see weaknesses). I am expecting TRA to run slower, but it might be worth bringing this up as a limitation if indeed the runtime is (significantly) slower. \n\n- As the work of Barbero et al. (2025) is repeatedly mentioned, it might be useful for completeness to add their method of p-RoPE to Table 1 as it aims to address length generalisation issues of RoPE to some extent. I see this as an optional change as there are already a number of comparisons to different PEs.\n\n- Could the authors report typical values of $\\delta_i$ in Eq. 8 in a trained model. As the function decays exponentially fast, this will quickly go to 0 as D grows. It would be interesting to see how the model prefers to use this mechanism.", "broader_impact_concerns": "None that I am aware of.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work presents a very insightful modification to prominent attentional mechanisms, with a special emphasis on an important downstream problem (improved length generalisation). All reviewers are in agreement that the work clearly meets (and exceeds!) the bar for publication in TMLR. Well done!
Decision: Accept (poster)