Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Authors: Jiajun Zhu
Peihao Wang
Ruisi Cai
Jason D. Lee
Pan Li
Zhangyang Wang
Venue: ICML 2025 poster (2025)
Keywords: Positional Encoding
Large Language Models
Equivariant Machine Learning
Primary area: deep_learning->large_language_models
TL;DR: We enhance the reasoning ability of large language models by representing positional encoding with contextual information under equivariance constraints.

Abstract: Transformers rely on both content-based and position-based addressing mechanisms to make predictions, but existing positional encoding techniques often diminish the effectiveness of position-based addressing. Many current methods enforce rigid patterns in attention maps, limiting the ability to model long-range dependencies and adapt to diverse tasks. Additionally, most positional encodings are learned as general biases, lacking the specialization required for different instances within a dataset. To address this, we propose con**T**extualized equivari**A**nt **P**osition **E**ncoding (**TAPE**), a novel framework that enhances positional embeddings by incorporating sequence content across layers. TAPE introduces dynamic, context-aware positional encodings, overcoming the constraints of traditional fixed patterns.  We show that TAPE can provably facilitate LLM reasoning ability by emulating a broader class of algorithms. By enforcing permutation and orthogonal equivariance, TAPE ensures the stability of positional encodings during updates, improving long-context ability. Our method can be easily integrated into pre-trained transformers, offering parameter-efficient fine-tuning with minimal overhead. Extensive experiments show that TAPE achieves superior performance in language modeling, arithmetic reasoning, and long-context retrieval tasks compared to existing positional embedding techniques.

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Jiajun Zhu 12 Peihao Wang"
Ruisi Cai
Jason D. Lee²
Pan Li
Zhangyang Wang'
Abstract
1. Introduction
Attention mechanisms are core component of many mod- em deep learning architectures, enabling models to selec- tively focus on relevant information within a given context. Transformer models (Vaswani et al., 2017) and their nu- merous variants (Carion al., 2020; Dosovitskiy al., 2021; Zhao et al., 2021), which fundamentally driven by attention, have revolutionized tasks involving sequential and spatial data, such as text (Kitaev et al., 2020), image (Dosovitskiy et al., 2021), and point cloud (Zhao et al., 2021). More recently. large transformer models have be- come dominant in natural language understanding, language generation, and complex reasoning (Brown al., 2020).
Transformers rely on both content-based and position-based addressing mechanisms to make predictions, but existing positional encoding techniques often diminish the effectiveness of position-based addressing. Many current methods enforce rigid patterns in attention maps, limit- ing the ability to model long-range dependencies and adapt to diverse tasks. Additionally, most positional encodings are learned as general bi- ases, lacking the specialization required for dif- ferent instances within dataset. To address this, propose conTextualized equivariAnt Position Encoding (TAPE), novel framework that en- hances positional embeddings by incorporating sequence content across layers. TAPE intro- duces dynamic, context-aware positional encod- ings, overeoming the constraints of traditional fixed patterns. We show that TAPE provably facilitate LLM reasoning ability by emulating broader class of algorithms. By enforcing permu- tation and orthogonal equivariance, TAPE ensures the stability of positional encodings during up- dates, improving long-context ability. Our method can be easily integrated into pre-trained transform- ers, offering parameter-efficient fine-tuning with minimal overhead. Extensive experiments show that TAPE achieves superior performance in lan- guage modeling, arithmetic reasoning, and long- context retrieval tasks compared to existing posi- tional embedding techniques.
Delving into attention's 's underlying paradigm, the prediction made for each token is expressed as weighted aggregation over the representations of other tokens. Due to the softmax function, attention often generates sparse mask, extracting a limited subset of tokens for interaction. Through this interpretation, attention can be understood as an addressing mechanism (Hopfield, 1982; Pagiamtzis & Sheikholeslami, 2006) that searches the context, locating and retrieving token representations deemed most relevant important.
Since the attention score is computed upon token features and positions (see Sec. 2), transformers' addressing ability can be further decomposed into two fundamental mecha- nisms: content-based addressing and position-based ad- dressing. Content- addressing recognizes relevant tokens through feature similarity, while position-based ad- dressing is facilitated by positional encoding techniques, de- signed to (ideally) enable random access along the sequence via indexing. It is important to let the two mechanisms cooperate to tackle more complex tasks, such as in-context retrieval (Hinton Anderson, 2014; et al., 2016), arith- metic (Lee al., 2023; McLeish al., 2024b), counting (Golovneva al., 2024), logical computation (Liu al., 2024), and reasoning (Wei et al., 2022; Rajani al., 2019; Dziri et al., 2024). However, we contend that the role of position- based addressing is limited, if not diminishing, in modern transformer architectures (Ebrahimi et al., 2024).
'Equal contribution. Work was partially done when J. Zhu was undergraduate ZJU. 'University of Texas at Austin ZZhejiang University Princeton University 'Georgia Tech. Correspondence to: Jiajun Zhu <jiajunzhu @ edu>..
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
It has not escaped our notice that most existing positional encodings weakens the position-based addressing capabil- ity. Recent works (Press et al., 202 Su et al., 2024; Chi et al., 2022b; Sun et al., 2022) impose fixed and some- what artisanal pattern on attention maps, typically adopting a decaying pattern in relation to relative distances, thereby enforcing a locality bias. This rigidity limits the ability of attention to model long-range dependencies and operate on flexible group of tokens. Although some positional encod- ings have trainable parameters (Vaswani et al., 2017; Shaw et al., 2018 8; Chi et al., 2022a; Li et al., 2023), the hypothesis space is often excessively constrained.
lowing attention values be computed based on not only token similarities but also positional proximities. We addi- tionally customize an MLP layer that directly mixes token features with positional encodings, while preserving orthog- onal equivariance.
We demonstrate the superior performance of TAPE on arith- metic reasoning tasks (McLeish et al., 2024a), which require LLMs to effectively locate/address and retrieve specific to- kens, as well including SCROLLS (Shaham et al., 2022) and passkey retrieval (Mohtashami & Jaggi, 2023), to validate the gener- alizability of the framework.
representative natural language tasks,
Perhaps more crucially, most existing positional encodings are handcrafted and learned as general bias across the entire dataset, lacking specialization and adaptability to spe- eific instances informed by the context. In more advanced tasks, such as general-puree purpose algorithmic reasoning, LLMSs rely on attention to internalize the representation of com- putation graph, which must vary according to the prob- lem specified in the context (Elhage et al., 2021; Ameisen et al., 2025). However, transformers without context-aware position-based addressing have been theoretically shown to represent only restricted family of algorithms (Merrill & Sabharwal, 2023).
Our contributions are summarized as follows:
We introduce TAPE, novel framework learning to represent token positions in the feature space jointly with sequential learning. TAPE contextualizes posi- tional embeddings with sequence content across layers to enhance the position-addresssing ability of transform- ers. We further enforce TAPE with permutation and orthogonal equivariance to guarantee the stability of positional encodings during the update. We theoreti- cally show that TAPE achieves higher expressivness for algorithmic reasoning.
To unleash the power of position- addressing, we en- deavor to design a more universal and generic position en- coding for language transformers. We introduce Contextu- alized Equivariant Positional Encoding (TAPE), a novel framework designed to contextualize positional embeddings by incorporating sequence content. TAPE continually pro- gresses information flow between positional embeddings and token features via specialized attention and MLP layers. To ensure the stability during model updates, we enforce permutation and orthogonal group equivariance properties on attention and MLP layers. This approach is inspired from the studies for geometric deep learning which repre- sents structural patterns (e.g., graphs) by integrating token features with their geometric properties while preserving inherent physical symmetries (Wang et al., 2024b; Huang et al., 2024). Building on this motivation, we encode algo- rithmic structures into transformers and theoretically show that TAPE achieves greater representational power, captur- ing algorithms of higher complexity than those representable by standard transformers. Further on, by enforeing equiv- ariance priniples, TAPE maintains the relative relationships between encoded positions, ensuring transformer's long- context ability and length generalization.
We propose practical implementations for our TAPE, which extends conventional positional embeddings into multi-dimennsonal and facilitates attention and MLP in transformers with two levels of equivariance. We also show that TAPE has hardware-efficient implementation and can be used as a drop-in component into extant pre-trained models for parameter-efficien fine-tuning.
Extensive experiments showcase TAPE's superiority in both training from scratch and parameter- efficcin fine-tuning scenarios for language modeling as well as downstream tasks such arithmetic reasoning and long-context retrieval. TAPE achieves state-of-the-arr performance in language modeling, surpassing base- lines in perplexity reduction for long sequences. We also report the state-of-the-art performance of TAPE in long-context tasks such as passkey retrieval tasks with LLM fine-tuning, and in arithmetic learning.
2. Preliminaries
In this work, we aim to design expressive and generalizable positional embeddings for transformers to address complex language tasks. Let X [ ... ZNd N represent the input sequence of tokens, where N is the context length and C' is the feature dimension, and M E {0, 1 de- notes the cross-token dependencies, often used as attention mask in transformers. Transformers lear token representa- tions using the attention mechanism (Vaswani et al., 2017),
Technically, we extend conventional vectorized positional embeddings into a dimenioammesio tensor, which enriches interactions between positional embeddings and token fea- tures. In the attention mechanism, TAPE incorporates the pairwise inner product between positional encodings, al-
2

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
which propagates information across tokens by computing pairwise correlations. Since pure attention is inherently language models integrate posi- tional information into the attention computation to differ- entiate tokens based on their positions.
2.2. Attention Bias Positional Encoding
An alternative method for encoding positional information involves applying bias to the attention map, conditioned the relative distances between tokens during the attention computation. The pre-softmax attention value with bias can be formulated as:
permutation-equuiaarii
2.1. High- High-Dimmesional Features as Positional Encoding
One common approach is to leverage high-dimensions fea- tures to represent positions. Positional encoding can be formulated as a series of embeddings attached to each token index Ei 2g...e EN, with the shape of e; determined by the specified positional encoding schemes. When computing the attention value, the pre-softmax attention value can be in general formulated as below". For every j E [N] such that M; زان 1:
Oi (W (o:))) (W k.,) - b(i.)),
where b(i,)) N X R is bias regarding the token indices i and j. Many existing positional encoding methods can be interpreted as various instantializations of b(i, j). We follow Li et al. (2023) to summarize a few examples. (i) In T5 (Raffel et al., 2020), b(i,j) = (imin(i-jj ). where Imaz denotes the maximal relative distance considered, and (rr E RR : E [0. [maz]] are learnable scalars. (ūi Alibi (Press et al., 2021b) simplifies the bias term to b(,,j) -rli - jl. where > 0 is a hyperparameter that acts as the slope, imposing linear decay pattern based on the relative distance. (iii) Kerple (Chi et al., 2022a) enforces a logarithmic or power decay rate: b(i,j) = -ri log(1 + respectively, where
as q(,,, ei) Ti) (x),,))
where -) and .): are generalized query and key trans- formations that incorporate positional features. The original transformer paper (Vaswani et al., 2017) assigns each abso- lute token index a vector of length identical to token embed- dings, either learnable or fixed as sinusoidal waves: e, E RC. The query and key transformations directly add the positional information into token features at the first layer: (()., = W @ + e;) and k( e₁) = -W Wx( for some query and key matrices We WK E ROX' Shaw et al. (2018) introduces learnable embeddings for relative distances, which are applied to the key vector dur- ing attention computation. More recently, Rotary Position Encoding (RoPE) (Su et al., 2024) has gained widespread adoption in modern LLMs (Touvron al., 2023a;b; Bider- man et al., 2023; Chowdhery et al., 2023; Jiang et al., 2023). RoPE encodes absolute positions using series of block- wise rotation matrices E E XC/2xx22 while implicitly capturing relative distances during dot-product attention. Formally, the positional embeddings and the transformation -(:-)) are defined as shown below, with k(.--) adhering similar formulation:
r2li - -il) and b(i.j) = -rili- - r1, r2 > 0 are hyperparameters. (iv) FIRE (Li al., 2023) learns neural network with parameters O to model the bias: b(i,j) fo( (oo(ii))//(maseee L)),, where (xx)
j -
log(cz + 1), and L > 0 is a hyperparameter.
3. Our Approach
3.1. Motivations and Design Principles
In the paper, we interpret the attention mechanism as an addressing system, where row-wise attention scores can be viewed as an indicator vector locating important tokens in the context to inform predictions for the current token. The underlying addressing mechanisms include both content- based addressing, which locates tokens via feature similarity, and position-based addressing, which leverages positional encodings to extract location-based information. Content- based addressing is often prioritized in language modeling which is evidenced by series of simplifications on posi- tional encoding in the literature (Press et al., 202 bb; Haviv et al., 2022; Wang et al., 2024a; Kazemnejad et al., 2024) due to the fact that natural language semantics primarily depend on the meaning of constituent words rather than their arrangement order (Sinha et al., 2021). However, position- based addressing can sometimes be crucial for many ad- vanced tasks. Ebrahimi et al. (2024) demonstrates that in arithmetic tasks (Lee et al., 2023), token's position is as im- portant as its value. An ideal attention map for performing addition needs to exclusively rely on token indices.
q(z;, e;) = R. W qai,
R = diag(ei,1, ... (i,0/2),
sin(lm)))
e,.
cos(Omi))
(2)
where diag(-) constructs block-diagonal matrix by concatenating the arguments on the diagonal. In RoPE, the hyper-parameters 6m ranges from 6m -100002m/C m E [C/2]. Subsequent works explore meth- ods to extend the context length for RoPE-based L.LMs through the adoption of damped trigonometric series (Sun et al., 2022), positional interpolation (Chen et al., 2023a) and adjustments coefficients (0m) ) me(C/2] (M/Loo 2023; Peng et al., 2023b; Liu et al., 2023).
Moreover, we observe that the interaction between token features and positional embeddings is lacking in current transformers. Golovneva al. (2024) demonstrate that in-
Forr simplicity, we ignore the denominator / F by default.
3

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Oulput
Oulput Probabillies
Feed Forward for PE
Probabilties
Lineer
Language Head
Lenguage Head
Broadcast
Prod.
Taken Embeddlig
Takkn Embedding
Position Embedding
Feed Forward
Lineer
Norm
Add Norm
Feed Forward
Feed Forward
Feed Forward for PE
Causal Attention with PE
X
N
N
Norm
Mat. Mul.
Causal Altention
Causal Altention with PE
Scale Mask SoftMax
Dot Product
Token Embedding
Postton Embedding
Token Embedding
Position Embedding
Mst.
Input
Input Token
K VIE
(a) Traditional position embedding.
(b) TAPE with enhanced causal attention and feed forward layers.
Figure 1. Overview of proposed TAPE in standard decoder-only architecture. Different from traditional positional encoding, TAPE represents and updates positional features layer-wisely through interactions and joint training with token representations.
corporating the interplay between context and positional information allows for more flexible addressing, leading to improvements in complex compositional tasks such as algorithm execution and logical reasoning (Liu et al., 2024). In domains with data characterized by intricate geometries, such as graphs and point clouds, capturing the interaction between node point features and their geometric loca- tions is essential for effectively learning structural patterns relevant to tasks (Wang et al., 2024b; Huang et al., 2024). Algorithmic reasoning requires attention to be endowed with similar properties to internalize representations of compu- tation graphs in order to emulate corresponding algorithms for solving specified tasks (Elhage al., 2021; Ameisen al., 2025). However, Merrill & Sabharwal (2023) shows that purely content-based addressing attention is subject to limited algorithm class.
ized transformer block consisting of two separate operations: token mixing and position contextualization. The token mix- ing is formulated as a function f X x & x M + X, which combines token features and positional embeddings to represent each token. The position contextualization g X x M + & encodes the context information into the positional embeddings. layer does not have component corresponding to g that layer-wisely represents position encodings.
standard transformer
Tensorial Positional Encoding. Our first enhancement extends positional encodings to a multi-dimensionll for- mat and diversifies their coupling with token features to allow for richer interactions among token and position rep- resentations, drawing inspiration from positional encodings used in geometric learning (Deng et al., 2021; 1; Wang et al., 2024b; Huang et al., 2024). We first divide the hidden di- mension into M blocks, resulting = I1, ...... [EN] E x.NxMxB with æ; E M and B = C/M. We propose assign each block L-many R-dimension posi- tional embeddings. Therefore, embeddings as E [er, ... ,eN]T NxMxXR with e; E M and D = Mx
Based on above observations, we aim to establish a more expressive family of positional encoding, which can be ef- fectively informed by the context to facilitate rosition-based addressing LLMs. The main idea to customize atten- tion and MLP modules in transformers such that they can iteratively update positional embeddings each layer with sequence content, and use the updated embeddings as the positional encoding for the next layer. We formally outline a couple of key design principles below.
reorganize positional
Equivariance Principles. Second, we establish two fun- damental criteria for the design of functions f and g. Con- ceptually, by representing each token tuple comprising its token and positional embedding, the entire sequence can be viewed as unordered set. This implies that permuting these tuples arbitrarily will not alter the outputs of f and g, aside from a corresponding change in order (Zaheer et al., 2017; Lee et al., 2019). We note that this is an intrinsic property of standard attention. Furthermore, we aim for the positional embeddings to effectively model relative dis-
General Formulation. Let tuple (x, E, M) represent sequence, where X E X' C N are the token fea- tures, E E & C RND are the positional embeddings, M E M = (NNNN encodes the intra-data dependen- cies, often used as attention mask. For causal language modeling, M becomes constantly a binary lower-trianuulat matrix (M,, * if and only if j S i). We define general-

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
tances, necessitating that / remains invariant translations in the token positions (Sun al., 2022). This invariance can be achieved by structuring f and g to depend the positional embeddings in a manner invariant and equivari- ant, respectively, to orthogonal transformations along the last dimension (Villar et al., 2021). Formally, let us denote (N) as a permutation group over N elements, and ( as an orthogonal group over the R-dimension Euclidean space. The two aforementioned criteria require f and g to satisfy that: for vP E II(N), R O(R),
Among these, we select the widely-used sinusoidal series embedding, RoPH (Su et al., 2024), as our initialization, as detailed in Sec. 3.3.
O ( II)ariia Token Mixing, In each transformer block, f updates token features through attention and MLP following the principles of germutation-equivariaae and ( -ivvariirie. We define pre-softmax attention value between the i-th and j-th tokens (M,, F 0) as:
( PMP = E,
(4)
04,
Σ a..j.m
(6)
gPP,, PER, PMP - E,
(5)
Oi.j.m = (W Qu: ( (m. x
where left- ceft-multiplication of P II(N) permutes the first dimension of X and E, while right-multipliicaion of Re E (R) applies to the last dimension of tensor E. When permutation is applied for the attention mask, it permutes both columns and rows accordingly. We note that attention with ROPE inherently satisfies Eq. 4, with the invariant orthogonal group being O(2). We formalize the merits of position contextualization and group invariance in Sec. 4. In summary, we prove that contextualized positional encoding improves flexibility and adaptability to represent a broader class of algorithms than either raw attention or RoPE. Additionally, orthogonal group equivariance ensures that the positional encodings remain relative, facilitating generalization across varying context lengths.
where s(.) R LIX. + R BBx can be any function. Permutation-equivarianee is inherently preserved in pair- wise attention, regardless of the method used to derive atten- tion values. O (())-invariance is achieved by computing the inner product of positional embeddings (Villar et al., 2021; Wang et al., 2022a; 2024b). We note that O -invarrireence stems from the separation of the inner product calculations between features and positional embeddings, in contrast to Vaswani et al. (2017). In practice, we can let B- = I. and & be an identity mapping, which simplifies Eq. 6 to a series of tensor multiplications. After applying attention, a stan- dard MLP layer employed to transform token embeddings without using positional embeddings.
3.2. Contextualized Positional Encoding with Equivariance
((R))-qqivarian Position Contextualization. The pri- mary contribution of this work is the introduction of a method to condition positional embeddings on sequence content. We employ an (R))eequiv function g to en- sure structure conservation of this update. A key insight is that linearly combining positional coordinates preserves O -qquivvaaae,,,,,, provided the weights are invariant to the orthogonal group (Villar et al., 2021; Wang et al., 2022a; Huang et al., 2024). This observation leads us to leverage attention maps, which capture content- content-based token rela- tionships, to integrate positional embeddings. Hence, the attention layer can update positional embedding via:
In this section, we instantiate design principles discussed in Sec. 3.1 as practical neural architecture. We note that although there are lots of ways to achieve conditions in Eq. 4 and 5 (Dym & Maron, 2020; Bogatskiy et al., 2020; Yarotsky, 2022), the proposed method focuses on enhancing existing components used in standard transformers with consideration of computational efficiency. We term our proposed approach of informing positional encoding with context through enforcing equivariance as Con Textualiee EquivAriant Positional Encoding (TAPE).
N
Mij exp(ass..)
Σ
Model Structure and Initialization. We adhere to the conventional architecture of the standard trans- former (Vaswani et al., 2017), wherein each layer comprises an attention module for token mixing and a Multi-Layer Perceptron (MLP) for channel mixing. Both the attention and MLP components are tailored update positional em- beddings at each layer. We depict the overall architecture in Fig. 1. The initial positional features may encompass vari- cty of representations, including but not limited to learnable features ( (Vaswani et al., 2017), sinusoidal series (Vaswani et al., 2017: Su et al., 2024; Sun et al., 2022), or random Fourier features (Rahimi & Recht, 2007; Yu et al., 2016).
Ci.m
N Ej= M, exp(as.r..) j=1
Gjmm
where fēr.m (Bs.mjE[N]].me [M[[] denotes an intermediate output of the attention layer. In practice, we share the attention map between Eq. and We can re-use aij.m computed in Eq. 6 because we have shown that attention weights ai.j.m (II)-invarin
We further propose layer similar to the function of MLP, which directly transform matrix-fome positional embed- dings with token features incorporated. Specifically, we first flatten the first two dimensions of ē; E MMxxxX to the
5

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
RoPE
RandPE
NoPE
FIRE
TAPE
100
80
60
40
60
20
80
20
40
60
80
20
40
60
80
20 40
80
20
40
60
80
20
40
60
80
0
Figure 2. Accuracy on addition task between different methods on 2x context length. The x- and y-axes represent the sequence lengths of the two operands respectively. Models trained sequence with length up 40 while tested sequence with length up 80. The average accuracy across the heatmap is 26.32%, 26.56%, 22. 22.45%, 26.98% and 32.82% respectively for RoPE (Su et al., 2024). RandPE (Ruoss al., 2023), NoPE (Kazemnejad al., 2024), FIRE (Li et al., 2023), and our TAPE.
shape L then apply linear transformation constructed by token features, and finally unflatten the first dimension of the resultant matrix to és E MXLX Vj E [N]:
that the augmented model remains identical the original at initialization, follow Hu et al. (2022) by setting the initialization of W in Eq. 8 to zero such that all updates to the positional encoding inside the block will then be reset via a residual connection. To allow for parameter-efficie fine- sneettuning we enable gradients to update the weights for position encoding contextualization including W1, 1, W2 and the weights for the post-attention linear layer, while keeping all other parameters frozen.
é, unflatten (22 W flatton((;))
(8)
where we denote I; E RC as the output of attention af- ter token mixing. Let be the intermediate dimension, W1 and W₂ have shape ML I. We further define mapping between token features to linear transformations. To reduce computation overhead, we adopt an MLP to transform ã; into an T-dimensional vec- tor and form a diagonal matrix with the resultant vector as its diagonal. The detailed computational flow is illustrated in Fig. in Appendix C. By applying linear transformations only to the first two dimensions of e;, this layer maintains O O(R) equivarriine
RC RIxI
4. Theoretical Analysis
In this section, we formally state the theoretical benefits of our proposed architecture, analyzing two key aspects: (I) its enhanced expressiveness for algorithmic reasoning, enabled by the contextualization mechanism, and (2) its guaranteed stability under equivariance constraints.
We summarize the overall geometric properties of our archi- tecture in Proposition 3.1 below.
Benefits of Contextualization. To characterize the ex- pressive power of transformer architectures on reasoning problems, we adopt cireuit-complexin lens following Mer- rill (2019); Merrill & Sabharwal (2023); Liu et al. (2022); Yang et al. (2025), identifying families of algorithms these models can and cannot simulate. Our premise that formal reasoning entails multi-step algorithmic computation. Ef- fective reasoning with LLMs therefore requires sufficient capacity represent these algorithms within the model.
Proposition 3.1. The proposed TAPE including: (i) atten- tion in Eq. with normal MLP for token mixing, and (ii) attention Eq. 7 with MLP defined Eq. 8 for position contextualization, satisfies Eq. 4 and Eq. 5.
3.3. Parameter-Efficien Fine-tuning with TAPE
In this section, we demonstrate that our TAPE can be seamlessly integrated into pre-trained models, en- abling Eefameter-fffiieen Fine-Tuning (PEFT) enhance position-based addressing in existing architectures. Notably, the widely adopted RoPE (Su et al., 2024) can be considered a special case of TAPE. This can be seen by letting L sin((mi)] m, (S)n( cos(Om [)] T. With this configuration, Eq. becomes equivalent to Eq. 2. As result, RoPE can serve as the initialization for TAPE, while the model is further enhanced by incorporating the contextualization component specified in Eq. 7 and 8. This initialization is applied across all experiments, encompassing both pre- vee-traning and fine- tnnin stages. Specifically, during fine-tuning, to ensure
Circuit complexity quantifies the difficulty of different prob- lems, with NC, AC and TC being the key classes. The classes NC, AC, and TC consist of problems (with input size N) solvable by circuits depth O(log NY Specifi- cally, NC allows only bounded fan-in AND, OR, and NOT gates; AC* extends NC* by permitting unbounded fan-in gates; and TC* further includes unbounded fan-in threshold gates. We refer interested readers to formal definitions in Li et al. (2025); Vollmer (1999); Arora & Barak (2009). It is known that AC Ç TC c NC but it remains open whether TC = NC! achievable (Cook, 1985).
R and 1,,m,1
(0)))
Below we show the algorithmic expressivity of TAPE
6

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
reaches NC -Cmmcompeete, which implies all the problems as hard NC² -complee problem within the NC' class can be solved by transformers with TAPE.
determined by the relative distances between token positions (Sinha et al., 2022). This property also ensures stability and supports generalization beyond the training context length.
Theorem 4.1 (Cireuit Representational Power). There exists a constant-depth and constant-widff transformer equipped with TAPE and satisfying Eqs. 4 and 5 such that it is capable of solving all NC¹ -complete problem under AC -reductions.
5. Experiments
In this section, we first validate our method on arithmetic tasks, which relies on better position-addressing ability for prediction (Sec. 5.1). We also show our effectiveness in natural languages, in both pre-training (Sec. 5.2) and fine- tuning case (Sec. 3.3). More experiments, visualization, and model interpretation can be found in Appendix D and E.
It is known that the representation capacity of constant-layer transformers is upper bounded by TC Theorem 4. shows a clear separation in expressivity between transformers with and without TAPE unless TC NC¹, Equipping attention with TAPE can potentially enhance its computational power.
5.1. Arithmetic Learning
The proof of Theorem 4.1 shows TAPE can solve NC complete state-tracking problem. When understanding LLM reasoning from the automaton perspective (Merrill, 2019), state-tracking ability is essential for modeling sequential computations and carrying out formal logic entailment (Mer- rill et al., 2024; Peng et al., 2025; Grazzi et al., 2024). Our proof also highlights the importance and the role of layer-wise update of positional encoding, contextualized positionallinformeed token mixing. The multi-layer construc- tion illustrates how layer-wise update of positional encoding progresses positional and tokens features into the positional encoding, which learns to be a more complex representation encoding inherent algorithmic reasoning structures.
As demonstrated by prior research (Lee et al., 2023; Zhou et al., 2024), even large transformer models struggle with arithmetic tasks. Recent studies suggest that this limita- tion may stem from their constrained position-addresssng capabilities (Ebrahimi et al., 2024). In particular, arithmetic tasks treat every digit as equally important to the equa- tion, regardless of its distance from the output. In contrast, traditional positional embeddings for language tasks often assume a distance-decay effect, where words farther apart have less significance in the output. Positional contextualiza- tion potentially addresses this by dynamically reweighting positional importance based on the task context. To evaluate the ability of LLMs of performing arithmetic tasks with our position embedding, dataset (McLeish et al., 2024a) which contains 20 million samples with i x i (i < 40) operand lengths. We train transformers from scratch using the arithmetic data, and during evaluation, we sample 100 samples for each pair of operand lengths. Following the existing attempt (MeLeish et al., 2024a), the operands in the training set are not nec- essary to have the same length, but the maximum length of two operands are the same. We then report model accuracy for each (i. J) length pair. Note that accuracy is measured strictly, counting only exact matches of all output digits as correct. The transformers are standard decoder-only archi- tecture with config detailed in Appendix B. We compare our method with four baselines, including RoPE (Kitaev et al., 2020), RandPE (Ruoss et al., 2023) NoPE (Kazemnejad et al., 2024), and FIRE (Li et al., 2023).
This result is largely inspired by the concurrent work of Yang et al. (2025). It proves seminal result showing that positional encodings constructed by accumulating input- conditioned Householder matrices (Yang et al., 2024) can uniformly represent all algorithms within NC¹ class. Our proof connects the accumulation of Householder matrices with our proposed geometric constraints.
use the Addition Bucket 40
Importance of Equivariance. It has been shown that positional embeddings that encode relative distances among tokens are crucial for ensuring long-context capabilities in LLMS (Li et al., 2023). Although initializing positional encodings with RoPE or random Fourier features at the first layer guarantees relativity under inner product (see Eq. 6), it remains unclear whether this property preserved after the embeddings are updated in intermediate layers. To address this concern, we provide the following justification.
The heatmaps further demonstrate TAPE's superior gener- alization to longer sequences, as indicated by the concen- trated dark-colored regions representing higher accuracy across a wider range of operand lengths. TAPE outper- forms other methods with the highest average accuracy of 32.82%. Compared to FIRE, which achieves 26.98% and previously held the strongest length generalization in arithmetic tasks (Mcl et al., 2024a; Zhou et al., 2024), TAPE shows a remarkable 21.6% relative improvement. This shows TAPE's effectiveness in maintaining accuracy
Proposition 4.2 (Relativity of position encoding). Suppose a transformer consists of f,g satisfying Eqs. 4 and 5. Given inputs (X, E) with position encoding E initialized by RoPE or random Fourier features, then the transformer is invari- ant to shift on token indices.
Proposition 4.2 highlights the importance of the equivari- ance principle with respect to the orthogonal group. Under O equiiaaraar,, even after positional encodings are up- dated with token features, their effect on attention is still
7

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Table 1. Performance comparison on seven datasets from SCROLL benchmark. For all tasks, the performance is better if the reported metric is higher (1), We highlight the lop-performing methods via bold font.
QAS
CNLI
NQA
QuAL
QMS
SumS
GovR
Metric (%)
F1 (t)
EM (1)
F1 (1)
EM (1)
Rgm (t)
Rgm (1)
Rgm (t)
Median length
5472
2148
57829
7171
14197
9046
8841
RoPE
8.39
65.00
1.77
0.04
6.34
5.63
9.71
ALiBi
8.25
69.62
4.11
0.0
9.92
9.78
18.81
RandPE
13.44
62.01
4.63
0.38
8.43
8.31
8.93
FIRE
3.41
71.26
0.48
1.25
8.78
7.42
11.03
xPos
9.02
71.75
4.83
0.24
10.73
9.38
16.38
TAPE (Ours)
11.52
72.80
6.79
11.60
12.42
10.34
15.18
as sequence lengths increase, making it particularly suitable for long-range dependency tasks.
Table 2. Evaluation perplexity across 1k to 8k context lengths. Lower perplexity means better performance (4). Top results are marked bold. Each model is first pre-trained from scratch and later fine-tuned on the downstream long-context datasets.
5.2. Pre- Fre-Training from Scratch
Pre-training a language model on a corpus followed by fine- tuning on downstream tasks is the standard methodology for evaluating the performance of positional embeddings in prior studies (Li et al., 2023; He et al., 2024). Simi- larly, we first pre-train transformers with 1024 context win- dow from scratch, using C4 dataset (Raffel et al., 2020), and then fine-tune those models in long-context benchmark SCROLLS (Shaham et al., 2022). We report three eval- uation metrics for seven different tasks: unigram overlap (F1) for Qasper and NarrativeQA, and exact match (EM) for QuALITY (QAS) and ContractNLI (CNLI), and Rgm score (the geometric mean of ROUGE-1.2L)) for the three summarization tasks: GovReport (GovR), QMSum (QMS), and SummScreenFI) (SumS). We compare our methods with RoPE (Kitaev et al., 2020), ALiBi (Press et al., 2021a), RandPE (Ruoss et al., 2023), FIRE (Li et al., 2023) and xPos (Sun al., 2022), and report the results in Tab. 1.
Method
1024
2048
4096
8192
Proof-pile
LoRA
3.828
3.369
3.064
2.867
LongLoRA
3,918
3.455
3.153
2.956
Theta Scaling
3.864
3,415
3.121
2.934
TAPE (Ours)
3.641
3.196
2.901
2.708
PG-19
LoRA
9.791
9.098
8.572
8.199
LongLoRA
9.989
9.376
8.948
8.645
Theta Scaling
9.257
8.640
8.241
7.999
TAPE (Ours)
8.226
7.642
7.278
7.063
models on passkey retrieval task which has been adopted by many literature (Chen et al., 2023b;a; Tworkowski et al., 2024). We choose a popular open-sourced LLM Llama2 7B (Touvron et al., 2023b), which uses ROPE, the base model and extend it to the 8192 context length. Three base- lines are selected to compare to our TAPE method: vanilla LoRA (Hu et al., 2022), LongLoRA (Chen et al., 2023b), Theta Scaling (Liu et al., 2023).
Our method consistently outperforms all baselines, demon- strating significant improvements, particularly in scenarios with longer context lengths, as observed in QuAL and NQA. In terms of overall performance, xPos is the closest competi- tor to TAPE. While FIRE, RandPE, and ALiBi exhibit good results on a few datasets, they fall short across the board. RoPE struggles with all long-context datasets.
As shown in Tab. 2, TAPE consistently outperforms the other methods across all context lengths on both the Proof-pile and PG19 datasets. On Proof-pile, TAPE achieves a perplex- ity of 3.641 at 1024 tokens, improving over LoRA (3.828), LongLoRA (3.918), and Theta Scaling (3.864). At 8192 tokens, TAPE's advantage grows, reaching 2.708, surpass- ing LongLoRA (2.956), LoRA (2.867), and Theta Scaling (2.934). Similarly, on PG19. TAPE achieves 8.226 at 1024 tokens, improving up to 18.3% over competitors. At 8192 tokens, TAPE reaches 7.063, further showing superiority, especially at longer context lengths.
5.3. Context Window Extension by PEFT
We extend the context window of the pre-trained Llama2 7B model (GenAI, 2023) from 4096 to 8192 using the Red- pajama (Computer, 2023). For validation, we then compare the perplexity on sequence of length 8192, on the cleaned ArXiv Math proof-pile dataset (Azerbayev et al., 2022; Chen et al., 2023a) and the book corpus dataset PG19 (Rae et al., 2019). To further evaluate the models' performance of long context understanding, we report the accuracy of fine-tuned
We also evaluate the passkey retrieval accuracy of our model, following Landmark Attention (Mohtashami & Jaggi, 2023), which has also been adopted by other literature (Chen et al., 2023a; Tworkowski et al., 2024; Chen al., 2023b). In
8

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Passkey Retrieval Accuracy
1.0
0.9
0.8
0.7
0.6
0.5
0.4
1k
2k
3k
4k
5k
6k
7k
8k
LongLoRA
LoRA
Theta Scaling
TAPE
Figure 3. Accuracy on passkey retrieval from 1k to 8k context length with LJama2 7B. We adopt the parameter-efficient fine-tuning strategy for TAPE (see See. 3.3). In contrast to other parameter-effie efficceet fine- tne-fining methods (e.g. LoRA (Hu et al., 2022) and LongLoRA (Chen et al., 2023b)), TAPE achieves no accuracy drop at 8k context length. TAPE performs even on par with full- parameter tuning with Theta Scaling (Liu et al., 2023).
Table Comparison of FL.PS,, MACs, and the number of param- eters for models with different position embeddings.
Table 4. System measurement. We report execution time per step in the Time row and iteration per second in the Throughput row.
The values are averaged over 100 inference steps.
Method
TAPE
RoPE
FIRE
TS's relative bias
TAPE
Method
RaPE
FIRE
TS's relative bias
FLOPS (G)
365.65
321.10
331.97
321.10
w/ Fusion
w/o Fusion
MACs (G)
180.69
160.46
165.69
160.46
Time (x10-
2.56
5.63
2.08
5.56
6.90
Params. (M)
155.33
154.89
154.90
154.90
Throughput
3910
1775
4810
1799
1449
this task, the models are required to locate and retrieve random passkey hidden in a long document. We test the passkey retrieval accuracy ranging from 1 to 8k. The re- sults of long-context passkey retrieval task is presented in Fig. As shown, TAPE consistently achieves near-po sarr-perfect accuracy across all context lengths, outperforming other methods. Theta Scaling shows relatively stable perfor- mance while LoRA and LongLoRA exhibit fluctuating and lower accuracy. Notably, Theta Scaling is widely employed in popular open-source long-context models like Llama3 8B Instruct 262k (AI@ Meta, 2024) and MistralLite (AWS, 2024). TAPE demonstrates a similar superior capability to be applied in long-context tasks.
Flash Atlention
/
X
x
gle A100 GPU. We run 100 inference steps and report the average execution time. Both RoPE and TAPE leverage the acceleration provided by Flash Attention (Dao, 2024b), whereas FIRE and T5's relative bias are not fully compatible with Flash Attention, as it currently lacks support for gradi- ent computation in relative bias. In contrast, we observe that the computations for position embeddings and token fea- tures in TAPE are highly parallelizable, making it suitable for further acceleration using kernel fusion techniques. To capitalize this, we implemented version of TAPE with kernel fusion, referred to as TAPE w/ Fusion. As shown in Tab. 4, the efficiency of the original TAPE (w/o Fusion) already surpasses TS's relative bias and is comparable to FIRE. With additional kernel fusion applied, TAPE achieves 2.2x speedup, approaching the efficiency of RoPE with Flash Attention.
5.4. Efficiency Analysis
In this subsection, we analyze the complexity of our meth- ods in comparison to traditional position embedding tech- niques. Using the models from the pretraining experiment in 5.2, we report three key metrics: FLOPs, MACs, and the number of parameters. The metrics are evaluated with a batch size of and sequence length 1024. As shown in Tab. 3, our architectural modifications introduce a negligible increase in FLOPs, MACs and number of parameters, com- pared to the standard Transformer with RoPE. Moreover, our TAPE is fully compatible with Flash Attention (Dao et al., 2022; Dao, 2024a), widely adopted accelerated at- tention mechanism with IO-awareness, which introduces extra efficiency.
6. Conclusion
This paper introduced TAPE, framework that enhances transformer models by contextualizing positional embed- dings with sequence content across layers. Through incorpo- rating permutation and orthogonal equivariance, we ensured stability and adaptability in positional encoding updates. TAPE can also be easily integrated into existing models, introducing negligible computation and inference overhead. Extensive experiments confirmed TAPE's effectiveness in both arithmetic reasoning and long context language model- ing tasks. One current limitation lies in our exclusive focus on decoder-only models, with limited training scale.
For simplicity. we evaluate the running time of attention layers with different position embedding methods on a sin-

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Acknowledgements
Barrington, D. A. Bounded-width polynomial-size branch- ing programs recognize exactly those languages in nc1, In Proceedings of the eighteenth annual ACM symposium on Theory of computing, pp. 1-5, 1986.
The work reported in this paper was substantially performed using the Princeton Research Computing resources at Prince- ton University, which is a consortium of groups led by the Princeton Institute for Computational Science and Engineer- ing (PICSciE) and Research Computing, PW thanks Kaiyue Wen and Songlin Yang for discussions on the role of posi- tional encoding in circuit representational power, Shenghao Yang for mentioning the relevance with algorithmic reason- ing, and Chuanyang Zheng for discussions on context-awarr position encodings. PW and PL also thank Yinan Huang and Siqi Miao for discussing positional encoding for geometric data. PI Li is supported by NSF awards IIS IIS22239565 and IIS-2428777 for this project. PI Wang is in part supported by NSF Awards 2145346 (CAREER), 021 02133861 (DMS), and the NSF AI Institute for Foundations of Machine Learning (IFML).
Batzner, S., Musaclian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. (3)equivarrian graph neural networks for data-efficient and accurate interatomic potentials. Nature commanications, 13(1):2453, 2022.
2022.
Bevilacqua, B., Nikiforou, K., Ibarz, B., Bica, I., Paganini, M., Blundell, C., Mitrovic, J., and Veličkovié, Neu- ral algorithmic reasoning with causal regularisation. In International Conference on Machine Learning, pp. 2272- 2288. PMLR, 2023.
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O' OrBien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., et al. Pythia: A suite for ana- lyzing large language models across training and scaling. In International Conference on Machine Learning, pp- 2397-2430. PMLR, 2023.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Bischof, C. and Van Loan, C. The wy representation for products of householder matrices. SIAM Journal on Sci- entific and Statistical Computing, 8(1 8(1):s2-s13, 1987.
References
Bogatskiy, A., Anderson, B., Offermann, J., Roussi, M., Miller, D., and Kondor, R. Lorentz group equivariant neural network for particle physics. In International Conference on Machine Learning, pp. 992-1002. PMLR, 2020.
Al@Meta. Llama model card. 2024. URL https://gitthee com htmmmmtalllla 1allaammmmammmaa /1/ama3///// 3lob/main//MDDE CARD md.
Ameisen, E., Lindsey, J., Pearce, A., Gurnee, W., Turner, N. L., Chen, B., Citro, C., Abrahams, D., Carter, S., Hos- mer, B., Marcus, J.. Sklar. M., Templeton, A., Bricken, T., McDougall, C., Cunningham, H., Henighan, T., Jermyn, A., Jones, A., Persic, A., Qi. Z., Ben Thompson, T., Zim- merman, S., Rivoire, K., Conerly, T., Olah, C., and Bat- son, J. Circuit tracing: Revealing computational graphs in language models. Transformer Circuits Thread, 2025. /ubb URL http://///taaaaffeeeereiiii .htll. 2025 20t5///ttr.bbttioongggasps/////
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Butoi, A., Khalighinejad, G., Svete, A., Valvoda, J., Cot- terell, R., and DuSell, B. Training neural networks as recognizers of formal languages. arXiv preprint arXiv:2411.07 2024.
Arora, S. and Barak, B. Computational complexity: mod- em approach. Cambridge University Press, 2009.
Cai, R., Tian, Y., Wang, 7., and Chen, B. Lococo: Dropping in convolutions for long context compression. arXiv preprint arXiv:2406.05317, 2024.
AWS. Mistrallite model card. 2024. URL https:// huggingface.c /tt//aaaoon//ii amazon/Misstrllite
Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., and Zagoruyko, S. End-to-end object detection with transformers. In European conference computer vision, pp. 213-229, Springer, 2020.
Azerbayev, Z., Ayers, E., and Piotrowski, B. Proof- pile, 2022. URL.
com/
Ba, J., Hinton, G. E., Mnih, V., Leibo, J. Z., and lonescu, C. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016.
Chen, S., Wong, S., Chen, L.,, and Tian, Y. Extending context window of large language models via positional interpolation. arXiv preprint arxv:22006. 15595, 2023a.
10

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z.,, Han, S., and Jia, J. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309.120 12307, 2023b.
Learning (ICMI),, 2025. doi: 10448550karXxx.2210001666 URL Utt.httpss//arr/w..orrgg 01866.
Delétang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, E.. Cundy, C., Hutter, M., Legg, S., Veness, J., al. Neural networks and the chomsky hier- archy. arXiv preprint
Chi, T.-C., Fan, T.-H., Ramadge, P.J., and Rudnicky, A. Ker- ple: Kernelized relative positional embedding for length extrapolation. Advances in Neural Information Process- ing Systems, 35:8386-8398 8399, 2022a.
2022.
Deng, C., Litany, O., Duan, Y., Poulenard, A., Tagliasacchi, A., and Guibas, L. J. Vector neurons: A general frame- work for so (3)-equivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 12200-12209 2021.
Chi, T.-C., Fan, T.-H., Rudnicky, A. I., and Ramadge, P. J. Dissecting transformer length extrapolation via the lens of receptive field analysis. arXiv preprint arXiv:2212.10356, 2022b.
Diao, C. and Loynd, R. Relational attention: Generalizing transformers for graph- staphsstruttured tasks. arXiv preprint arXiv:2210.05002, 2022.
Chien, E., Pan, C., Peng, J., and Milenkovic, O. You are allset: multiset function framework for hypergraph neural networks. arXiv preprint arXiv:2106. 0ett21066.3266 13264, 2021.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An Image is Worth 16x16 Words:
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113 2023.
Transformers for Image Recognition at Scale. In Proceedings of ICLR, 2021.
Dudzik, A. J. and Veličković, P.
Graph neural networks are dynamic programmers. Advances in neural information processing systems, 35:20635-2066 20647,, 2022.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv: 2018
Dwivedi, V. P. and Bresson, X. A generalization of
transformer networks to graphs. arxiv. arXiv preprint arXiv:20/2.09699, 2020.
Cohen, T., Weiler, M., Kicanaoglu, B., and Welling, M. Gauge equivariant convolutional networks and the icosa- hedral cnn. In International conference on Machine learn- ing, pp. 1321-1330, PMLR, 2019.
Dym, and Maron, H. On the universality of rota- tion equivariant point cloud networks. arXiv preprint arXiv:2010.02449, 2020.
Computer, T. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023. URL https: http:///ggttuub.cotgeeeeeeeeeeee comm/oogeeheercommteeeeee RedPajama-Data.
Dziri, N., Lu, X., Sclar, M., Li. L., Jiang, L., Lin, B. Y.. Welleck, S., West, P., Bhagavatula, C., Le Bras, R., et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36, 2024.
Cook, S. A. A taxonomy of problems with fast parallel algorithms. Information and control, 24(1-3):2-22, 1985.
Ebrahimi, M., Panchal, S., and Memisevic, R. Your context is not an array: Unveiling random access limitations in transformers. arXiv preprint arXiv:2408. 245.0005000,0 2024.
Dao, T. FlashAttention- Faster attention with better paral- lelism and work partitioning. In International Conference on Learning Representations
CICOaassIICICAAA (ICLR), 2024a.
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z.. Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Trans- former Circuits Thread, 2021. https//transformer http::/ww/oo////aeoooooo/oo/
Dao, T. Flash attention. 2024b. URL https: github. com/ Dta//Dao-ai Becto-AAlb//ea htt.//llasaaaatttiooooo
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAt- tention: Fast and memory efffcient exact attention with 10-awareness. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
de Luca, A. B., Giapitzakis, G., Yang. S., Veličković, P., and Fountoulakis, K. Positional attention: Expressivity and learnability of algorithmic computation. In Proceed- ings of the 42nd International Conference on Machine
Engelmayer, V., Georgiev, D. G., and Veličković, P. Parallel algorithms align with neural execution. In Learning on Graphs Conference, pp. 31-1. PMIR, 2024.

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Fuchs, E, Worrall, D., Fischer, V., and Welling, M. Se (3) trannfforrmrs:: 3d noto-translation equivariant attention networks. Advances in neural information processing systems, 33:1970-1999 2020.
Huang, Y., Lu, W., Robinson, J., Yang, Y., Zhang, M., Jegelka, S., and Li, P. On the stability of expressive positional encodings for graph neural networks, arXiv preprint arXiv:2310.02579, 2023.
GenAI, M. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
Huang, Y., Miao, S., and Li, P. What can we lear from state space models for machine learning on graphs? arXiv preprint arXiv:2406.058: 05885, 2024.
Georgiev, D. G., Lio, P., Bachurski, J., Chen, J., Shi, T., and Giusti, L. Beyond crdos-renyi: Generalization in algorithmic reasoning graphs. In The Second Learning on Graphs Conference, 2023.
Ibarz, B., Kurin, V., Papamakarios, G., Nikiforou, K., Ben- nani, M., Csordás, R., Dudzik, A. J., Bošnjak, M., Vitvit- skyi, A., Rubanova, Y., et al. A generalist neural algorith- mic learner. In Learning on graphs conference, pp. 2-1. PMLR, 2022.
Golovneva, O., Wang, T., Weston, J., and Sukhbaatar, S. Contextual position encoding: Learning to count what's important. arXiv preprint arXiv:2405.18771 2024.
Immerman, and Landau, S. The complexity of iterated multiplication. Information and Computation, 116(1): 103-116, 1995.
Grazzi, R., Siems, J., Zela, A., Franke, J. K., Hutter, F., and Pontil, M. Unlocking state-tracking in linear rmss through negative eigenvalues. arXiv preprint 1222:2441.12537, 2024.
Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas. D. d. I., Bressand, F., Lengyel, G., Lample, G., Saulnier, I.., et al. Mistral 7b, arXiv preprint arXiv:2310.06825 2023.
Haviv, A., Ram, O., Press, O., Izsak, P., and Levy, O. Transformer language models without positional encod- ings still learn positional information. arXiv preprint arXiv:2203. 16634, 2022.
Kazemnejad, A., Padhi, I., Natesan Ramamurthy, K., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
He, 7., Feng, G., Luo, S., Yang, K., He, D., Xu, J., Zhang. Z.,, Yang, H., and Wang, L. Two stones hit one bird: Bilevel positional encoding for better length extrapolation. arXiv preprint arXiv:2401. 16421, 2024.
Kitaev, N., Kaiser, Ł., and Levskaya, A. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0444 2020.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021.
Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, Y. W. Set transformer: A framework for iitention-based permutation-invariant neural networks. In International conference on machine learning, pp. 3744-3753. PMLR, 2019.
Hinton, G. E. and Anderson, J. A. Parallel models of as- sociative memory: updated edition. Psychology press,
Lee, N., Sreenivasan, K., Lee, J. D., Lee, K., and Papail- iopoulos, D. Teaching arithmetic small transformers. arXiv preprint M.XXi:2307.038// 2023.
2014.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference machine learning, pp. 8867- 8887. PMLR, 2022.
Li, S., You, C., Guruganesh, G., Ainslie, J., Ontanon, S., Zaheer, M., Sanghai, S., Yang, Y., Kumar, S., and Bho- janapalli, S. Functional interpolation for relative posi- tions improves long context transformers. arXiv preprint arXiv:2310.0444 2023.
Hopfield, Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2554-2550 1982.
Li, X., Liang, Y., Shi, Z., Song, Z., Wang, W., and Zhang, J. On the computational capability of graph neural networks: A circuit complexity bound perspective. arXiv preprint 2025.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z.,, Li, Y., Wang, S., Wang, L., and Chen, W. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL https:// openrevice
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, Transformers learn shortcuts to automata. arXiv preprint arXiv:2210. 10749, 2022.
12

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Liu, B., Ash, J., Goel, S., Krishnamurthy, A., and Zhang, C. Exposing attention glitches with flip-flpp language modeling. Advances in Neural Information Processing Systems, 36, 2024.
Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv: 20t.:::20...2440 12409, 2021a.
Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv: 108. 12409, 2021b.
Liu, X., Yan, H., Zhang, S., An, C., Qiu, X., and Lin, D. Scaling laws of rope-based extrapolation. arXiv preprinl arXiv:2310.05209, 2023.
Rae, J. W., Potapenko, A., Jayakumar, S. M., and Lillicrap, T. P. Compressive transformers for long-range sequence modelling, arXiv preprint arXiv: /9/ 05500,557200 2019.
Maron, H., Ben-Hamu, H., Shamir, N., and Lipman, Y. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2018.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., Liu, P. J., et al. Exploring the limits of transfer learning with unified text-to-text transformer. J. Mach. Learn. Res., 21(14)):1677 2020.
McLeish, S., Bansal, A., Stein, A., Jain, N., Kirchenbauer, J., Bartoldson, B. R., Kailkhura, B., Bhatele, A., Geiping, J., Schwarzschild, A., and Goldstein, T. Transformers can do arithmetic with the right embeddings. arXiv preprint 1rXiv:2405.17399 2024a.
Rahimi, A. and Recht, B. Random features for large-scale kemel machines. Advances in neural information pro- cessing systems, 20, 2007.
McLeish, S., Bansal, A., Stein, A., Jain, N., Kirchenbauer, J., Bartoldson, B. R., Kailkhura, B., Bhatele, A., Geip- ing, J., Schwarzschild, A., et al. Transformers can do arithmetic with the right embeddings. arXiv preprint mrXiv:2405.17399 2024b.
Rajani, N. F., McCann, B., Xiong, C., and Socher, R. Ex- plain yourself! leveraging language models for common- sense reasoning. arXiv preprint arXiv:1906.02361, 2019.
Merrill, W. Sequential neural networks as automata. arXiv preprint NEKV:19906006131 2019.
r/LocalLLaMA. / //www. http://tt..ooo////oaaaaaaaa comments/141227///ttkaeee rope_a http:/owwwwlammmmdds___sss_sssee
Ntk-aware scaled rope.
https:
Merrill, W. and Sabharwal, The parallelism tradeoff:
2023.
Limitations of log-pree res-prccision transformers. Transactions of the Association for Computational Linguistics, 11:531- 545, 2023.
Ruoss, A., Delétang, G., Genewein, T., Grau-Moya, J., Csordás, R., Bennani, M., Legg, S., and Veness, Ran- domized positional encodings boost length generalization of transformers. In 61 Annual Meeting of the Associa- tion for Computational Linguistics, 2023.
Merrill, W., Petty, J., and Sabharwal, A. The illusion of state in state-space models. arXiv preprint 2rXiv:2404.08819, 2024.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International con- ference machine learning, pp. 9323 93322. PMLR, 2021.
Mohtashami, A. and Jaggi, M. Landmark attention: Random-acees access infinite context length for transformers. arXiv preprint arXiv:2305. /6300, 2023.
Pagiamtzis, K. and Sheikholeslami, A. Content-addressable memory (cam) circuits and architectures: A tutorial and survey. IEEE journal of solid-state circuits, 41(3):112- 727, 2006.
Schütt, K., Unke, O., and Gastegger, M. Equivariant mes- sage passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-9388, PMLR, 2021.
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Cao, H.. Cheng, X., Chung, M., Grella, M., GV, K. K., et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.1 13048, 2023a.
Shaham, U. Segal, E., Ivgi. M., Efrat, A., Yoran, O., Haviv, A., Gupta, A., Xiong, W., Geva, M., Berant, J., et al. Scrolls: Standardized comparison over long language sequences. arXiv preprint arXiv:2201.032. 203523333.0333,,,,0 2022.
Peng, B.. Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 2023b.
Shaw, P., Uszkoreit, J., and Vaswani, A. Self-attention with relative position representations. arXiv preprint arXiv: 10X3::8803.025550 2018
Peng. B., Zhang, R., Goldstein, D., Alcaide, E., Du, X., Hou, H., Lin, J., Liu, J., Lu, J., Merrill, W., al. Rwkv-7" goose" with expressive dynamic state evolution. arXiv preprint Arviv:2003. /4456, 2025.
Sinha, K., Jia, R., Hupkes, D., Pineau, J., Williams, A., and Kiela, D. Masked language modeling and the distribu- tional hypothesis: Order word matters pre-training for little. arXiv preprint arXiv:2 104. 06644, 2021.
13

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Sinha, K., Kazemnejad, A., Reddy, S., Pineau, J., Hupkes, D., and Williams, A. The curious case of absolute posi- tion embeddings. In Findings of the Association for Com- putational Linguisties: EMNLP 2022, pp. 4449-4472, 2022.
Vollmer, H. Introduction to circuit complexity: uniform approach. Springer Science & Business Media, 1999.
Wang. H., Yin, H., Zhang, M., and Li, P. Equivariant and stable positional encoding for more powerful graph neural networks. arXiv preprint arXiv:2203.000 20090099990000 2022a.
Su. J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063. 2024.
Wang. J., Ji, T.,, Wu, Y., Yan. H., Gui, T., Zhang, Q..
Huang, X., and Wang, X. Length generalization of causal transformers without position encoding, arXiv preprint arXiv:2404. /2224, 2024a.
Sun, Y., Dong, L., Patra, B., Ma, S., Huang, S., Ben- haim, A., Chaudhary, V., Song, X., and Wei, F. A ettetleetraaplltabble transformer. arXiv preprint arXiv:2212. 10554, 2022.
Wang, P., Yang, S., Liu, Y., Wang, Z., and Li, P. Equivariant hypergraph diffusion neural operators. arXiv preprint arXiv:2207. 06680, 2022b.
Tang, H., Huang, Z., Gu, J., Lu, B.-L., and Su, Towards scale-invariant graph-related problem solving by iterative homogeneous gnns. Advances in Neural Information Processing Systems, 33:15811-1582 115822, 2020.
Wang, X., Li, P., and Zhang, M. Graph as point In Forty-first International Conference Machine Learn- ing. 2024b.
Thomas, N., Smidt, T., Kearnes, S., Yang. L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and translation-eqqivrrian neural networks for 3d point clouds. arXiv preprint 102.02180,.08219, 2018.
Wei, J., Wang. X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, et al. Chain-of-thouge prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24833 2022.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A,, Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.13977 2023a.
Weiler, M. and Cesa, G. General e (2)-equivariant steerable cnns. Advances in neural information processing systems, 32, 2019.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288. 2023b.
Worrall, D. E., Garbin, S. J., Turmukhambetov, D., and Brostow, G. J. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 00028-503, 2017.
Tworkowski, S., Staniszewski, K., Pacek, M., Wu, Y., Michalewski, H., and Mitos, P. Focused transformer: Contrastive training for context scaling. Advances in Neural Information Processing Systems, 36, 2024.
Yan, Y., Swersky, K., Koutra, D., Ranganathan, P., and Hashemi, M. Neural execution engines: Learning to execute subroutines. Advances in Neural Information Processing Systems, 33:17298-17308 2020.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, t.., and Polosukhin, I. Attention is All You Need. In Proceedings of NeurIPS, 2017.
Yang, S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Par- allelizing linear transformers with the delta rule over se- quence length. Advances in neural information process- ing systems, 37:1 1e:115491-111 15491-115522 2024.
Veličković, and Blundell, C. Neural algorithmic reasoning. Patterns, 2(7), 2021.
Yang, S., Shen, Y., Wen, K., Tan, S., Mishra, M., Ren, L., Panda, R., and Kim, Y. Path attention: Position encoding via accumulating householder transformations. arXiv preprint arXiv:2505. 16381, 2025.
Veličković, P., Badia, A. P., Budden, D., Pascanu, R., Ban- ino, A., Dashevskiy, M., Hadsell, R., and Blundell, C. The clrs algorithmic reasoning benchmark. In Inter- national Conference on Machine Learning, pp. 22084- 22102. PMLR, 2022.
Yarotsky, D. Universal approximations of invariant maps by neural networks. Constructive Approximation, 55(1): 407 407,474, 2022.
Villar, S. Hogg, D. W., Storey-Fisher, K., Yao, W., and Blum-Smith, B. Scalars are universal: Equivariant ma- chine learning, structured like classical physics. Advances in Neural Information Processing Systems, 34:28848- 28863, 2021.
Yu, X. x., Suresh, A. T., Choromanski, K. M., Holtmann- Rice, D. N., and Kumar, S. Orthogonal random features. Advances in neural information processing systems, 29, 2016.
14

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S. J., and Kumar, S. Are transformers universal approximators of sequenne--o--eqquencc functions? arXiv preprint arXiv:1912.10077, 2019.
Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R., and Smola, A. J. Deep sets. Ad- vances in neural information processing systems, 30, 2017.
Zhang, R. Making convolutional networks shift-invariant again. In International conference machine learning, pp. 7324-7334. PMLR, 2019.
Zhao, H., Jiang, L.,, Jia, J., Torr, P. H., and Koltun, V. Point transformer. In Proceedings of the IEEE/CVF Interna- tional Conference Computer Vision, pp. 16259-16268, 2021.
Zheng, C., Gao, Y., Shi, H., Huang, M., Li, J., Xiong, J., Ren. X., Ng, M., Jiang. X., Li, Z., and Li, Y. Dape: Data-adaptive positional encoding for length extrapola- tion. In Advances in Neural Information Processing Sys- tems (NeurIPS), 2024.
Zhou, Y., Alon, U., Chen, X., Wang, X., Agarwal, R., and Zhou, D. Transformers can achieve length generalization but not robustly. arXiv preprint arXiv:440..0937// 2024.
15

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
A. More Related Work
Context Length Extrapolation. The length extrapolation ability of transformers are limited mainly in two aspects: (1) the high memory usage caused by quadratic memory usage; and (2) the poor generalizability to unseen sequence length during inference. To address the memory usage during long sequences training, LongLoRA (Chen al., 2023b) introduced shifted sparse attention and leveraged parameter-efftcient tuning. LoCoCo (Cai et al., 2024) introduce KV cache compression mechanism. To help generalizability of positional embedding to unseen sequence length, Chen et al. (2023a) explores zero-shot linear interpolation on rotary embedding; r/LocalLLaMA (2023); Peng et al. (2023b) enhance simple interpolation by retaining high-frequency encoding ability; Liu et al. (2023) investigates the relationship between rotary base and extrapolation ability. While the previously mentioned methods focus primarily on extending rotary positional embeddings, Li al. (2023) introduced a functional relative position encoding framework that enhances generalization to longer contexts. However, these methods generally impose fixed pattern on attention maps, often adopting a decaying pattern based distance. In contrast, we propose a learnable and generic position encoding framework that primarily focuses arithmetic reasoning.
Equivariant Learning. Equivariant machine learning is a broad field that leverages task-specific symmetries to introduce inductive biases into neural networks, reducing learning complexity and improving generalization. Here, we focus on foundational works in this domain that are highly relevant and provide key motivation for our study. Prior research has primarily focused on data representations with intrinsic symmetries. such as graphs (Satorras al., 2021: Schütt et al., 2021; Batzner et al.,, 2022; Maron et al., 2018), hyper-graphs (Wang et al., 2022b; Chien et al., 2021), and point clouds (Zaheer et al., 2017; Fuchs et al., 2020; Thomas et al., 2018; 8; Hoogeboom et al., 2022), which primarily require permutation equivariance. Recent progress models graph representation learning as a joint invariance of permutation and orthogonal groups (Wang et al., 2022a; Huang et al., 2023; Wang et al., 2024b). Beyond geometric data, another stream of work (Worrall et al., 2017; Zhang, 2019; Weiler& & Cesa, 2019; Cohen et al., 2019) ensures symmetric inputs yield consistent outputs, reflecting the same label under symmetric transformations. For instance, Worrall et al. (2017) introduce onation-equivariant feature transformations after mapping images to the continuous fourier domain, while Zhang (2019) enhance translation invariance in CNNs by incorporating low-pass filter in the pooling layer. To the best of our knowledge, we are the first to introduce equivariance in language models, bridging the geometric properties of positional embeddings to modeling complex language tasks (e.g., retrieval and reasoning) and learning stability.
Concurrent works on Contextualized Positional Encoding. Recent studies have moved beyond fixed absolute or relative encodings toward context-adaptive position signals. DAPE (Zheng et al., 2024) pioneers the design of context-dependent positional embeddings. where positional information is decomposed into a learned prior modulated by the current attention context. This yields a data-adaptive signal that improves both in-length performance and length extrapolation. A subsequent work CoPE (Golovneva et al., 2024) proposes positional contextualization that accumulates intervals determined by token features to represent token positions in the sequence. PaTH Attention (Yang et al., 2025), parameterizes positional effects as an accumulated product of input-conditioned Houscholder- transformations, offering a compact parallel formulation and a FlashAttentionn-tty blockwise implementation. Importantly, Yang et al. (2025) highlights the role of positional encoding in enhancing the algorithmic representational power of transformers and strengthening their computational completeness. Our theoretical results are highly motivated by Yang et al. (2025).
Algorithmic Reasoning and Geometric Learning. A large volume of work studies training neural networks to execute algorithms. In these settings, models take as input various data structures - often encoded as graphs and compute solutions to tasks such as path finding, graph connectivity, and dynamic programming (Veličković & Blundell, 2021; Ibarz et al., 2022; Yan et al., 2020; Veličković et al., 2022; Dudzik & Veličković, 2022; Bevilacqua et al., 2023; Georgiev et al., 2023; Engelmayer et al., 2024; Diao & Loynd, 2022: Tang et al., 2020). A recent line of work investigates transformers' capabilities as general-purpose algorithmic performers (Delétang et al., 2022; Butoi et al., 2024). More recently, Positional Attention (de Luca et al., 2025) formalizes transformer architectures whose attention weights depend solely on positional encodings rather than token features. At high level, attention endows models with the capacity to realize computation graphs over tokens (Elhage et al., 2021; Ameisen et al., 2025). The structure of our proposed TAPE can be viewed as natural way of embedding the discrete, symbolic, and input-dependent graph structures of algorithms into continuous representations with orthogonal-gret -gropp equivariance. A similar transformation has also been seen in designing graph transformers with Laplacian embeddings Dwivedi& & Bresson (2020); Wang et al. (2022a). The duality between graph topological representation and geometric embeddings has been investigated in Wang et al. (2024b).
16

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
B. Deferred Proofs
Notations. The vector ò; denotes the i-th canonical basis vector, i.e., a vector with in the i-th position and 0 elsewhere. Its length can be often inferred by context without cause of ambiguity. We denote the permutation group over N entries as [I(N), and orthogonal group of Rt dimension as O(R). Associated with permutation matrix P E II(NV, we represent the We also work with the swap (i.e., transposition) group S(N) C IT(N) which exchanges up to two elements in []. Note that S(N) is generating set of [[(N).
permutation map as function [N] - [N], i.e., P. = 80())
B.1. Formal Statement and Proof for Theorem 4.1
In this section, we formalize the statement of Theorem 4.1 with the hypothesis class yielded by transformers with TAPE and the precise definition of representational power of circuits.
B. ARCHITECTURAL FAMILY
First of all, we define a family of transformer architectures equipped with generalized TAPE subject to invariance principles (Eqs. 4 and 5). Recap that we denote a sequence of token features as X E X' S N positional encoding as E E S RN and masking matrix M E E C RNN Without loss of generality, we assume the shape of token and positional representations is unchanged in hidden layers. Following our general formulation in Sec. 3.1, each layer of such transformer consists of two components: (1) token mixing layer f()) (x (1-1) (E MM) + X (9) and (2) position contextualization layer g() (x('-1) E + E where we use superscript to denote the layer index. The token mixing layer follows standard design of the attention block with residual connections and element-wise feedforward mapping: f() = JFFN o fALun such that fAlm E.) = Σ/= A M),, W
v Dy + Res (!)
JAna æi
where Res E {0.1) denotes whether skip connection is adopted across the module * at the 1-th layer, and A is a conventional attention modulated by positional encoding:
My exp 2 W
Q (e,,e, k Wyæ,
( E,
Σラ Ma exp x. W (e W (1) xTy)
We denote W
W W (1) V E RCx as key, query, and value matrices, and di0) MXLXR R MXLXR + RCXC injects information of positional encoding into attention logits. In addition, JEPN JFFN (X); (F (a:)) + Res (1)
K
frmm 2z with
(FN(") element-wise MLP of two layers, independent of positional encoding. Clearly, permutation equivariance holds due to the nature of attention, while e²(1) needs to be ((R)-invariant to let f invariant to orthogonal groups (Eq.
5). Likewise, we consider g()) as a sequence-level propagation composed with an element-wise feedforward mapping:
(1) (0) JFFN DAttn where JAmn [X, M); Σÿ= (X, E, UP e, + Res! S/amm e; performs pairwise aggregation where w E ML and the aggregation weights Bdd M),, E utilize both contextual and positional information, and 9FFN (X, E); = (x,)e; + Res (1) e; modulates positional embeddings by token representations point- wisely with function (5(); EC + M.MLMM To have g satisfies Eqs. 4 and 5, B(()) needs to satisfies the following equation for every Pe II(N) and R O(R): B (PX, PER, PMP PB()) [X B.M.) M )
(1) (2)
B. 1.2. 11(5) GROUP WORD PROBLEM
The group word problems are family of state tracking tasks, asking that, given sequence of elements sampled from a generating set of a finite group, does it evaluate to the identity element after reduction? Formally, given a group (G,) and its generating set G", we form a sequence 91, ... .N such that 9i E G for every i E [N]. The decision problem is whether 91 .... !N equals identity. The circuit complexity for this decision problem depends on the algebraic structure of the target group G. is known that the the word problem with any non-solvable group is strictly -cooopeee (Immerman & Landau, 1995; Barrington, 1986). In particular, typical example for non-solvable group is I[55,, the permutation group over 5 entries (Merrill et al., 2024; Peng et al., 2025). We call this specific problem I1(5) group word problem.
To this end, we focus on handling the sequences synthesized from S(5)- generating set of (1()). Note that |S(5) = 10. Without loss of generality, we fix a canonical order in S(5) and define the n-th element in S(5) as [src(n) + dst(m)) (suppose src(n) < dst(n)) or S, E R5x5 in matrix form for 1 S n S 10. Consider a vocabulary V : [11] - {(bos]
where 1,.., 10 are mapped to S(5) and 11 is mapped to the "beginning of sequence" token [bos). The input to the model
17

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
sequence = (u[1),... ... u[N] e [11j%, where fix u[1] = 11 (the [bos] token). The input sequence can be then illustrated as:
[bos], [src(u[2]) + dst(u[2])), .M [src(u[[]] + dst(u[N]]].
We denote the resultant permutation after reduction over the beginning transpositions as ग, : [3] + [5], which corresponds
to the matrix form = Sulj- The ground-truth labels are a boolean sequence (y[1)... =((N) E {0.1) 1) such that
u[i] = 1 if ny is identity, otherwise y|i] = O. We call model is capable of solving (1( group word problem for arbitrary
input size N and imputs (u[1),.-. ..- N[], the model can precisely predict the ground-trum truth labels [[[]).....lll
B.1.3. MAIN RESULT AND PROOF
To show Theorem 4.1, it is equivalent to prove the following theorem, since (1(5) group word problem is NC¹ NC'-complets thus all the problems within NC¹ family can be reduced to under AC reductions.
Theorem B.1. There exists constant-depti and constant-widd transformer architecture from the model class defined in Sec. B.I.I, subject to permulation equivariance defined in Eq. 4 and orthogonul equivariance defined in Eq. 5, such that it is capable of solving I1(5) group word problem, defined in Sec. B.1.2.
Proof. Proof by the following construction.
(0) Embedding Layer. We represent input sequence u as a matrix of one-hot vectors: (0) E {0, 1) where x We initialize tensorial positional encoding ( E RNCx6 i.., M = (thus omitted), L = C, R = 6, such that E = | [ for every i E [N], where E E RC6 E; = (ōsrc(e) )du(t)) E IR6 for 1 <i 10, and E = Og- In addition, we denote M E RNN as the masking that describes the topology and dependencies among tokens. For word problems, M is lower-triangular matrix, i.e., Mis = when j S while M,, = when j > i.
6m((]] E R (i.e., ==
= 11).
The First Layer. We leverage the identity token mixer ((1) = SAeme and FFN = SFFN., as in Lemma B.4 such that x (0) For positional contextualization layers, skip the sequence-level positional contextualizer, i.e.,
(1)
JAttn = ld by letting B(I) = and (25(1)
(1)
gamn = 1, while leveraging the point-wise positional contextualization with function
x
(0) x +
0
0
16x6
(O)T x
(O)T
(1) = 9FFN (X) E{
(1)
0
0
e
(0)
0
Og T6x6
= I [
0g
E
R
I 16x6_
I 6
Since the effective operations are only point-wise, permutation invariance in Eqs. 4 and 5 holds for this layer. Orthogonal group invariance can be also verified by seeing e R gives e(") R for every R E 0(6).
The Second Layer. We again leverage an identity attention block /12 SAtemu (Lemma B.4) to process token features. Slightly different from identity feedforward layer fFFNu we intend to construct f(2) such that a new zero channel is concatenated to the inputs, i.c., x/²) [0 x (O'TTT E CC+, We achieve this goal by setting the first linear layer the
identity matrix, activation function as ReLU, and the second linear layer as [0c Icxc] Note that ReLU does not change inputs here because all inputs are in the positive orthant.
Furthermore, we consider the following sequence-level contextualization (2) J with Res (2) 1 for positional embed-
dings. First, we compute pre-masking "attention" logits by "key" and "query" matrices U(22 = U(2) 5 such (1)
that (x() ( (U e
(U
e T ول Eup! E R, and next we define B x (1) E(1) M)
Kee
[ (INXN+ + M O (((( E()))
We let ( [0s 81 0sxe] T, and then the output of sequence-level
ij
18

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
contextualization be written as:
0
E u
M IINN+ M & )
= bi
0g
R
[6x6_
6x6_
I Lx6
where we define
N
INNN + M B (x(1, E
Eubl E RR
7:
(2) We further adopt feedforward contextualization JFFN via v/( +
[ x)102
( z)6d2
83
8s]
x) 1
Eu
c
2 9FFN (X (2))
( 2)66
n.
RR 12x6
83
n:
I 1x6
Tx6
[6x6_
6g
For the sake of coherence, we defer the justification of invariance for this layer to Lemma B.6.
The Third Layer. We leverage the weight configuration of Lemma B.5 to construct the attention layer: [(3] s and re-use the attention weights for cross-token positional contextualization (similar to Eq. 7):
B(3)
X (2)
E M);;
Further we let U (3) 112x12 and W (3) 0
6
A
X
M) kij
We enable residual
v
Oc
Ocxc
connection for token-mixing attention Res Jaca
(3)
1, while disabling it for cross-token contextualization Res (3)
0:
gAse
a((3)
1
M,
(
(3)
e
(2)
1
M
(2)
1
Mij Mij
Euii
E R
(9)
M
JM
I I
j=1
x(33)
1
N
Mij W V 2
x
(2)
Mij
1
(0)
811
+ +
(2)
E RC++
(10)
|
[
Oc
We note that for the word problem, M is always lower triangular matrix, and only the beginning token [bos] can equal 11,
i.e. u[1] 11, u[i] o 11 for all > In this scenario, we can express E and X (3)
below:
(3) = + Eulj (2)
Ej
(3) =
1/i x
E RC+
(1I)
I 16x6
Moving forward, we choose the identity feedforward mapping for element-wiss token transformation: JFFN fFFN
(3)
i.e. x; (3)
[1/i
x [0)T JT. The element-wise positional contextualization JFFN is given by a nonlinear transformation
(3)
(8, x) T6x6 T6xe]:
E E.g = 1666 I 16x6
OFFN
(3)
== [- - 1Ix6
[6x6]
Σ E j=1
(12)
where the second equality holds for the word problem case. By Lemma B.5, we show that Eqs. and 10 satisfy permutation cquivariance (Eq. 4). Cross-token contextualization (Eq. is orthogonal equivariant because B(3) invariant (or
even independent of) E(2), and Eq. 9 a linear combination of e(2)'s. The point-wise contextualization a left matrix
multiplication on E (2) (Eq. 12), preserving orthogonal equivariance as well.
19

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Significantly, by Lemma B.7, show that the transpose of Eq. 12 is cumulative Householder product (Grazzi et al., 2024; Yang et al., 2024; 2025):
S
0
0 1
() II [6x6
=uny - - mb
0
j=1
where the second equality is by Lemma B.8 (Peng al., 2023a).
The Fourth Layer. We follow the proof of Yang al. (2025), Theorem 2.1 construct the output layer. For the attention,
we define weight matrices for A
W (4) 06x 10
61 26:
363
484 + 50,
56] [cc Icxo] E R6x(C+1),
W (4) (ő) + 2-2 + 303 + 484 5 + +5.556))
Oc Ixxol E R
W (4) [((++))x10 - 61] 00c
Icxo] E [CC+))((C+1)
where [0c Icxc] is applied to extract the lower block of a [0c Iexe] [1/ x; ! x; a;!), We also utilize the positional encoding via g(4) defined as below:
(3) E 66x6
First of all, (6() is invariant to orthogonal transformations on {e; (3) we analyze the pre-softmax attention logits A (x (3)
[i[ due to its inner-product form, so is ((4) fAttn' Next,
E3) M) = (W q ( (e) ,e (W xx (8)). For every
i [N] and j * 1, Ā (3)
,M = 0 because W x (3) = 0g. For j =
1
2
A (3) ( MMN.1
TT j=1 Swil
0
3
5 (n,(V) -4.55
1
5
54.5
where we use the fact that ei
(3) 1x6 and 21
(3) 6T.]T. Observe that Σ (v) S 55 with equality hold if and
only if T. (v) = v for every v E [5). i.e., n, identity permutation. This is, (x E()) M) 2 0.5 if m identity, and (1) (X (3) (3) M) S -0.5 otherwise. After softmax normalization, followed by applying the attention, value matrix, and skip connection (i.e., (es = 1), the output of this attention layer becomes:
fame
x
A (4)
x (3)
E M) W
(4) x. + x.
(3)
(3)
(x (x ( M 1) x. (3)
expt i x E M).))
exp(A14((X gin M
E c
One can see that if π; is identity, then (x (x E(3) M).1) 2 0. and 2 Otherwise, (3) E M)
(4) 0. By this observation, we construct JFF to read the first element from a!") and map the corresponding and [.
label via two-layer MLP: adopt -ē the first linear layer, hard sigmoid min(1, max(0,.)) -)) the non-linear activation. and identity as the second layer. The final output can be analytically written as:
if n. is identity otherwise
BardSigmoid(-bb )=b
which concludes the proof.
20

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Remark B.2. Different from the construction of Yang al. (2025), Theorem 2.1, do not encode the sequence length in the query matrix. In fact, if the construction is length dependent, the architecture could potentially lose uniformity of the target circuit family when input size is unbounded. We note that the length dependency in the construction of Yang et al. (2025), Theorem 2.1 is likely due to the consideration of the computational precision of attention. Under the similar setting, is necessary to modify v/(3) as the function: æ + [- 2 T T to scale the pre-softmax attention logit by n. Remark B.3. We emphasize that our proof constructive and does not rely on the universal approximation property of Attention (Yun et al., 2019) or MLPs (Yarotsky, 2022). As a result, the network has constant width and depth and represents a precise classifier solving the (5) word problem.
B.1.4. AUXILIARY RESULTS
Lemma B.4. (Identity mappings.) There exists an attention layer faume and feedforward layer JEFN, that meets Eq. 4 and Eg. such that SAttma X, E, M) = X and JFFN (X) = X for every inputs X, E, M.
Proof. Proof by construction. For the attention, we let W Ocxe in attention and Res fxum 1. For the feedforward layer, we choose the associated F-N as a two-layer MLP with the last weight matrix all zeros, and let Re5 = 1. Eq. 4 and Eq. 5 hold automatically.
Lemma B.5. (Neighhorhood mean.) There exists an attention layer fAttnng with free value matrix Wv that meets Eq. 4 and Eq. 5 such that fAmtming (X, E,E, = (E Mi, W v. / Msp).
Proof. Proved by letting Wa Q and W k all zero and disable the residual connection Res farmAgg 0. Permutation equivariance are satisfied automatically due to the nature of masked attention. To be more specific, for every permutation matrix P E II(NV) and its corresponding permutation map: : [N] [N]:
fAttnAnn (PX,PE, PMP 2j=1 M (0)=00) W VV E/ M (()) Wyz; Σp=1Ma M
fActnaug (X, E, E,M, (()): xi))
E M(0,0))
where the second equality is because exchanging terms in summation does not change the outcome. Hereby, we have shown fALtnann (PX, PMP f (X, E, M). Orthogonal invariance is given by the independence of E.
(2) Lemma B.6. The second positional contextualization layers g and 9FFN Theorem B.1 satisfy permutation equivariance (Eq. 4) and orthogonal equivariance (Eq. 5).
Proof. First of all, 9FFN (X. E, M) is elesent-wise, thus satisfies permutation equivariance (Eq. 4). Also, OFFN [X, E, M)
(2)
(2)
is left matrix multiplication on every positional encoding e,, hereby, orthogonal equivariant (Eq. 5). It remains to
demonstrate the two invariances for J For every P E TT(N) and its corresponding permutation map æ [N] - [N], have that:
(2)
(PX, -x(1)e
= ((2) (X, E) (()
which indicates 3
(PX, 2 (X. E.P.P Next we analyze 222) X,E,M)
BCa) PP. PMP
[I+ [I+PMP O (PX = -(I+ PMP O P
-IPP
-I+P M M BOX pT) -1 =
PT)
(X.
PT.
21

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
By above derivation, we have:
N 9Atmn PX, PE,PPM ); = Σ
N
PE )(C((())
Σ
PB (x P (X,
Ca(i)
i=1
1=1
N Σ
(
(x, E, a()
( E, M) mi)
j=1
j=1
(2) Att
(X E, M m()))
which proves the permutation equivariance: Jeten equivariance, we note that, for every R E (OR):
(2)
PE,PPM
Pg (2) (X, E, M). To show orthogonal
B X, ER) = e, (1)
RR ej = e.e = = (X,
This implies ((2) ER) ((2) (x, E) and henceforth,
N
AAttn X, ER M); = Σ B (X, (R Σ (2) (x, E, R 9Amm E, R)
R
==1
which shows the orthogonal equivariance.
Lemma B.7 (Bischof & Van Loan (1987); Yang el. (2024)). Consider sequence of vectors E = the following equation holds for every i [N]:
Ei -...... E.v'" ERNR,
i
II 3=1
[RXR - = RXR Σ [ M o [ 3 E] j=1
where M is lower-triangular mask matrix, ie.., Mij 1 ifj S i, otherwise Mij 0.
Proof. The proof follows from the standard derivation of WY representation for accumulating Householder matrices (Bischof & Van Loan, 1987). A clear derivation can be seen in Yang et al. (2024). First of all, let w, = & )E (,),,, E R for 2 2 and w1 1 We W W1,'"', N E R Then expressing W in matrix form:
W = = M + EE W.
which gives W = NINXN + M o EE ) '.. We denote P; = IT (T - E,,_]) and perform the following induction.
The induction hypothesis P, I IRXR n Σj= w Obviously, P1 RRX E i 2, we can derive that:
I Inxn - w Then for
P,
II j=1
I Tnxn
E P.- (Inxn- -
InxR Σ w
(IuxR- E,E, -
i-1
=IRxR
Σ w - - + Σ
w.s E.ET = I IRXR
Σ
3. Σ
E,)e, ET
f=1
f=1
==1
j=l
InxR
Σ
w.eT
I n
Σ w.y.
which justifies the hypothesis. We completes the proof by writing w, in the matrix form:
Σ [TNxN+
i
MMxN++M@EB O BETE = E
P; IRXE
(13)
as desired.
22

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Lemma B.8 (Peng al. (2025)). Given a swap group of elements S(n), every swap <> f). its matrix representation
S E S(n) be written as 5 I, (ō; 5;)(6, 6;)
Proof. Proof by noticing that S, ,,. S ò and S 6% for every + i and & # j.
B.2. Formal Statement and Proof for Proposition 4.2
We first formulate a general computational paradigm of transformers with TAPE.
Inputs. Let ( NxC be the input sequence of token embeddings. We consider two types of initialization of positional encoding:
RoPE. Construct positional embeddings as a tensor: E(0) E BRNNC/22XX22 For every i E [N], m E [C/2], we let (0) cos((m)))
sin(fm:) where 8m = - -100002m/c.
e,
sin(0nn
cos(mmi)
((0) &RNXXRR For every Reweighted Random Fourier Features. Construct positional embeddings as a tensor: E [N], E [M], E [L]
(0) m
2 R
Wm,dr cos(@m, Wm.d. Sin..,. 1 , [R/2]
where (0m.r [0mm.mmMM)) are often chosen as (oute--MMonte-Carloo samples from a distribution, and (wmarh (oeerrrrmm())).ee))eeemmm are a collection of coefficients.
Model. We consider transformer F R.C x RNNMLLLRR + NxC consisting of 7 transformer blocks. For the t-th block, we consider it employs function f(() NNC x NxM + NxC to update features, and function g( NC x PNNMXLXR + N MMXXXX to update positional embeddings:
x (1) = (x E(-")),
E (x E(T
t E [T].
We denote (T) as the final output of F(X (0) ( We assume f()) and (()) jointly satisfies our invariance properties Eqs. 4 and 5.
Phase Shift. We say transformer is invariant to translation the final outputs X (2) remains unchanged when the input token indices are shifted by an offset. This implies the attention mechanism inside depends on positional information based on the relative distances instead of absolute indices. Formally, when an offset A E IR is applied to all positions, the initial positional embeddings undergo phase shift. We denote the positional embeddings with translation A E whose per-token representations {e] (2)(0) [e]") can be written as:
[on (i A))
sin A)) cos(Om (i + A))
RoPE. 2(0)
for every E [N], [C/2].
sin((m (i A))
Random Fourier Features. For every E [N], m E [M], [4]
2 R
Wm.I.r cos (0, + A)) Wm,dr sin(m,, A)) __.
L.m
V
We denote the intermediate outputs yielded by shifted positional embeddings (X (1) (t)
E ). for every t [T].
Main Result. We provide formal version and proof of Proposition 4.2 below:
Proposition B.9 (Formal version of Proposition 4.2). Assume (!) and (()) satisfies Eqs. 4 and 5 for every E [T]. Then
for any shift A R, we have F(X (0) ((0)) F(X (0) E (0),
23

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Proof. First, observe that shift the token indices translates to orthogonal transformation the embedding space for both RoPE and random Fourier features. For RoPE, this can be seen by:
(0) i,m
= (0mm)
sin(0mi) cos(Omi)
(os(Omm)) 4) sin A) sin A) COs cos A)
sin 2m)
ORPPE.A.
where the extracted matrix ORoPE,,, is an orthogonal matrix. For random Fourier features, we observe that ēi.m (0) O OFFFA,, where O OFF.A. = diag (angOFFFFA...."" O ORFF.A,,,R22) and each O ORFFAA, E R2ax2 defined as: ii
(oos(0m, sin
- sin
ORFFAM,
cos(Om,r {
which shows the orthogonality of O OFFFA.
Now we apply the orthogonal invariance in an inductive argument. We make hypothesis that X (t) = x² and E = E o for some E [T] U {0} with O corresponding to the initialization specific orthogonal transformation. is obvious that this holds for & = 0. Then by the orthogonal invariance and equivariant of f((+1) and g(t+1), we have that for every e [T 1 {0]:
(t)
X ((+1) = (((+)
E0)) (()))))) p((+1) (x((', (x)) (141)
0)
E (t+1) g E EE = g ((+1) (x E 0) ((++1) (xl1) = E
which concludes the proof by induction.
C. Implementation Details
Experiment Settings. In Sec.5.1, the model architecture includes 16 layers, a hidden dimension of 1024, an intermediate dimension of 2048, and 16 attention heads, resulting in approximately 120M parameters. In Sec.5.2, the architecture features 12 layers, hidden dimension of 768, an intermediate dimension of 3072, and attention heads, totaling approximately 155M parameters. The training recipe in three experiments are presented in Tab.
Table 5. Training recipe for language model pre-training and fine-tuning in experiments.
Arithmetic ($5.1)
C4 Pre-training ($5.2)
SCROLLS ($5.2)
Context Extension ($5.3)
Sequence length Batch size Number of iterations
40 + 40
1024
1024
8096
512
512
64
64
20k
10k
1k
1k
Attention dropout prob.
0.0
0.0
0.0
0.0
Optimizer
Adam W
AdamW W
Adam W 1 x 10
Adam W 2 x 10 10-5
Learning rate
1 1 10
1 x 10
Masked Multi-Head Mechanism. The masked multi-head attention is key design in the original Transformer and is well compatible with our method. To enforce causality in language generation, the Transformer masks out (sets to -co) all values in the input to the softmax that correspond to illegal connections from future tokens to current tokens. This is similarly implemented in our enhanced Transformer for language modeling. To allow the model to jointly attend to information from different representation subspaces at different positions, multiple attention outputs are computed parallel with multiple attention heads, and then mixed through concatenation and linear transformation. In our enhanced Transformer, the head dimension is added to both token embeddings and positional embeddings, resulting in X E N and E E N XXXXXLXR where I/ denotes the number of heads.
Parameterization in Position Contextualization. The shapes of W1 and We allow for considerable flexibility. Given e; E x To achieve maximal expressiveness, e; can be flattened into R HMLXR with W, and W, E H Alternatively, to minimize parameter usage, we set W1 and We as IR xXxJ with weights shared across the M and dimensions. & is implemented through a standard MLP and the I dimension set to 4H in all experiments.
24

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Visualization of Tensor Operations. To provide clearer understanding of TAPE and the operation within the attention and feed-forward layers, visualize the process in Fig. 4.
Attention
Feed Forward
B = 3
L R 2
Linear
N B 3
= 4
Linear
Linear
=
R
1=4
Figure 4. Visualization of TAPE's operations. The channel dimension is omitted for simplicity all operations can be channel-wise. In the attention layer, the input token embeddings have a shape x B, and the positional embeddings have a shape of N x L x R. For the feed-forward layer, the N dimension is omitted as its operations are position-wise. The input token embeddings then have shape of (or B x 1), and the positional embeddings have shape x
D. Ad Additional Experiments
Ablation Study on Architecture. We ablate our architecture design for both attention layer and MLP layer in position contextualization. We conduct ablation studies on our architectural design for both the attention layer and the MLP layer in position contextualization. Additionally, we ablate two aspects of the design: rotation equivariance, by setting W1, W2 E IR Hxx1, which disrupts the (RR))-quivvaiance the use tensorial embeddings, by flattening L = = 2 into and R 4; and both properties simultaneously, by setting L 4 and 1. We use the same pre-training setting as Sec. 5.2 and directly report its perplexity in the test dataset of Github following He et al. (2024).
Table Ablation study TAPE architecture. We evalute pre-trained models' perplexity across varying sequence lengths the GitHub test set.
Architecture
Perplexity
Attention
Feed Forward
128
256
512
1024
X
X
139.2
92.8
69.3
57.2
X
143.3
95.0
70.7
58.4
142.7
94.3
70.1
57.6
132.0
86.6
63.9
52.2
Rotation Equivariance
Tensorial Embedding
X
x 140.7
92.1
68.2
56.2
X
138.4
91.3
67.8
55.7
X
132.9
87.8
65,4
54.1
132.0
86.6
63.9
52.2
As shown in Tab. 6, incorporating position contextualization in both the attention layer and the MLP layer results in the lowest perplexity across different positions within the training sequence length. Removing position contextualization from
25

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
either layer increases perplexity, even exceeding that of the traditional positional embedding without any architectural modifications. This outcome is reasonable, applying position contextualization to only one component introduces an architectural inconsistency. Furthermore, ablating rotation equivariance allows all neurons in the positional embedding to undergo linear transformations, increasing the number of parameters but leading to worse results compared to TAPE. Similarly, reducing the tensorial embedding to a vector embedding leads to higher perplexities and a decline in performance.
Ablation Study on TAPE Hyperparameter. We aim to investigate the impact of varying I on learning performance. Using the same pre-training settings as deseribed in Section 4.2, we directly report the perplexity on the GitHub test dataset. As shown in Tab. 7, there is no significant difference when using different values of although a trend of first decreasing and then increasing can be observed. This suggests that a range of I values from 2H- 24 to 3H 48 may yield better performance compared other settings. Therefore, general guideline, we recommend considering T {2,3, 4) H to optimize TAPE's performance.
Table 7. Ablation study on TAPE hyperparameter I. We evalute pre-trained models' perplexity across varying sequence lengths on the GitHub test set.
TAPE
Perplexity
Added Params. (M)
I
128
256
512
1024
0.11
12
133.2
87.9
65.2
53.6
0.22
24
133.0
86.1
63.2
51.8
0.44
48
132.0
86.6
63.9
52.2
0.88
96
133.2
87.5
64.5
52.7
1.76
192
133.0
87.3
64.5
53.0
Stability of TAPE under Positional Shifts. Stability in this context refers to the consistency of a sequence's 's representation under positional shifts (Sun et al., 2022). To evaluate the stability of TAPE, examine two types of positional shifts: (1) appending [BOS] tokens at the beginning of the sequence and (2) initializing positional indices with non-zero values to simulate a phase shift (Sinha et al., 2022). We analyze two aspects of the representation: the attention weights and the dot product of positional embeddings, quantifying their changes after applying positional shifts. For comparison, we include RoPE, which also exhibits O (())--quivariane. (R = 2) and remains consistent across layers, as well as TAPE without equivariance, as explored in previous ablations.
Table Comparison ROPE, TAPE, and TAPE without equivariance (wlo EQ) under positional shifts. The table shows differences attention weights (top) and positional embedding dot products (bottom) across layers for two shift methods: adding three [BOS] tokens ("Add Tokens") and starting position IDs at 3 ("Shift IDs").
Atten. Diff.
Add Tokens
Shift IDs
(x10-3)
Layer
Layer 2
Layer 4
Layer 8
Layer
Layer 2
Layer 4
Layer 8
RoPE TAPE
8.93
8.51
12.29
11.46
0.01
0.02
0.02
0.03
9.08
11.24
12.23
13.78
0.01
0.02
0.04
0.04
w/o EQ
11.30
11.38
13.32
14.55
0.01
0.24
0.37
0.51
PE Dot Prod. Diff. (%)
Add Tokens
Shift IDs
Layer
Layer 2
Layer 4
Layer 8
Layer
Layer
Layer 4
Layer 8
RoPE
0.03
0.03
0.03
0.03
0.03
0.03
0.03
0.03
TАРE
0.03
0.37
2.75
6.62
0.03
0.02
0.03
0.04
w/o EQ
0.03
2.29
3.34
6.37
0.03
0.54
0.44
0.86
As shown in Tab. 8, TAPE demonstrates stability comparable to RoPE, maintaining consistent attention weights and positional embedding dot products across different layers. Among these, the near-zero change (not exactly zero, attributable
26

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
numerical error observed in RoPE as well) in the dot-product when shifting IDs serves empirical evidence for Proposition 4.2. However, when equivariance is removed from TAPE, the differences increase significantly, especially in deeper layers, highlighting the importance of equivariance in preserving stability.
Additional Evaluation on Fine-tuned Llama-7b. Modern benchmarks provide a comprehensive means to assess large language models' advanced capabilities in language understanding and reasoning. Accordingly, we further evaluate our fine-tuned Llama-7b (See. 5.3) standard benchmarks, including ARC (Clark et al., 2018) and MMLU (Hendrycks et al., 2021).
Table 9. Accuracy in Percentage Across Methods and Benchmarks
MMLU (%)
ARC (%)
Method
Humanities
Social Sciences
STEM
Other
Challenge
Easy
LoRA
39.09 = 0.69
46,47 # 0.88
33.65 # 0.83
45.83 + 0.89
45.31 # 1.45
74.28 + 0.90
LongLoRA
37.53 0.69
43.55 = 0.88
3
43.84 + 0.88
4
74.16 + 0.90
ThetaScaling
37.45 + 0.69
43.16 0.88
33.05 # 0.83
44.64 0.88
45.65 1.46
74.24 0.90
TAPE
37.96 0.69
45,40 # 0.88
33.27 # 0.83
45.06 + 0.88
46.25 1.46
74.16 0.90
As Tab, 9 shows, TAPE demonstrates notable performance compared to other methods on MMLU and ARC benchmarks. While TAPE's accuracy on MMLU is slightly lower than that of LoRA, it consistently outperforms others. On the ARC benchmark, TAPE performs comparably to other methods on the "Fasy" subset but exhibits an advantage on the "Challenge" subset, underscoring its potential in complex reasoning tasks. Remarkably, these results are achieved using only fine-tuning, without pretraining TAPE, despite the presence of certain degree of architectural shift.
Additional Evaluation in Arithmetic Learning We also evaluate the effectiveness of TAPE in Sec. 5.1 using a different
training and testing length: 20/40 instead of 40/80. This setup is easier for the model to learn, with convergence achieved in less than half the steps. As shown in Fig. TAPE outperforms FIRE with a marginal improvement of 5%. However, this improvement is less pronounced compared the case with train/test length of 40/80, suggesting that TAPE may be more effective in tackling complex and challenging tasks than simpler ones.
RoPE
RandPE
FIRE
TAPE
100
80
10
60
20
40
30
20
40
0
10
20
30
40.
10.
20
30
40
10
20
30
40
10-
20
30
40
Figure 5. Accuracy on addition task trained with length test on context length. The average accuracy across the heatmap 26.12%.
26. 12%, 39.44% and 41.42% respectively for RoPE, RandPE, FIRE and TAPE.
Integration with Extrapolation Technique. Inspired by the demonstrated potential of NTK- methods (Peng et al., 2023b) enhance the length extrapolation ability of RoPE, we have explored integrating TAPE with such techniques when initialized as RoPE. Specifically, integration with TAPE to evaluate its performance in length extrapolation. The experiments were conducted under the same settings as described in Sec. 5. As shown in Fig. 6, the diagonal region exhibits darker colors, indicating higher accuracies. Quantitatively, YaRN effectively enhances the length extrapolation performance of TAPE with RoPE initialization, achieving a modest relative improvement 3.4%. However, it still struggles to generalize to unseen sequences with significantly longer digit lengths.
selected the most recent method, YaRN (Peng al., 2023b), and implemented its
27

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
FIRE
TАРE
TAPE + YarN
100
80
20
40
60
20
80
0
20
40
60
20
40
60
80
20
40
80
Figure Accuracy on addition task on 2x context length. The average accuracy 26.98% 32.82% 33.92% respectively FIRE, TAPE and TAPE YaRN.
Table 10. Comparing answers of different methods example questions in QuALITY.
Question A
Question B
Method
Answer
EM
Answer
EM
Ground
The secret service budget was small
Only the private quarters or the office restroom
Truth
TAPE
The secret service budget was small They were all they were waiting for
Only the private quarters X
xPos
Only tiny part of the right of the right leave foreverish Only handsome man
RandPE
Their human opinion was trusted by others who have trust the services of their people
X
X
RoPE
Their orless them together with their repories did not only they didn's never done was never done was never done... (repeating)
X
The/O only the full-College All of the full-College All of the full-College... (repeating)
ALiBi
Jimmy Carter is the president's de facto president
X
Jimmy Carter is the presi- dent's de facto president
E. Illustrations and Interpretation
Visualization of PE Patterns. To better understand the impact of TAPE, we analyze its attention and positional embedding (PF) dot-product patterns. Fig. 7 compares the patterns of TAPE and RoPE in the last layer, while Fig. 8 illustrates the evolution of TAPE's s dot-product patterns from shallow to deeper layers. The x-axis and y-axis correspond to the token positions of sampled input sequence.
As shown in Fig. 7, TAPE demonstrates emphasize token locality. In Fig. 8, TAPE behaves similarly to RoPE in the first layer but gradually reduces the dominance of diagonal patterns as the depth increases. This transition results in the formation of grid-like patterns, indicating that the model starts to focus on distant tokens in structured and periodic manner.
evenly distributed long-range attention patterns, whereas RoPE tends
Examples on QuALITY. To further validate TAPE's superior performance on the SCROLLS benchmark, we present two example questions from the QuALITY dataset within the SCROLLS benchmark. As shown in Tab. 10 and the detailed questions in Tab. 11, TAPE consistently generates either the correct answer or response similar to the correct answer, even if not exact match. In contrast, xPos and RandPE produce meaningful sentences that are unrelated the specific question. RoPE and ALiBi, however, generate incoherent outputs: RoPE tends repeat certain phrases, while ALiBi fails to recognize the presence of question, producing the same irrelevant answer regardless of the input.
28

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
RoPE
TAPE
Attention Difference
0013
2013
11
45
480
U5
(()) ) Dot-product patterns of positional embeddings TAPE and RoPE.
((b) ) Difference between TAPE and RoPE
Figure 7. Comparison of TAAE and RoPE methods in terms of positional embedding not-product patterns and their resulting attention differences. (a) TAPE demonstrates a systematic attention to surrounding tokens with relatively small dynamic ranges, whereas RoPE exhibits highly significant diagonal pattern with distinctively black regions. (b) TAPE effectively attends to Longer-range tokens, avoiding excessive attention to the self-token, in contrast to RoPE.
Layer
Layer
100
80
60
40
Figure Dot-product patterns of positional embeddings layers 1, 4, 8, and 12 (last) of TAPE.
Table 11: Example Questions in QuALITY
Qu. (ID: 20007 RZDMZJYW 2)
Qu. B (ID: 20007 RZDMZJYW 4)
What made easier previous presidents get away with adultery?
Where in the White House feasible president to meet a woman?
(A) Their staff did not know
(A) Only the East Wing
(B) They always tried to hide it well
(B) Only the private quarters
(C) The secret service budget small
(C) Only the oval office. bowling alley. East Wing
The reporters never found
(D) Only the private quarters the office restroom
Continued next page...
29

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
Article Content:
The logistics of presidential adultery.
The Washington Times could hardly contain excitement: "A former FBI agent assigned to White House describes in new book how President Clinton slips past his Secret Service detail in the dead of night, hides under blanket in the back of a dark-colored sedan, and trysts with woman. possibly celebrity, the JW Marriott Hotel in downtown Washington." For Clinton-haters Gary Aldrich's s tale sounded loo good be true. And So--s-Seccee--erviie agent's "source" turned out to be thirdhand rumor passed on by Clinton scandalmonger David Brock. Those who know about White House securttyCllintov staffers, the Secret Service, former aides to Presidents Reagan and desh--demlishhd Aldrich's claims. Clinton couldn't I give his Secret Service agents the slip (they shadow him when he walks around the White House), couldn't i arrange private visit without tipping off hotel staff, and couldn't 't re-enter the White House without getting nabbed. (Guards check all at the gate- gute-especially those that urrive 4 a.m.)
Even so, the image resonates. For some Americans, is an article faith: Bill Clinton cheated his wife when was governor. and cheats on her as president. But can he? Is it possible for the president of the United States to commit adultery and get away with it? Maybe, but it's tougher than you think.
Historically, presidential adultery is common. Warren Harding cavorted with Nan Britton and Carrie Phillips. Franklin Roosevelt "entertained" Lucy Rutherford
White House when Eleanor was away. America was none the wiser, even if White House reporters
were.
Those who know Clinton is cheating often point to the model of John F. Kennedy, who turned presidential hanky-panky into a science. Kennedy invited mistresses the White House for afternoon (and evening, and overnight) liaisons. Kennedy seduced women on the White House staff (including, it seems, Jackie's own press
Continued next page...
30

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
secretary). Kennedy made assignations outside the White House, then escaped his Secret Service detail by scaling walls and ducking out back doors, If Kennedy did it, so can Clinton.
Well, no. Though Clinton slavishly emulates JFK in every other way, he'd de'd be a fool to steal Kennedy's 's MO d'amour, Here's 's why: 1) Too many people would know. Kennedy hardly bothered to hide his conquests. According to Kennedy mistress (and mob moll) Judith Campbell's autobiography, those who knew about their affair included: Kennedy's 's personal aides and secretary (who pandered for him), White House drivers, White House gate guards, White House Secret Service agents, White House domestic staff, most of Campbell's friends, a lot of Kennedy's 's friends, and several Kennedy family members. Such broad circulation would be disastrous today because: 2) The press would report it. Kennedy conducted his affairs brazenly because he trusted reporters not to write about them. White House journalists knew about, at least strongly suspected, Kennedy's 's infidelity, but never published a story about it. Ask Gary Hart if reporters would exercise the same restraint today. Clinton must worry about this more than most presidents. Not only are newspapers and magazines willing to publish an adultery story ahout him, but many are pursuing it.
For the same reason, Clinton would find it difficult to hire a mistress. A lovely young secretary would set off' alarm hells in any reporter investigating presidential mishehavior. Says former Clinton aide, "There has been real tendency have good-looking women on the staff in order to protect him..
3) Clinton cannot avoid Secret Service protection. During the Kennedy era, the Secret Service employed fewer than 500 people and had an annual budget of about $4 million. Then came Lee Harvey Oswald. Squeaky Fromme, and John Hinckley. Now the Sccret Service payroll tops 4,500 (most of them agents), and the annual budget exceeds $500 million (up 300 percent just since 1980), At any given time, more than 100 agents guard the president the White House. Top aides from recent administrations are adamant: The Secret Service never lets the president escape its protection.
So what's 's randy president to do? Any modern presidential affair would need to meet stringent demands. Only a tiny number of trusted aides and Secret Service agents could know of it. They would need to maintain complete silence about it. And no reporters could catch wind of it. Such an affair is improbable, but-take heart, Clinton-haters-iit it's not impossible. Based on scuttlebutt and speculation from insiders the Clinton, Bush, Reagan, and Ford White Houses, here are the four likeliest scenarios for presidential adultery. 1) The White House Sneak. This a discreet variation the old Kennedy/Campbelll liaison. It's late night. The president's 's personal aides have gone home. The family is away. He is alone in the private quarters. The private quarters, a.k.a. "the residence." " occupy the second and third floors of the White House. Secret Service agents guard the residence's 's entrances on the first floor and ground floors, but the first family has privacy in the quarters themselves. Maids and butlers serve the family there, but the president and first lady ask them to leave when they want to be alone. The president dials a "friend" on his private line. (Most presidents placed all their calls through the White House operators, who kept record each one; the Clintons installed a direct-dial line in the private quarters.) The president invites the friend over for a cozy evening at the White House. After he hangs up with the friend, he phones the guard at the East Executive Avenue gate and tells him to admit a visitor. He also notifies the Secret Service agent and the usher on duty downstairs that they should send her up the residence.
A taxi drops the woman near the East gate. She identifies herself to the guard, who examines her ID, runs her name through computer (to check for outstanding warrants), and logs her in a database. A White House usher escorts her into the East Wing of the White House. They walk through the East Wing and pass the Secret Service guard post by the White House movie theater. The agent on duty waves them on. The usher takes her to the private elevator, where another Secret Service agent is posted. She takes the elevator the second floor. The president opens the door and welcomes her. Under no circumstances could she enter the living quarters without first encountering Secret Service agents.
Let us pause for a moment to demolish two of the splashier rumors about White House fornication. First, the residence is the only place in the White House where the president can have safe (i.e.. uninterrupted) sex. He can be intruded upon or observed everywhere excepp perhaps, the Oval Office bathroom. Unless the president is an exhibitionist a lunatic, liaisons in the Oval Office, bowling alley, or East Wing are unimaginable. Second, the much-touted tunnel between the White House and the Treasury Department all-but-useless the presidential adulterer. It is too well-guarded. The president could smuggle mistress through it, but it would attract far more attention from White House staff than a straightforward gate entry would.
Meanwhile, back in the private quarters, the president and friend get comfortable in one of the 14 bedrooms (ur, perhups, the billiard room). After pleasant 15 minutes (or two hours?), she says goodbye. Depending how long she stays, she may pass a differrnt shift of Secret Service agents as she departs. She exits the White House grounds, unescorted and unbothered, at the East gate. The Risks: A gate guard, an usher, and handful of Secret Service agents see her. All of them have a very good idea of why she was there. The White House maid who changes the sheets sees other suspicious evidence. And the woman's 's- real name is entered in a Secret Service computer. None of this endangers the president too much. The computer record of her visit is private, least for several decades after leaves office. No personal aides know about the visit. Unless they were staking out the Fast gate, no journalists do either. The Secret Service agents, the guard, the steward, and the maid owe their jobs to their discretion. Leaks get them tired.
That said, the current president has every reason not to trust his Secret Service detail. No one seriously compares Secret Service agents (who are pros) to Arkansas state troupers (who aren't (). But Clinton might not trust any Continued on next page...
31

---

Rethinking Addressing in Language Models via Contextualized Equivariant Positional Encoding
security guards after the beating he took from his Arkansas posse. Also, if other Secret Service agents are anything like Aldrich, they may dislike this president. One Secret Service leak the lamp-throwing
damaged Clinton. Agents could tattle again.
2) The "Off-the- Recood" Visit. Late at night, after his personal aides and the press have gone home, the president tells his Secret Service detail that he needs to take un "off-the-record" trip. He wants to leave the White House without his motorcade und without informing the press. He requests two agents and unobtrusive sedan. The Secret Service shift leader grumbles but accepts the conditions. Theoretically. the president could refuse all Secret Service protection, but would be far more trouble than it's worth. would have to inform the head of the Secret Service and the secretary of the Treasury.
The president and the two agents drive the unmarked car to woman friend's 's house. Ideally, she has covered garage. (An apartment building or hotel would raise considerably the risk of getting caught.) The agents guard the outside of the house while the president and his friend their thing. Then the agents chauffeur the president back the White House, re-entering through the Southwest or Southeast gate, away from the press station.
The Risks: Only two Secret Service agents and their immediate supervisor know about the visit. It is recorded in the Secret Service log, which is not made public during the administration's tenure. Gate guards may suspect something fishy when they the car. reporter passer-by could spy the president-evee even through tinted windows-as as the car enters and exits the White House. The friend's neighbors might spot him, or they might notice the agents lurking outside her house. neighbur might call the police to report the suspicious visitors. All in all, a risky, though not unthinkable. venture.
3) The Camp David Assignation. A bucolic, safer version of the White House Sneak. The president invites a group of friends and staffers includiin his paramour but not his wife to spend the weekend at Camp David. The girlfriend is assigned the cabin president's lodge. Late at night, after the Hearts game has ended and everyone has retired to their cabins, she strolls next door. There a Secret Service command post outside the cabin. The agents on duty (probably three of them) let her enter. A few hours later, she slips back to her own cabin.
the
The Risks: Only a few Secret Service agents know about the liaison. Even though the guest list is not public, all the Navy and Marine personnel Camp David, well the other guests, would know that the presidential entourage included attractive woman, but not the first lady. That would raise eyebrows
got back to the White House press room.
4) The Hotel Shuffle. The cleverest strategy, and the only one that cuts out the Secret Service The president is traveling without his family. The Secret Service secures an entire hotel floor, reserving elevators and guarding the entrance to the president's ' suite. The president's s. personal aide (a man in his late 20s) takes the room adjoining the president's. s. An internal door connects the two rooms, the aide can enter the president's s room without alerting the agents the hall. This standard practice. Late the evening, the aide escorts comely young woman back the hotel. The
Secret Service checks her, then waves her into aide's room. She emerges three hours later. slightly disheveled. She kisses the aide the hall as she leaves. Someone got lucky- hut who?
The Risks: The posted Secret Service agents might see through the charade. More awkwardly, the aide would be forced to play the seamy role of procurer. (He would probably do it. Kennedy's assistants performed this task dutifully.) In short, presidential adultery is just barely possible in 1996. But would be extremely inconvenient, extremely risky. and potentially disastrous. It seems, fict,, lot more trouble than it's worth. A president these days might be wiser to imitate Jimmy Carter, not Jack Kennedy, and only lust in his heart.
32


Reviews:
[{"review_id": "u9A5LTtbuN", "reviewer": "Reviewer_qpJK", "summary": "This paper proposes a new method for learnable positional encodings, where they are allowed to depend on context/content. The positional encodings, termed TAPE (“conTextualized equivariAnt Position Encoding”), can be added to pre-trained transformers, with only the TAPE-relevant parameters fine-tuned. TAPE is permutation and orthogonal equivariant and uses higher order tensors, and performs well in experiments spanning arithmetic reasoning, long-context retrieval, and language modeling.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly thank Reviewer qpJK for appreciating our contributions. We address the concerns as follows.\n\n> W1: The motivation of the technique is a bit confusing. The authors claim that relative positional encodings are crucial for “stability and generalization to varying sequence lengths” (L223-224), and use relative positional encodings in their formulation (equation 6) and prove that the transformer is invariant to token index shift (Prop 3.1), but then highlight an arithmetic task where absolute positions are necessary. \n\nWe apologize for the inaccuracy in our original statement (Line 316) regarding absolute positions that caused confusion. Please refer to our response to Q1 to see whether this clarification addresses the point.\n\n---\n>W2: Also, the use of tensors is not very well-motivated. Overall, although the experimental performance is seemingly very good, the design of TAPE seems a bit ad-hoc.\n\nWe acknowledge that certain implementation choices in TAPE, such as the tensorial embedding design, were empirically driven to optimize performance. As noted in Appendix D, our ablation studies validate the effectiveness of these architectural decisions.\n\n---\n>W3: The paper also cites geometric learning as inspiration — “This approach is inspired from the studies for geometric deep learning which processes graphs and point clouds by integrating token features with their geometric properties while preserving inherent physical symmetries” (L83) — but isn’t this property (positional encodings depending on relative distances) already satisfied by RoPE? Some additional ablations that could help clarify the design of TAPE include ablating the tensorial nature (reducing dimensions), ablating the dependence on context in the positional embeddings, etc.\n\nYes, RoPE also satisfies this property, which is precisely why we employ it as one of our instantiation methods. We do not claim this property as our novelty, but rather use it as a principle to motivate our design choices regarding contextualized position embeddings. As shown in Appendix D, we provide ablation studies that: (1) validate the design of equivariant and tensorial embeddings, (2) analyze the effects of Attention and MLP layers, and (3) investigate the impact of hyperparameter choices. Note that ablating the context-dependence in our positional embeddings (our core design) reduces to RoPE, which is consequently included as a baseline in our experiments.\n\n---\n> Q1: The arithmetic task is one where absolute positions are necessary (L316), but TAPE uses relative positions (Prop 3.1) — how then does TAPE do well? (Clarification of TAPE’s advantages and understanding of experimental results)\n\nWe appreciate this opportunity to clarify. There was an inaccuracy in the original L316 statement regarding absolute positions - we have now corrected this, as absolute positions are not necessarily required for arithmetic. The critical factor is learning the relative importance of different positions within the sequence. As detailed in L328 onwards, TAPE is able to learn these position-dependent importance relationships within the task context. A detailed explanation is also attached:\n\nIn arithmetic tasks, every digit has equal importance to the equation, regardless of its distance from the output. Traditional positional embeddings often assume a distance-decay effect, where words farther apart are less significant in the output. While this assumption is valid for most language tasks, it does not hold for arithmetic tasks. Positional contextualization enables dynamic reweighting of positional importance based on the task context, preserving effective distance decay for language tasks while addressing arithmetic contexts appropriately. This highlights TAPE’s potential advantages in arithmetic tasks.\n\n---\n> Q2: Is there an indexing problem with equation 7? Don’t the two sums cancel each other out? (Clarification for assessing paper’s accuracy)\n\nNo, the two sums in Equation 7 form a linear combination of vectors with weights summing to 1, making they cannot cancel out.\n\n---\n> Q3: Do the same hyperparameter choices (dimensions, etc) work across different tasks? (Affects my evaluation of advantages/disadvantages of the method)\n\nYes, we maintain consistent hyperparameters across all tasks (with L=R=2 in our main experiments). Additionally, Appendix D provides detailed ablation studies analyzing the impact of different hyperparameter choices.\n\n---\n*Response to Suggestions:*\n1. Prop 3.1:  The orthogonality of R is indeed crucial for Proposition 3.1 to hold, as non-orthogonal transformations would violate the invariance properties demonstrated in Appendix B. While generalizing the assumptions is an interesting direction, we believe this extension merits dedicated future research.\n2. Figure 2: The x- and y-axes represent the sequence lengths of the two operands respectively.\n3. Typos: Thank you for your careful review. We have fixed all of them.", "claims_and_evidence": "L051, “This rigidity [the fixed distance dependence / locality bias] limits the ability of positional encodings to model long-range dependencies and makes it challenging to attend to distant query-key pairs.” Is there a citation or experiment to support the claim that it is the positional encodings, in particular, that make long-range dependencies hard for LLMs?\n\nL686: “To the best of our knowledge, we are the first to introduce equivariance in language models, recognizing the symmetry in positional embeddings.” This is a strong claim. Is this not already done by RoPE, as noted in the paper?", "methods_and_evaluation_criteria": "Yes, although the highlighted benchmark datasets seem to be chosen as tasks where TAPE is expected to perform well (arithmetic and long context). However, other datasets are included in the appendix.", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "The authors explain the success of TAPE in terms of permutation and orthogonal equivariance, as well as the tenderized representations, but only ablate orthogonal equivariance — this is the only area I see for improvement, as the other experiments seem very thorough.", "relation_to_broader_scientific_literature": "The paper presents a novel positional encoding, building on RoPE and others. It draws on insights from papers like Ebrahimi et al 2024 and Sinha et al 2022.", "essential_references_not_discussed": "There are other references that deal explicitly with general group equivariance in positional encodings for geometric applications, but I don’t think they are essential.", "other_strengths_and_weaknesses": "Strengths: \n\nThe proposed positional encodings, TAPE, perform well compared to baselines on arithmetic and long-context tasks. They also admit an efficient implementation, and enable parameter-efficient fine-tuning that works better than LoRA and LongLoRA on passkey retrieval. Several additional experiments that I might have requested were already in the appendix, including ablations of orthogonal equivariance, and evaluation on other LLM tasks where long context is not necessarily the main challenge.\n\nWeaknesses:\n\nThe motivation of the technique is a bit confusing. The authors claim that relative positional encodings are crucial for “stability and generalization to varying sequence lengths” (L223-224), and use relative positional encodings in their formulation (equation 6) and prove that the transformer is invariant to token index shift (Prop 3.1), but then highlight an arithmetic task where absolute positions are necessary. Also, the use of tensors is not very well-motivated. Overall, although the experimental performance is seemingly very good, the design of TAPE seems a bit ad-hoc. The paper also cites geometric learning as inspiration — “This approach is inspired from the studies for geometric deep learning which processes graphs and point clouds by integrating token features with their geometric properties while preserving inherent physical symmetries” (L83) — but isn’t this property (positional encodings depending on relative distances) already satisfied by RoPE? Some additional ablations that could help clarify the design of TAPE include ablating the tensorial nature (reducing dimensions), ablating the dependence on context in the positional embeddings, etc.", "other_comments_or_suggestions": "I would recommend making Proposition 3.1 more general, by explicitly stating the necessary assumptions on E rather than immediately specializing to RoPE and random Fourier features. Also, would the result of Prop 3.1 not hold if f and g satisfied (4) and (5) for only R in the set of permutation matrices, rather than orthogonal matrices? This counterfactual would be good to include as part of the statement, if Prop 3.1 is supposed to motivate orthogonal invariance specifically. \n\nFigure 2: can the authors specify what the x and y-axis are in the figure caption? \n\nAs a minor comment, there were a lot of typos. Here are some of them:\n\nTypos:\nL99, “adiditionaly” \nL77, “superioroty”\nL150, “transformer” -> “transformers”\nL162, “encoding” -> “encodings”\nL237, “conTexturalized” \nL246, “O(r)-invariance” —> “O(r)-invariant”\nL285, “test” —> “tested”\nL297: “arthimetic” \nL663, “quardratic”\nL672: “focus” -> “focuses”\nL792, “Contextulization”", "questions_for_authors": "1. The arithmetic task is one where absolute positions are necessary (L316), but TAPE uses relative positions (Prop 3.1) — how then does TAPE do well? (Clarification of TAPE’s advantages and understanding of experimental results)\n2. Is there an indexing problem with equation 7? Don’t the two sums cancel each other out? (Clarification for assessing paper’s accuracy)\n3. Do the same hyperparameter choices (dimensions, etc) work across different tasks? (Affects my evaluation of advantages/disadvantages of the method)", "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": "bwEjGaKcqR", "reviewer": "Reviewer_G28F", "summary": "This paper introduces a new approach to processing language sequences using transformer blocks, where token features and positional embeddings are combined and contextualized. The authors extend traditional positional encoding by dividing it into multiple blocks, allowing for more flexible associations between tokens and their positions. They also ensure that their functions for token mixing and position contextualization are equivariant to permutations and orthogonal transformations, addressing limitations in existing models. The overall goal is to improve how transformer models process both token features and positional information in sequences.\n\n## update after rebuttal\n\nI think this is a good paper and keep my positive score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate Reviewer G28F's positive assessment of our contributions and strong endorsement for acceptance. The reviewer provided one suggestion, to which we respond below:\n\n> I do not see any major weaknesses. One suggestion is about the study of hyper-parameters. I think the authors choose B, L, R to align with RoPE for comparison. Since the proposed embeddings are learnable, it would be interesting to see how those hyper-parameters could affect the performance.\n\nThank you for your valuable suggestion. We have conducted hyperparameter exploration (detailed in the Appendix D), and our experiments demonstrate that the RoPE-like initialization yields the best performance among the configurations we tested.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes, I check the proof of Proposition 3.1.", "experimental_designs_or_analyses": "I check all experiments.", "relation_to_broader_scientific_literature": "It provides a new way to handle positions in Transformer, which can be wildly applied to diverse domains.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths\n- Instead of a specialized architecture, this paper modifies attention and MLP layers of Transformer, allowing easy integration into existing transformers.\n- Contextualized positional embeddings allow the model to dynamically adjust how positional information is interpreted depending on the surrounding tokens, improving its ability to capture more nuanced relationships between tokens.\n- The authors conduct thorough evaluations on various tasks, including passkey retrieval, arithmetic learning, and training from scratch with different context lengths. Additionally, they provided clear and insightful visualizations of attention maps, which further enhance the understanding of their model's behavior.\n\nWeaknesses\n- I do not see any major weaknesses. One suggestion is about the study of hyper-parameters. I think the authors choose B, L, R to align with RoPE for comparison. Since the proposed embeddings are learnable, it would be interesting to see how those hyper-parameters could affect the performance.", "other_comments_or_suggestions": "None.", "questions_for_authors": "None.", "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": "nPZiGZH5CP", "reviewer": "Reviewer_poAq", "summary": "This paper introduces TAPE ,a novel approach to enhancing position-based addressing in Transformers by dynamically adapting positional encodings across layers based on sequence context. TAPE ensures stability and robustness by enforcing permutation and orthogonal equivariance. Experimental results demonstrate that TAPE outperforms existing positional encoding techniques in language modeling, arithmetic reasoning, and long-context retrieval tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly thank Reviewer poAq for appreciating our contributions, providing valuable suggestions on improving the work, and supporting the acceptance of this work. We address the questions as follows.\n\n>W1: The authors provide the running time of attention layers as experimental results in Table 4. However, since the proposed method updates positional features layer-wise through interactions and joint training with token representations in every Transformer layer, what is the experimental runtime during training and inference for large-scale language models?\n\nThank you for your suggestion. We have included the experimental runtime for full model inference in the updated table below.\n| Method | TAPE | RoPE | FIRE | T5’s relative bias |\n| --- | --- | --- | --- | --- |\n| Samples Per Second | 58.6 | 71.8 | 33.4 | 46.8 |\n\n>W2: Since the proposed method updates positional features layer-wise, how does this affect gradient in the model? Could the authors provide an analysis or empirical results on gradient changes to better understand the impact of this method on training process?\n\nWe appreciate this insightful question. During our experiments, we monitored gradient norms and did not observe significant differences compared to RoPE. However, as this aspect falls outside our primary focus (motivation, methodology, and significance), we did not conduct an in-depth analysis of gradient behavior.\n\n>Q1: In Table 2 and Table 4, the authors present experimental results showing that the proposed method performs well on long-context tasks. I am curious whether this strong performance in long-context scenarios has theoretical support.\n\nTAPE's superior performance in long-context scenarios can be attributed to two key factors: First, its learnable nature enables dynamic adaptation to varying context lengths. Second, as formally established in Proposition 3.1, our relative position encoding scheme possesses inherent generalization capabilities to unseen sequence lengths.\nIn contrast, conventional positional encoding methods exhibit fundamental limitations in long-context settings, as they either rely on predetermined distance-decay patterns, or lack this crucial relativity property.", "claims_and_evidence": "Yes, the claims are.", "methods_and_evaluation_criteria": "Yes, the proposed methods do.", "theoretical_claims": "Yes, I checked the formulas in Section 3.", "experimental_designs_or_analyses": "Yes, I checked the experimental results tables and their corresponding analysis.", "relation_to_broader_scientific_literature": "The Positional Encoding proposed in this paper aims to unleash the power of position-based addressing, as existing methods have gradually weakened this capability[1,2].\n\n[1]Roformer: Enhanced transformer with rotary position embedding. Neurocomputing,2024\n\n[2]Sun, Yutao, et al. \"A length-extrapolatable transformer.\" 2022", "essential_references_not_discussed": "No, there aren’t.", "other_strengths_and_weaknesses": "Strengths：\n\n- The proposed method is novel and effective, with corresponding theoretical support.\n\n- The extensive experimental setup, especially in terms of performance on long-context tasks provides strong evidence of the effectiveness of the proposed method.\n\nWeaknesses：\n\n- The authors provide the running time of attention layers as experimental results in Table 4. However, since the proposed method updates positional features layer-wise through interactions and joint training with token representations in every Transformer layer, what is the experimental runtime during training and inference for large-scale language models?\n\n- Since the proposed method updates positional features layer-wise, how does this affect gradient in the model? Could the authors provide an analysis or empirical results on gradient changes to better understand the impact of this method on training process?", "other_comments_or_suggestions": "- Figure 3 could be represented using a different type of chart and would look more visually appealing if it occupies half a column.", "questions_for_authors": "- In Table 2 and Table 4, the authors present experimental results showing that the proposed method performs well on long-context tasks. I am curious whether this strong performance in long-context scenarios has theoretical support.", "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": "u4v3Yd8gXc", "reviewer": "Reviewer_udZG", "summary": "This paper proposes a new positional encoding method for LLMs, which could enhance the position-addressing ability of transformers. Permutation and orthogonal equivariance are also applied to enforce the positional encoding. This method demonstrates superior performance on various tasks, especially long-context tasks, such as passkey retrieval tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly thank Reviewer udZG for appreciating our contributions, providing valuable suggestions on improving the work, and supporting the acceptance of this work. We address the questions as follows.\n\n> W1: Existing positional encoding methods have introduced improvements to RoPE to better adapt it for long-sequence tasks, such as NTK-aware RoPE [4]. Could you provide comparative experimental results of this method on long-context tasks?\n\nThank you for your suggestion. We have implemented YaRN as an additional baseline in our SCROLLS benchmark experiments. The training is currently in progress, and we will update the results once they are available.\n\n> Q1: Computing a positional embedding at each layer increases the computational cost for larger-scale language models. Is the added time acceptable, and does the proposed method have scalability? Please provide a time complexity analysis.\n\nAs shown in Table 3, TAPE introduces <1% additional parameters and approximately 12% increased computational cost (FLOPs). Since scaling typically involves stacking these layers, the overhead remains proportional. We believe the modest increase in memory usage and ~10% additional compute time is a reasonable trade-off for the benefits provided.", "claims_and_evidence": "the claims are supported by experiments results and theoretical proof.", "methods_and_evaluation_criteria": "the authors propose context-aware positional encodings aiming to improve positional encoding and enhance the performance of LLMs during both pre-training and fine-tuning.", "theoretical_claims": "I checked the proof of the proposition, and I think no issues with it.", "experimental_designs_or_analyses": "I have checked the author’s experimental results and the baselines used for comparison. Regarding the need for more comparative methods on long-context tasks, I will elaborate on this specifically in the weakness.", "relation_to_broader_scientific_literature": "This paper contributes to the broader scientific literature by introducing a novel context-aware positional encoding method for LLMs, enhancing their position-addressing ability through permutation and orthogonal equivariance. It builds upon prior work on positional encoding techniques like RoPE[1], ALiBi[2].\n\n[1] Su, Jianlin, et al. \"Roformer: Enhanced transformer with rotary position embedding.\" Neurocomputing, 2024\n\n[2]Press, Ofir, Noah A. Smith, and Mike Lewis. \"Train short, test long: Attention with linear biases enables input length extrapolation.” 2021.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: The authors propose a dynamic, context-aware positional encoding, which can be applied during both the pre-training and fine-tuning stages. They provide extensive experimental results in the main paper and appendix, demonstrating the effectiveness of the method across various tasks. Additionally, the paper is well-structured, making it easy to understand.\n\nWeakness: Existing positional encoding methods have introduced improvements to RoPE to better adapt it for long-sequence tasks, such as NTK-aware RoPE [4]. Could you provide comparative experimental results of this method on long-context tasks?\n\n[4] Peng, Bowen, et al. \"Yarn: Efficient context window extension of large language models.\" 2023.", "other_comments_or_suggestions": "Several typos need to be corrected :\n1. “superioroty” should be “superiority” in Line 76.\n2. “fine-fining”should be “fine-tuning” in Line 399.", "questions_for_authors": "Computing a positional embedding at each layer increases the computational cost for larger-scale language models. Is the added time acceptable, and does the proposed method have scalability? Please provide a time complexity analysis.", "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:

Meta-review: The paper introduces TAPE, a method designed to enhance positional embeddings by incorporating sequence content across layers. TAPE can be integrated into pre-trained transformers, with only the relevant parameters fine-tuned. The proposed method is novel, and it demonstrates strong performance compared to baselines on arithmetic and long-context tasks. Extensive experiments are conducted, and the technology quality is solid. Overall, this is a strong paper. For the final version, the authors should clarify the motivation further and provide additional results as per the reviewers' comments.
Decision: Accept (poster)