Positional Attention: Expressivity and Learnability of Algorithmic Computation
Authors: Artur Back de Luca
George Giapitzakis
Shenghao Yang
Petar Veličković
Kimon Fountoulakis
Venue: ICML 2025 poster (2025)
Keywords: Theory
Transformers
Attention
Parallel Algorithms
Expressivity
Sample Complexity
Out-of-distribution Generalization
Massively Parallel Computation
Primary area: deep_learning->theory
TL;DR: This work shows that Transformers where the attention weights depend exclusively on positional encodings maintain expressivity and learnability in the representation and learning of parallel algorithms.

Abstract: There is a growing interest in the ability of neural networks to execute algorithmic tasks (e.g., arithmetic, summary statistics, and sorting).
The goal of this work is to better understand the role of attention in Transformers for algorithmic execution. Its importance for algorithmic execution has been studied theoretically and empirically using parallel computational models. Notably, many parallel algorithms communicate between processors solely using positional information. Inspired by this observation, we investigate how Transformers can execute algorithms using positional attention, where attention weights depend exclusively on positional encodings. We prove that Transformers with positional attention (positional Transformers) maintain the same expressivity of parallel computational models, incurring a logarithmic depth cost relative to the input length. We analyze their in-distribution learnability and explore how parameter norms in positional attention affect sample complexity. Our results show that positional Transformers introduce a learning trade-off: while they exhibit better theoretical dependence on parameter norms, certain tasks may require more layers, which can, in turn, increase sample complexity. Finally, we empirically explore the out-of-distribution performance of positional Transformers and find that they perform well in tasks where their underlying algorithmic solution relies on positional information.

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Artur Back de Luca"
George Giapitzakis" Shenghao Yang Petar Veličković? Kimon Fountoulakis
Abstract
algorithms (Veličković & Blundell, 2021). Notable appli- cations range from basic arithmetic to solving problems on abstract data structures such as graphs and lists (Lee et al., 2024; Tang et al., 2020; Tay et al., 2020; Yan et al., 2020; Veličković et al., 2020; Veličković et al., 2022; Dudzik & Veličković, 2022: Ibarz et al., 2022: Cappart et al., 2023;
There is growing interest in the ability of neural networks execute algorithmic tasks (e-g., arith- metic, summary statisties, and sorting). The goal of this work is to better understand the role of attention in Transformers for algorithmic execu- tion. Its importance for algorithmic execution has been studied theoretically and empirically using parallel computational models. Notably, many parallel algorithms communicate between proces- solely using positional information. Inspired by this observation, we investigate how Trans- formers can execute algorithms using positional attention, where attention weights depend exclu- sively on positional encodings. We prove that Transformers with positional attention (positional Transformers) maintain the same expressivity of parallel computational models, incurring a log- arithmic depth cost relative to the input length. We analyze their in-distribution learnability and explore how parameter norms in positional atten- tion affect sample complexity. Our results show that positional Transformers introduce a learning trade-off: while they exhibit better theoretical de- pendence parameter norms, certain tasks may require more layers, which can, in turn, increase sample complexity. Finally, we empirically ex- plore the out-of of-distribution performance of po- sitional Transformers and find that they perform well tasks where their underlying algorithmic solution relies on positional information.
Diao & Loynd, 2023; Bevilacqua et al., 2023; Rodionov Prokhorenkova, 2023; Minder et al., 2023; Georgiev et al., 2023; Engelmayer et al., 2023a; Georgiev et al., 2024; Li et al., 2024). Due to this success, there is growing interest in understanding neural networks as computational models that can solve complex algorithmic and computational tasks.
In this context, Pérez et al. (2021); Wei et al. (2022) show that Transformers are Turing Complete, and Giannou et al. (2023); Back De Luca & Fountoulakis (2024); Yang et al. (2024) illustrate how Transformers instructions in their parameters to solve linear algebra and graphs problems, Additionally, it has been shown that Trans- formers can perform computational tasks using far fewer layers than the number of sequential steps (Liu et al., 2023), indicating a connection between Transformers and parallel algorithms. Building on this, Sanford al. (2024) demon- strates that Transformers can simulate the Massively Parallel Computation (MPC) model (Andoni et al., 201 8), which is based on the MapReduce framework for large-scale data processing (Dean & Ghemawat, 2008).
effectively encode
While parallel computational models like MPC do not re- strict how communication is established, we observe that in many parallel algorithms, the communication between processors is independent of the processed data. Instead, this communication relies solely on processor identifica- tion (see Appendix A.3 for examples). Motivated by this observation, we study a Transformer variant that reflects this data-agnostic communication property. Specifically, we investigate positional attention, where the attention weights are determined exclusively by fixed positional encodings.
1. Introduction
Neural network models, particularly Transformers (Vaswani et al., 2017), have achieved remarkable success across vari- ous domains, especially in their ability to learn and execute
Our contributions. We examine Transformers with posi-
'Equal contribution 'David R. Cheriton School of Computer Science, University of Waterloo, Waterloo, Ontario, Canada. *Google DeepMind, London, UK.. Correspondence to: Artur Back de Luca <abackdel@ uwaterloo.ca>>.
Whiie positional encodings remain constant across layers, their attention weights can vary - unlike Graph Neural Networks (GNNs), which take a fixed connectivity structure as input. Fur- thermore. the tasks in this paper lack predefined graph structure, necessitating artificial graph for application to GNNs. For experimental comparisons, see Appendix D.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Positional Attention:
Expressivity and Learnability of Algorithmic Computation
tional attention (positional Transformers) from both theo- retical and empirical perspective in the context of algorith- mic computation. Our contributions are given below:
formers offers further insights into algorithmic learning in structured domains (Vinyals al., 2015; Prates et al., 2019; Joshi et al., 2019; Bai et al., 2020; Selsam & Bjørner, 2019; Karalias & Loukas, 2020; Gasse et al., 2019; Cappart et al., 2023; Veličković, 2023). Recent studies have also explored the broader algorithmic capabilities of Transformer- based models, shedding light on their potential for general algo- rithmic reasoning (Delétang et al., 2023; Butoi et al., 2025; Barbero et al., 2025). A parallel stream of research has focused on performing arithmetic operations (e-g, addition, multiplication) (Kaiser & Sutskever, 2016; Klindt, 2023; Shen et al., 2023; Lee et al., 2024; Ruoss et al., 2023). Specifically, Rodionov & Prokhorenkova (2023) leverages problem-spectie information within self-supervised train- ing framework, whereas the other studies use intermediate labels to guide the learning process. For example, Engel- mayer et al. (2023a) shows that using intermediate labels derived from parallel algorithms leads to improved perfor- mance. While some research, such as Kazemnejad et al. (2023), investigates the impact of different positional en- codings across several tasks, to the best of our knowledge, no existing work examines the use of positional attention within the context of neural algorithmic reasoning. The closest formulation to the positional Transformer architec- ture presented in Section 4 is the position-basec based attention of Schmidt& & Di Gangi (2023), but it is used in the context of neural machine translation.
Expressivity. We prove that positional Transformers can sim- ulate a single round of the Massively Parallel Computation (MPC) model with (logg)) layers, where n is the input length. The expressivity result emerges from a connection between Transformers and parallel computational models. In particular, we utilize a proxy computational model, which only allows communication over a static network, bridging the simulation results between positional Transformers and MPC and, consequently, by Sanford et al. (2024), standard Transformers. Due to the restrictive nature of communi- cation positional Transformers and static networks, this equivalence with MPC is established with logarithmic cost.
In-distributiin generalization. We provide a based generalization bound for positional Transformers, which, combined with our expressivity result, implies that they can learn algorithmic tasks to any desired accuracy via empirical risk minimization. Our analysis reveals a trade-off between the number of layers and parameter norms: positional Trans- formers exhibit better dependence on parameter norms but may require more layers for certain tasks, potentially lead- ing to higher overall theoretical complexity. This result is established using the Provably Approximately Correct (PAC) framework, where the complexity bounds of the hy- pothesis class are expressed as a function of the network parameter norms (Edelman et al., 2022).
Theoretical: From a theoretical perspective, the most closely related work to ours Sanford al. (2024), which presents simulation results for standard Transformers and the Massively Parallel Computation (MPC) model. Our work also presents simulation results demonstrating that positional Transformers can simulate MPC. The simulation proof involves introducing a proxy computational model with an alternative communication paradigm that can be shown to simulate MPC. Subsequently, we establish that the positional Transformer architecture can simulate this proxy model and, by extension, MPC itself. Other relevant stud- ies demonstrate the expressive power of neural networks through simulation results. For example, Siegelmann & Sontag (1995) establishes the Turing completeness of re- current neural networks (RNNs), while Hertrich & Skutella (2023) presents specific RNN constructions that solve the shortest paths problem and provide approximate solutions to the knapsack problem. Additionally, other simulation results focused on Transformers have shown their Turing completeness (Pérez et al., 2021; Wei et al., 2022) as well as demonstrated constructive solutions to linear algebra and graph-related problems (Giannou et al., 2023; Back De Luca & Fountoulakis, 2024; Yang et al., 2024).
Out-of-distribution (OOD) generalization. We empirically analyze the OOD performance of positional and standard Transformers. We observe that positional Transformers exhibit better OOD generalization for tasks where the under- lying algorithmic solution relies solely on positional infor- mation. However, for tasks that do not meet this criterion, such as induction heads tasks, both models perform poorly. We hypothesize that the inherently dynamic nature of com- munication in such tasks affects the loss landscape, making it challenging to discover good parameters for the positional Transformer during training.
2. Related work
Several studies focus on the experimental as- pects of training neural networks to execute algorithms. Notable examples include algorithms operating on common data structures, such as lists (e.g., sorting, binary search) and graphs (e.g., shortest paths, connected components) (Yan et al., 2020; Tang et al., 2020; Veličković et al., 2022; Ibarz et al., 2022; Bevilacqua et al., 2023; Engelmayer et al., 2023b; Georgiev et al., 2023; Minder et al., 2023; Diao & Loynd, 2023). Additional research on combinatorial optimization tasks using graph neural networks and Trans-
Empirical:
2

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Self-attention
Positional attention
x X P
Linear
Q
Linear
Mul
Softmax
A (9
P
Mul
Softmax
-0
x
Linear
K
Muil
Linear
Linear
K
Mui
Linear
-0"1
Linear
VO
(
Linear
V0
Figure 1: Diagram comparing the operations of self-attention in Transformers with positional attention. The figure illustrates single attention head, but in multi-head attention, multiple sets of queries, keys, and values are processed in parallel and then combined. In Transformers, the model's input, (O),, is combination of input values X and positional encodings P.
In positional attention, however, these components are processed separately, At layer f, the query (Q(0) and key (K)
are derived solely from the positional encodings P. where P remains fixed across layers. These are multiplied (denoted by Mul) and passed through a softmax function to produce the attention matrix A As in self-attention, the value (())
in positional attention is computed from the previous layer's input, X (t-1), The attention matrix (4()) and the value (()
then multiplied to form the weighted representation, which is linearly transformed into the output O This output passed to Multilayer Perceptron (MLP) for further processing, as detailed in Section 4.
3. Preliminaries and notation
ated with an attention weight matrix A E (0, 1) and value matrix w-(e.A) [ddxxdv xdv, Here, l denotes the layer index and h the head index, allowing a specific attention head within layer to be identified as (f, h). The outputs of these attention heads are concatenated and then transformed by an output matrix W. E jf xdo, This result is con- catenated with a residual connection of the input X and then passed through a multilayer perceptron (MLP), represented as (())_ RRdo+xx -> R Be
Let N- (1, denote natural numbers and [m] .. for n E N. For a set S, let P(S) be its power set (i.e. the set containing all subsets of S). For a matrix Z, let Zi,: and Zr, be the i-th row and column of Z. respectively. -12 denotes the spectral norm for ma- trices, and || llr. denotes the (p, q) matrix norm where the p-norm is over columns and q-norm over rows (e.g. || = Σ, l12). For vectors, || llp denotes the l norm. For a sequence of vectors X E Ra we may use x[]] = X E Ied for the i-th vector in the sequence.
Traditionally, in standard Transformers, the attention matrix ((f.h) is derived using self-attention.
Definition 3.1 (Risk). Let 'D be a distribution over X' x R,, and let R x
(8M) (X) = softmax (XIW) ((xI)(("))(((wW//))))) (2)
IR + I be a loss function. For a given class of functions F = (f= X + R and E F, the risk (i.e. the generalization error) is defined as risk(f: DD) = E( (x() v)), and the empirical risk is risk([; (zx(9,, (1(0,,(())) (2)))))))))))))))))))) tn Em ((I(x.()). 10))
Here, the attention weights are derived by the input ma- trix X, using query and key matrices w-le.)) w((.))
xdm, where dm is the embedding dimension.
In contrast, we utilize positional attention, where attention weights are learned solely using positional encodings P. which are constant across all layers. This distinction also illustrated in Figure
4. The Positional Transformer architecture
We now define the positional Transformer, an adaptation of the Transformer model (Vaswani al., 2017). This adapta- tion is motivated by the fact that communication between machines is independent of the data processed in many real- world parallel algorithms. In particular, we decouple the input values from the computation of attention weights. For an input X e 1_xxx, we define the 8h layer as follows:
A (E,N) = softmax
(pWWEM)) (PW((())) (3)
We utilize node positional encodings defined by a matrix P Bnxdp and whose attention weights are computed sim- ilarly to traditional attention, using query and key matrices but with input dimension dp. The positional encodings are fixed across layers, indicated by the absence of the index.
F(0) (X) (2)
(@@
o M=1
11
(2,h) xW((
Theoretically, we evaluate whether these changes reduce the expressive power of positional Transformers. We prove that
The input is processed by H attention heads, each associ-
3

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
positional Transformers, like standard Transformers, can simulate the Massively Parallel Computation (MPC) model. This indicates that there is no loss in expressivity, though with an added logarithmic increase in depth relative to the input length, as discussed further in the next section.
Remark 5.2. Consider a standard Transformer M' with N nodes, I layers, and dy H sublinear in N. Then, there exists positional Transformer M with O(N2) nodes, OL log N) layers, and dy- H sublinear in N that can sim- ulate the computation of M'..
Note that such simulation result introduces an additional quadratic dependency on the number of nodes in positional Transformers. These dependencies are merely a byproduct of the simulation strategy employed in Sanford et al. (2024), and other, more efficient, simulation strategies may exist. Additionally, this quadratic cost arises only when simulat- ing the exact same sequence of operations of the standard Transformer. In practice, the model may converge to param- eter configurations that do not correspond to the standard Transformer. When solving a given problem, positional Transformers may adopt a completely different configura- tion to achieve the solution without incurring this additional cost, as demonstrated by our experiments in Section 7.
5. Expressivity of Positional Transformers
A popular approach to illustrate the expressivity of ar- chitecture is by establishing an equivalence with a com- putational model (Sanford et al., 2024; Pérez. et al., 2021; Loukas, 2020). For positional Transformers, we will utilize the Massively Parallel Computation (MPC) model, which was first used by Sanford et al. (2024) to analyze standard Transformers. We begin by stating our main result.
Theorem 5.1. Consider an instance of MPC T with R rounds, N machines with local memory s. Let M be model following the architecture in Equation (1) with N +1 1 nodes, 2R[oggN N1 layers and 2s attention heads. Then, for any instance T with Borel measurable local func- tions, there exists configuration of M that approximates T to any desired degree of accuracy.
6. Learnability of Positional Transformers
In this section, we present a norm-based generalization bound for the positional Transformer architecture, discuss its implications, and compare it with the previously derived bound for the Transformer architecture that employs self- attention (Edelman et al., 2022). In particular, by decoupling the attention weights from the input X and requiring them to depend exclusively remove the exponential dependence (w.r.t. the depth 7.) on the norms of key and query weight matrices.
Proof overview: We establish this result using a two-step procedure. This procedure involves introducing a proxy computational model which, in contrast MPC, uses static network for communication. We call this model Par- allel Computation with Oracle Communication (PCOC)?, and its formal definition is presented in Appendix A. In the first step of the proof, we demonstrate that the computations over a static network model such as PCOC can effectively simulate the dynamic communication of MPC. Specifically, we show that R. rounds of an MPC protocol can be simulated using 2R Rlog n rounds of PCOC, where n represents the number of processors in MPC. The additional computation rounds and communication coordinated by Beneš networks (Beneš, 1965) are sufficient to simulate MPC.
the positional encodings P, we
We consider multi-layer networks obtained by composing individual layers from Equation (1), denoted as
(X) F(L) o .... F(2) (X).
(4)
The resulting network Fi:L has L layers and each layer has H attention heads. For simplicity we assume that the MLP ((t) consists of two layers:
In the second step of the proof, we show that PCOC can be, in turn, simulated by positional Transformers. Specifically, the attention mechanism is shown simulate communica- tion between nodes, while the computation stage leverages universal approximation results for MIPs (Cybenko, 1989; Hornik et al., 1989a). By demonstrating our architecture's 's capability to approximate both communication and local computation, we establish that it can simulate any MPC instance of R rounds using 2R log n layers.
(T0 o ((W W VE E [Z],
(5)
where a is an L.pLipschitz activation with (0) = O. Us- ing MLPs with more layers would lead to a worse depen- dence on the spectral norm of weight matrices, but does not affect the results in this section. Following prior work, our generalization bound is derived for scalar-valued func- tions. We extract a scalar output from our architecture as w T x [n] with trainable weights w E Rddu where Xout = (X) E gaxdous, The index i at which we apply the linear function I + w x can be arbitrary, and
By Theorem 3.4 in Sanford et al. (2024), which demon- strates the simulation of standard Transformers by MPC, we can derive the following remark regarding the simulation of standard Transformers by positional Transformers.
'More layers in the MLP would simply increase the constant e in Theorem 6.1 to e' On the other hand, two-layer MLPs are already sufficient invoke the universal approximation properties required by the expressivity result of Theorem 5.1.
We note that PCOC is only used to establish our expressivity result. It is not meant model compete with established parallel models, such as MPC.
4

---

Positional Attention:
Expressivity and Learnability of Algorithmic Computation
here we simply set to the last index. This setting captures algorithmic computations whose output is scalar. For comparison purposes, it aligns with the setup considered previously by Edelman et al. (2022) for the analysis of self- attention Transformers. In general, for computational tasks that require vector or matrix output, one can easily extend our result to parallel architectures using a union bound on the failure probability ó and incur an additional logarithmic cost with respect to the dimension of the output.
Compared with the norm-based risk bounds of self-attention Transformer networks in Edelman al. (2022), the bound in Theorem 6.1 does not depend on the spectral norm of key and query weight matrices. Intuitively, this is due to restricting attention weights to be determined exclusively by positional encodings, as opposed to allowing them to be determined compositionally by the output from the pre- vious layer. If we replace positional attention with self- attention in the definition of F, then the bound in The- orem 6.1 would incur an additional factor BOk where
Given , B2 2 1, the class of L-layer H-head bounded- norm scalar-outple positional Transformer networks is:
|w wheh) S Box. Vt,h. This shows a poten- tial ad advantage of positional attention when all other model configurations are the same.
=-{wi f- (FLL(X))))) : || & B2, and Ve, [I
Note that achieving low empirical risk with self-attention or positional attention can require very different model con- figurations, which may depend on the computational task. For example, in the simple task of computing the minimum of an n-number array, standard Transformer can solve it approximately using just one attention layer with an er- ror controlled by how well softmax approximates hardmax. In contrast, a positional Transformer can solve it exactly but requires log layers with two attention heads per layer (e.g., via binary tree reduction). Here, the positional Trans- former's worst-case sample complexity depends on n, while the standard Transformer does not. However, the standard Transformer can have a much worse dependence on the mag- nitude of the input numbers due to the approximation. Therefore, one architecture might be better depending on the magnitude of and the input numbers. For more complex tasks, setting an "optimal" model config- uration (e.g., choosing L, H, dv is difficult.
T w(thh) 12,1, IW (
2,,1, (1 (8) 12,11 S B2,11
[Iw]""la, (la, s B2,1, IW /
22, ||
12, 11
" 112 S B2
with pl:k from Equation (4), A ( from Equation (3), and ((() from Equation (5). In order to avoid overloading the result with unnecessary notation that do not add value to the discussion in any meaningful way, we assume that the input X E xxxx and positional encodings P E Rxxd normalized row-wise that ||X,,:l22 |P. = 1 for all i E (n). We denote d = max dr, dv dout] }. and therefore d(H H + 1) is the largest possible size of any weight matrix along any axis.
Let D be a distribution on Raxxx xR and (x(1) ((()y)))))
(X (X(m), ()) E Raxdx x R be sequence of m samples.
Theorem 6.1 (Generalization bound). For any 8 > 0, with probability at least 1 8, simultaneously for all f E F, it holds that for some c 0,
7. Experiments
risk(s: - risk(J: (x(1), y )m))))))
In this section, we evaluate the performance of positional and standard Transformers across various algorithmic tasks and regimes. We first outline the tasks considered in this work and then describe the experimental setup in detail. Next, we present results for the in-distribution regime across different settings, followed by an analysis of the out-of- distribution performance. Finally, we introduce a more complex task that combines textual and numerical data. For additional analyses, including comparisons with other mod- els and ablation studies, refer Appendix D.
s δ (HLLO)))) oL B2.1
log(H ddd)
11g(1/8) m
In the context of learning algorithmic computations, Theo- rem 6.1 means that the generalization gap ( f)l goes to 0 as the number of samples tends to infinity. On the other hand, by Theorem 5.1, for any computational task that is computable in an MPC model with Borel measur- able local functions, there is f E F (where the parameters L, H, B2, 12, that specify F might depend on n but are independent of m) such that the empirical risk can be arbi- trarily close to 0 for all samples of size m and for all m 2 O, i.e. risk(]: (x (0) >m Ri O, Vmn Consequently, in this case, Theorem 6.1 implies that as the number of samples m tends to infinity, we have risk(f:D) D) + O. This means that positional Transformer learns algorithmic computations through empirical risk minimization.
Tasks: We consider the following algorithmic tasks:
1. Cumulative sum: Given 2 E R", output y E IR" where each element y4 is the sum of the first i elements of x, i.e. y/s Σ)= Zj.
2. Cumulative min: Given x E R", output y E R where each element yi is the minimum value among the first i elements of x, i.e. y: = min(2, | S i))
5

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Sum
Min
Median
Sort.
Max sum salvarray
10
10
10
Loss (MSE)
10
10"5
10
10
10
10-1
10
1056
20k
3.16
4UA
00K
3.06
4UK
0.K
10k 2016
201
10%
2026
30k
400
20k
30k
40%
60K
Number of Samples
Positional
Stendlrd
Figure 2: In-distribution loss across all five tasks for standard Transformers (red) and positional Transformers (blue) as function of the number of training samples (indicated on the x-axis).
3. Cumulative median: Given x E R", output y E R" where each element yi the median of the first elements of æ, i.e. y:- - median
Experimental setting: All tasks employ the same model structure of Equation (1), augmented with linear encoding and decoding layers. We compare the standard Transformer, which utilizes the attention mechanism in Equation (2), and the positional Transformer, which employs the attention defined in Equation (3). Standard Transformers incorporate positional encodings concatenated with the input values, unlike positional Transformers, where positional informa- tion is exclusively supplied through the matrix P. Ap- pendix D, we examine other configurations for standard Transformers, including relative positional encodings such as Rotary Positional Embedding (RoPE) (Su et al., 2024b) and find significant differences in performance. Both variants share the same number of layers and dimensional configurations, with any specific differences explicitly noted. For details on the particular layer configurations and training details, refer Appendix D. The training data is sampled in the range [-2, 2]. To ensure diversity in the data, for each input sample,
j S i}.
4. Sorting: Given x RR, output sort(ar), vector contain- ing the entries of sorted in ascending order.
5. Cumulative maximum sum subarray: given x E IR", out- put E IR" where each element Yi is the sum of elements a maximum sum subarray of the first elements of x, i.c. yı sisis E
The tasks selected were chosen to represent varying levels of complexity and illustrate the strengths and limitations of each architecture. We adopt cumulative versions of al- gorithms when feasible for several reasons: They naturally provide n to training settings and are more challenging than non-cumulative versions.
To test the non-numeric capabilitics positional Trans- formers, we further present two additional tasks that utilize textual data. The first the k-hop inductive heads of San- ford al. (2024), which represents higher-order version of the standard inductive heads task by recursively executing the completion of a bigram to determine the subsequent bigram. The second task employs the same computational tasks described earlier in the list, but incorporates additional textual data that serve as categories, which the model must learn to for conditional reasoning.
first select lower and upper bounds n and
Ta uniformly in [-2, 2, 2). and then for each of the elements of the input sample, we select its value uniformly in [nn, Yu].
k-hop induction heads (k 3)
Loss (Crosa Entroopy))
0.6
1.5
0.4
1.0
Although both tasks involve pattern matching, the nature of pattern matching differs. The k-hop induction heads task re- quires dynamic pattern matching, where each step depends on previous matches, whereas the mixed-type input task in- volves static pattern matching, which better aligns with our architecture. Despite our expressivity result guaranteeing that this task can be solved using static network, we hypoth- esize that the inherently dynamic nature of communication in this task is reflected in the loss landscape. Consequently, discovering good parameters for the positional Transformer is difficult. This is evident in our experiments, where we run multiple trials, all of which result in poor dito--ddittibbttion performance.
0.5
0.2
Number of Semples
Number of Layers
Preitional
Stondand
Figure 3: In- in-distribution loss for standard Transformers (red) and positional Transformers (blue) in the k-hop induction heads task of Sanford et al. (2024), plotted as function of the number of training samples (left, with eight layers) and the number of layers (right, with 50,000 training samples).
6

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Sum
Min
Modian
Sort
Max suat subaray
10
10-
10
11P
10
10
Loss (MSE)
10
10
20
10
10
10
10
11f
10
10
10
10
10:
10
10
10
2
6
9
10
4
5
88
9
1
4
6 Seale
8
9
10
5
1
10
3
5
8
16
Positional
Standard
Figure 4: In-distribution and out-of-distributinn loss for standard Transformers (red) and positional Transformers (blue) across all five tasks. The x-axis shows the OOD scale factor, with scale factor of one indicating in nn-distribution performance.
7.1. In-distribution regime
J: Soop induction heads (k 3)
40
We evaluate our architecture in a fixed input length regime, with a detailed analysis of in in-distribution generalization as a function of the number of samples and expressive power. The plots presented show the median over ten runs, with shaded areas representing the 10th and 90th percentile.
85
30-
25
20
Sample Size vs. Generalization: In this setting, we demon- strate the learnability of positional Transformers, established in Theorem 1,1, as a function of the number of samples. To this end, we fix the input length n = 8 and examine in- distribution loss as a function of the number of training samples, ranging from 5,000 to 50,000, Figure 2 shows that for all tasks, the in-distribution loss of both architectures consistently decreases with increasing number of samples.
Number Layens
Figure 5: Out-of-distribution loss for standard Transformers (red) and positional Transformers (blue) in the k-hop induc- tion heads task of Sanford al. (2024), plotted as function of the number of layers (with 50,000 training samples).
Expressive power vs. Generalization: In this experiment, we empirically validate the theoretical results of Remark 5.2, analyzing the in-distribution capabilities of standard and po- sitional Transformers as a function of their expressive power, measured by the number of layers. We evaluate their per- formance generalized version of the induction heads task (Olsson et al., 2022), which reflects the relational capa- bilities of self-attention. This task, named k-hop induction heads, applies a successive l number of bigram completions to determine the next bigram in the sequence. We adapt the setting of Sanford et al. (2024), restricting the number of hops to most three and the sequence length to 30. For complete description of the task and additional experi- mental details, we refer readers to Sanford et al. (2024) and Appendix D, respectively. We chose this task because it is designed to highlight the efficiency of self-attention in terms of the number of required layers. The in-distribution results in Figure 3 show that standard Transformers exhibit near-zero in-distribution loss with as few as three layers. In contrast, we find that positional Transformers generally perform worse. As established in Section 7, this results from the expressivity of the positional Transformers since achieving level of performance comparable to that of the standard Transformers requires an increasing number of layers, aligning with our theoretical results.
7.2. Out-of-distribuute regime
We further evaluate the capabilities of both architectures in generalizing out-of distiiiribtttion (OOD) for the tasks previ- ously outlined. To generate OOD data, we employ a similar sampling strategy as in training but extend the value range to [-2c, 2c]. where >> 1 is the OOD scale factor. Addition- ally, during the test sampling process, we apply rejection step to ensure that either < 2 or Yu > 2, while main- taining -2c < 7 S The S 2c. This ensures that a test sample will not lie in the training domain with high prob- ability.* This implies that our test results reflect the "true" duttrf-distribution performance of both architectures.
The results in Figure 4 show that positional Transform- ers achieve significantly more stable performance across increasing
ranges than standard Trans- formers across computational tasks. We hypothesize that such differences in performance can be attributed to algo- rithmic alignment (Xu al., 2020), a concept that suggests
"This sampling strategy does not guarantee that every test sam- ple falls outside the training distribution. However, as shown in Ap- pendix E, the probability of test sample being within the training distribution is at most O(1 /nn²). In practice, most test instances outside this domain. For further details, see Appendix E.
7

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Min
Sum
Multitask
10
10
Loss (MSE)
10
10a
10
10
10-
10
10
109
10
10
10-1
10-1
10
10-0
2
Scale
Positional
Standard
Figure 6: In-distribution and out-of-distribution losses for standard (red) and positional Transformers (blue) across mixed- type input task variations. The x-axis shows the OOD scale factor, where one indicates in-distribution performance.
a particular architecture demonstrates better generalization capabilities if its underlying computational structure aligns more closely with the target task. Specifically, the struc- ture of these tasks aligns well with the core motivation for positional attention, i.e., enabling communication through positional information alone. For example, the minimum task is known to have solution where the underlying com- putational graph is binary tree and independent of the input values, see Appendix A.3. In Appendix F, we discuss potential reasons why standard Transformers fail to match the performance of positional Transformers and present ad- ditional experiments that illustrate the differing behaviors of self-attention and positional attention.
33.5, 'Cat28', 9.17, 'Cat30', 55.90, 'Cat31', Cat= 23.70, 'att33, 12.47, 'Cat34', 8.45, cat 40', "cat40' 1.50, 'Find min of Cat28, Cat Cat33 and Cateo" J. Output 1.50
The imput prompt presents the model with number of tex- tual categories and associated values and asks for the result of an aggregation operation on subset of the categories. Noise is injected by adding irrelevant categories (with no associated value) in the prompt (colored blue in the exam- ples). Note that the category identifiers, the range of values, and the type of noise all change between training and test data. Specifically, the train category identifiers integers between 1 and 20, whereas the test category identifiers are integers between 21 and 40. The models are trained on category values in [0, 5] and tested on category values in [0, 5c], for c = 1, 10. We experiment with the fol- lowing aggregation operations: minimum (min), sum, and multi-tasking setting that combines minimum (min) and maximum (max). This task involves distinguishing relevant from irrelevant categories and pattern matching across dif- fering train and test distributions. The task also demands basic conditional reasoning since each query involves only a subset of all given categories, and algorithmic compu- tation to derive the correct outputs. The textual input tokenized and processed through an embedding layer, while the numeric input passes through linear layer. The results, shown in Figure 6, indicate that the positional Transformer outperforms the Transformer."
For the k-hop induction heads task, OOD data is generated by creating sequences sampled exclusively from a portion of the alphabet not used during training. Figure illustrates the OOD performance of both standard and positional Trans- formers. Note that the positional Transformer performs poorly. We hypothesize that this is because the k-hop in- duction heads task deviates from core motivation, its inherent task requirement suggests need for dynamic communication that changes with every input, see also our comments in the description of the tasks above.
7.3. Additional experiments with mixed-type inputs
In this section, we further investigate the OOD performance when the input data is a mixture of tokens and numbers. This task involves conditional reasoning and simple pattern matching on top of computing an aggregation function. We begin by illustrating training and test sample, followed by high-level overview of the task. For the definition of the task and the details about the setting, see Appendix D.3.
8. Limitations and future work
Our work provides a theoretical and experimental inves- tigation of learning algorithms with positional attention. However, further research is needed in some directions that
Training sample: Input - C'Cat'', 3.45, Cat+77, "att5' 1.23, 'Cat'', 0.65, 'Cat8', 2.23, "Catll', 'Cat-8', 4.10, 'Cat13', 1.10, 'Cat14', 0.10, 'Cat20', 2.75, 'Find min of Cat5, Cat8, Catll and Cat20'), Output 1.23
'In Appendix D.3 compare the performance of positional Transformer against fine-tuned version of GPT2 (Radford al., 2019). Positional Transformer has better OOD performance. Note that this task does not demonstratc/snggeet that positional Trans- former will outperform other Transformer-basee architectures (like LLMs) general- purrose reasoning text data.
Test sample:
Input- C'Cat23', 7.28, Cat22''
8

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
believe are important. Theoretically, more precise OOD analyses are necessary capture finer differences between architectures, existing OOD generalization bounds are not tight enough (see Appendix G for an extended discus- sion). Empirically, our approach relies on fixed positional encodings, limiting its applicability to longer inputs. De- veloping positional encodings that support arbitrary input lengths would enable us to explore the length generalization capabilities of positional attention.
Bai, Y., Ding, H., Gu, K., Sun, Y., and Wang. W. Learning- based efficient graph similarity computation via multi- scale convolutional set matching. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pP. 3219-3226, 2020.
Barbero, F., Vitvitskyi, A., Perivolaropoulos, C., Pascanu, R., and Veličkovié, P. Round and round we go! what makes rotary positional encodings useful? In 13th Inter- national Conference Learning Representations, 2025.
Acknowledgements
Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities risk bounds and structural results. Journal of Machine Learning Research (JMLR), 3:463-483 2003.
The authors would like to thank Federico Barbero (Google DeepMind) and Csaba Szepesvári (Google DeepMind) for their valuable comments and suggestions on this work.
Benes, V. Mathematical theory of connecting networks and telephone traffic. Academic press, 1965.
K. Fountoulakis would like to acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC). Cette recherche a été financée par le Conseil de recherches en sciences naturelles et en génie du Canada (CRSNG), [RGPIN-2019-04067, 2GECR-2019- 00147].
Bevilacqua, B., Nikiforou, K., Ibarz, B., Bica, I., Paganini, M., Blundell, C., Mitrovic, J., and Veličković, P. Neu- ral algorithmic reasoning with causal regularisation. In Proceedings of the 40th International Conference on Ma- chine Learning, 2023.
S. Yang would like acknowledge the support of the Natu- ral Sciences and Engineering Research Council of Canada (NSERC).
Bounsi, W., Ibarz, B., Dudzik, A., Hamrick, J. B., Mar- keeva, L.., Vitvitskyi, A., Pascanu, R., and Veličković, P. Transformers meet neural algorithmic reasoners. arXiv preprint arXiv:2406. 09308, 2024.
G. Giapitzakis would like to acknowledge the support of the Onassis Foundation- Scholarship ID: ZU 020-1/2024- 2025.
Butoi, A., Khalighinejad, G., Svete, A., Valvoda, J., Cot- terell, R., and DuSell, B. Training neural networks as recognizers of formal languages. In 13th International Conference on Learning Representations, 2025.
Impact Statement
Our theoretical and empirical results demonstrate the po- tential of neural networks in addressing complex reasoning tasks, which are traditionally deemed exclusive to classical computational methods. This research enhances our under- standing of neural network capabilities. While this work is primarily theoretical and relies solely on synthetic data, we must acknowledge that societal implications may emerge from findings. However, at this stage, we do not identify any specific consequences worth highlighting here.
Cappart, Q., Chételat, D., Khalil, E. B., Lodi, A., Morris, C., and Veličković, P. Combinatorial optimization and rea- soning with graph neural networks. Journal of Machine Learning Research, 24(130):1-61, 11,30)):-61,,0000 2023.
Cybenko, G. Approximation by superpositions of a sig- moidal function. Mathematics of control. signals and systems, 1989.
Dean, J. and Ghemawat, S. Mapreduce: simplified data pro- cessing on large clusters. Commun. ACM, 51(1):107-113, 2008.
References
Andoni, A., Song, 7.., Stein, C., Wang, 7., and Zhong, P. Parallel graph connectivity in log diameter rounds. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science (FOCS), pp. 674-685, 2018.
Delétang, G., Ruass, A., Grau-Moya, J., Genewein, T., Wen- liang, L.. K., Catt, E.,, Cundy, C., Hutter, M., Legg, S., Veness, J., and Ortega, A. Neural networks and the chomsky hierarchy. In 11th International Conference on Learning Representations, 2023.
Back De Luca, A. and Fountoulakis, K. Simulation of graph algorithms with looped transformers. In Proceedings of the 41st International Conference on Machine Learn- ing, volume 235 of Proceedings of Machine Learning Research, pp. 2319-2366 2363. PMLR, 2024.
Diao, C. and Loynd, R. Relational attention: Generalizing transformers for graph-structured tasks. In The Eleventh International Conference on Learning Representations, 2023.
9

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Dudzik, A. J. and Veličković, P. Graph neural networks are dynamic programmers. Advances in Neural Information Processing Systems, 35:20635-20647, 2022.
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, 2022.
Edelman, B. L.., Goel, S., Kakade, S., and Zhang, C. Induc- tive biases and variable creation in self-attention mecha- nisms. In International Conference on Machine Learning, pp. 5793-5831. PMLR, 2022.
Im, S., Kumar, R., Lattanzi, S., Moseley, B., Vassilvitskii, S., et al. Massively parallel computation: Algorithms and applications. Foundations and Trends in Optimization, 5 (4):340-417,22 2023
Engelmayer, V., Georgiev, D., and Veličković, P. Parallel algorithms align with neural execution. In The Second Learning on Graphs Conference, 2023a.
Joshi, C. K., Laurent, T., and Bresson, X. An efficient graph convolutional network technique for the travelling salesman problem. arXiv preprint arXiv:1906.0m (eo::/90002222 2019.
Engelmayer, V., Georgiev, D. G., and Veličković, P. Parallel algorithms align with neural execution. In The Second Learning on Graphs Conference, 2023b. URL https: //openreview
Kaiser, L. and Sutskever, I. Neural GPUs learn algorithms. In The Fourth International Conference Learning Rep- resentations, 2016.
Gasse, M., Chételat, D., Ferroni, N., Charlin, L.., and Lodi, A. Exact combinatorial optimization with graph convolu- tional neural networks. Advances in neural information processing systems, 32, 2019.
Karalias, and Loukas, A. Erdos goes neural: an unsuper- vised learning framework for combinatorial optimization graphs. Advances in Neural Information Processing Systems, 33:6659-666 6672, 2020.
Georgiev, D. G., Lio, P., Bachurski, J., Chen, J., Shi, T., and Giusti, L. Beyond erdos-renyi: Generalization in algorithmic reasoning on graphs. In The Second Learning on Graphs Conference, 2023.
Kazemnejad, A., Padhi, I., Natesan, K., Das, P., and Reddy, S. The impact of positional encoding on length general- ization in transformers. In T'hirty-seennth Conference on Neural Information Processing Systems, 2023.
Georgiev, D. G., Wilson, J., Buffelli, D., and Lio, P. Deep equilibrium algorithmic reasoning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL httpps:///pppeeeee http::////penreeee.neeee SrrumiiddSSuIxxxCEN
Klindt, D. A. Controlling neural network smoothness for neural algorithmic reasoning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Lee, N., Sreenivasan, K., Lee, J. D., Lee, K., and Papail- iopoulos, D. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024.
Giannou, A., Rajput, S., Sohn, J.-Y., Lee, K., Lee, J. D., and Papailiopoulos, D. Looped transformers as pro- grammable computers. In Proceedings of the 40th In- ternational Conference on Machine Learning, volume 202, pp. 11498-11442 2023.
Li, H., Peng, C., Xu, C., and Yang. Z. Open-book neu- ral algorithmic reasoning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL httpt::/t///seeet hwtp:///epeeeeeeeeeeeeeee. wwtp/wneeeeeeee forum? i id=6H033urpaT.
Habermann, N. Parallel neighbor-sort (or the glory of the induction principle). Carnegie Mellon University, 1972.
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Hertrich, C. and Skutella, M. Provably good solutions to the knapsack problem via neural networks of bounded size. INFORMS journal computing, 1505):107-11097, 2023.
Loukas, A. What graph neural networks cannot learn: depth vs width. International Conference on Learning Rep- resentations, 2020.
Horiik K., Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. Neu- ral Networks, 2(5):359-33))) 366,, 1989a. ISSN 0893-6080.
Mansour, Y., Mohri, M., and Rostamizadeh, A. Domain adaptation: Learning bounds and algorithms. In Proceed- ings of The 22nd Annual Conference on Learning Theory (COLT 2009), Montréal, Canada, 2009.
Hornik, K. Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. Neu- ral networks, 2(5):359-36) -3669936666163 1989b.
10

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Minder, J., Grötschla, E., Mathys, J., and Wattenhofer, R. SALSA-CLRS: A sparse and scalable benchmark for al- gorithmic reasoning. In The Second Learning on Graphs Conference, 2023.
Selsam, D. and Bjørner, N. Guiding high-performance sat solvers with unsat-core predictions. In Theory and Applications of Satisfiability Testing- SAT 2019: 22nd International Conference, SAT 2019, Lisbon, Portugal, July 9-12, 2019, Proceedings 22, pp. 336-353. Springer, 2019.
Nanda, N., Chan, I.., Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic inter- pretability. In The Eleventh International Conference on Learning Representations, 2023.
Shen, R., Bubeck, S., Eldan, R., Lee, Y. T., Li, Y., and Zhang, Y. Positional description matters for transformers arithmetic. arXiv preprint arXiv:231 14737, 2023.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., al. In-context learning and induction heads. arXiv preprint arXiv:2209. /00.099///99/ /895, 2022.
Siegelmann, H. and Sontag, E. On the computational power of neural nets. Journal of Computer and System Sciences, 50: 10:132-150, 1995.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomput. 568(CC 2024a. ISSN 0925- 2312.
Prates, M., Avelar, H., Lemos, H., Lamb, L. C., and Vardi, M. Y. Learning solve up-complete problems: A graph neural network for decision tsp. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 4731-4738, 2019.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing. 568:127063, 2024b.
Pérez. J., Barceló, P., and Marinkovic, J. Attention is turing- complete. Journal of Machine Learning Research, 22 (75):1-35, 2021.
Tang, H., Huang, Z., Gu, J., Lu, B.-L., and Su, H. Towards scale-invariant graph-related problem solving by iterative homogeneous GNNs, Advances in Neural Information Processing Systems, 33: 13115811-5882, 2020.
Radford, A.,
Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Tay, Y., Bahri, D., Yang, L.., Metzler, D., and Juan, D.-C. Sparse sinkhorn attention. In International Conference on Machine Learning, pp. 9438 1108,944, 2020.
Redko, I., Morvant, E., Habrard, A., Sebban, M., and Ben- nani, Y. A survey on domain adaptation theory: learning bounds and theoretical guarantees, 2022.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit. J., Jones, I.., Gomez, A. N., Kaiser, t.., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Rodionov, G. and Prokhorenkova, L. Neural algorithmic reasoning without intermediate supervision. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Veličković, P. and Blundell, C. Neural algorithmic reasoning. Patterns, 2(7), 2021.
Ruoss, A., Delétang, G., Genewein, T., Grau-Moya, J., Csordás, R., Bennani, M., Legg, S., and Veness, J. Ran- domized positional encodings boost length generaliza- tion of transformers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), 2023.
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, 2022.
Sanford, C., Hsu, D., and Telgarsky, M. Transformers, par- allel computation, and logarithmic depth. In Proceedings of the 41st International Conference Machine Learn- ing, volume 235 of Proceedings of Machine Learning Research, pp. 03276-43322 PMLR, 2024.
Veličković, P. Everything is connected: Graph neu- ral networks. Current Opinion in Structural Biology, 79:102538, 2023. ISSN 0959-440X. doi: URL science/arte httaaa///..ttceeeee///////94////
ttttps://wwwsccienctttt
com/
Schmidt, F. and Di Gangi, M. Bridging the gap between Dosition-base and content-based self-attentio for neural machine translation. In Proceedings of the Eighth Con- ference on Machine Translation, pp. 507-521, Singapore, 2023. Association for Computational Linguistics.
Veličković, P., Ying, R., Padovano, M., Hadsell, R., and Blundell, C. Neural execution of graph algorithms. In International Conference on Learning Representations, 2020.

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Vinyals, O., Fortunato, M., and Jaitly, N. Pointer networks. In Cortes, C., Lawrence, N., Lee, D., Sugiyama, M., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 28, Curran Associates, Inc., 2015. URL.
neurups cc pecpapeee_iillessppprr//////////
299200022200a3aaee3//12222288223 Paper. pdf.
Wei, C., Chen, Y., and Ma, T. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Infor- mation Processing Systems, 35:12071-1200 2022.
Xu, K., Li, J., Zhang, M., Du, S. S., ichi Kawarabayashi, K., and Jegelka, What can neural networks reason about? In International Conference on Learning Representations, 2020.
Yan, Y., Swersky, K., Koutra, D., Ranganathan, P., and Hashemi, M. Neural execution engines: Learning to execute subroutines. Advances in Neural Information Processing Systems, 13:17298-17341 2020.
Yang. L.. Lee, K., Nowak, R. D., and Papailiopoulos, D. Looped transformers are better at learning learning al- gorithms. In The Twelfth International Conference on Learning Representations, 2024.
12

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Appendix
Table of contents
A Background and definitions
A.I Massively Parallel Computation (MPC)
A.2 Parallel Computation with Oracle Communication (PCOC)
A.3 Illustration of parallel algorithms
A.4 Additional definitions
B Expressivity results
B.1 Proof of Theorem 5.1 (Positional Transformers can simulate MPC)
B.2 Proof of Lemma B.1 (PCOC simulate MPC)
B.3 Hardmax patterns using positional attention
B.4 Proof of Lemma B.2 (Positional Transformers simulate PCOC)
B.5 Softmax patterns using positional attention
C Proof of Generalization Bound (Theorem 6.1)
C.I Preliminaries
C.2 Lipschitz Property of the Architecture
C.3 -Covering Number Upper Bound
D Experiments
D.1 Numeric tasks
D.2 k-hop induction heads
D.3 Mixed-type input task
E - Probability of generating OOD test data in empirical setting
F Potential reasons for failure in self-attention for out-of-distribution
G Informal Discussion on Out-of-distribution Generalization Bounds and future work
A. Background and definitions
A.1. Massively Parallel Computation (MPC)
The MPC model is computational model of the MapReduce framework (Dean & Ghemawat, 2008), widely used for computations over massive datasets. It defines parallel computing model that jointly executes function across multiple machines, each constrained by limited memory capacity. The MPC model is capable of representing various parallel algorithms (Im et al., 2023) and is more powerful than other established parallel models, such as parallel random-access machine (PRAM) (Andoni et al., 2018).
For completeness, we provide a simplified version of the definition of the MPC protocol by (Andoni al., 2018), which makes the connection
proxy computational model more apparent
Definition A.1 (MPC protocol, Def. I.1 (Andoni al., 2018), simplified). Let be parameter. There are p 2 machines (processors), each with local memory of size The input is distributed on the local memory of some of the machines. The computation proceeds in rounds. In each round, each machine computes the data in its local memory and sends messages other machines at the end of the round. The total size of messages sent received by machine in round is bounded by In the next round, each machine only holds the received messages in its local memory. At the end of the computation, the output in the memory of some machines.
'The original definition includes assumptions about the relation between the input size, the memory size, and the number of machines which we omit hetter highlight contrast between MPC and PCOC.
13

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
A.2. Parallel Computation with Oracle Communication (PCOC)
In this section, we present the definition of PCOC used in Section 5 and Appendix B establish the expressivity of positional Transformers. We emphasize that PCOC is merely a proxy model to establish our expressivity results. It is not intended to serve competing parallel computational model, such MPC.
The PCOC model consists of two steps at each round. The first step is communication, where machines send and receive data from other machines. The communication pattern can change at every round. The second step is computation, where all machines perform some local computation on data stored in their local memory.
Oracle communication. For each length n and input data, we assume the existence of an oracle that provides the destination for each machine and message each round. The oracle executes communication pattern agnostic the data being processed. This contrasts with other parallel models, such as the Massively Parallel Computation (MPC) model (Andoni et al., 2018), where it is assumed that the processed data can influence communication patterns. At first glance, introducing such an oracle might not seem particularly useful, mainly because it is fixed for each imput data, where the data can be real-valued, which implies potentially unaccountably many communication oracles for a particular task. However, its importance lies in the fact that for a variety of parallel algorithms, the communication pattern at each round depends only on the length of the input and is independent of other input data. For example, in algorithms that compute basic statistics such the sum minimum or tasks such sorting list of numbers, the communication between machines is determined not by their values but by their positions. This means that if the values each machine change, the communication pattern between machines for given task and input length remains the same. In Appendix A.3, we illustrate communication patterns for some of these tasks, which are also addressed in the experiments in Section 7.
Definition A.2 (PCOC model). The PCOC model described by set of n machines labeled from to n, the number of rounds R. an integer& and an oracle RCV: [R] [n]] + [ni] (P([s) (0}) (which is fixed for given and input data) satisfying the following:
1. Each machine i has local memory MEM, E T of size 8, where some abstract data-type. The contents of the memory are indexed from to s and we use the notation [EM to refer to the element at position j on the memory of the i-th machine.
2. Each machine performs some local computation on the data it receives and overrides the results to its memory. A single machine can perform different local computations different rounds and different machines (generally) perform different local computations.
3. When (r, i), where r E [R] and i E [n]. is passed to the oracle it returns a subset M of the set [r|] The oracle essentially returns a (possibly empty) set of machines that machine i has to receive some data from in round r along with the exact positions on the memories of those machines to retrieve.
(P((s) \ (0)).
4. The total size of data sent and received by each machine is at most s. Size here is measured in terms of the number of "variables" of data-type T.
The protocol executed in R rounds. At the start, the input is distributed across the n machines. At the beginning of round r, each machine i simultaneously queries the oracle with input (r, i) and receives data from the machines it returns. The machines then simultaneously perform their local computations on the received data.
Input: Data = (Data),.... Data.) distributed across the memories of n machines, labeled in [n]] An oracle RCV, naata [R] x [n] + [n] (P([s) {0)).
1 For each round r = 1,.... It then
2 Each machine simultaneously queries RCV with (r, i) as input and receives
data from the machines and memory positions returned by the oracle.
3 The machines simultancously perform local computations on the received data and write the results in their local memories.
14

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
A.3. Illustration of parallel algorithms
In this section, we further expand on the discussion of Section 5, stating that communication in several parallel algorithms depends only the identification of the machines rather than their current values. To illustrate this, we provide some concrete examples.
(Cumulative) sum/min.
Sorting (Odd-Even sort)
(0)
(1)
12)
(0)
(1)
(2)
(3)
1
11
11
11
1
(0)
(1)
(2)
(0)
(3)
(2)
(8)
2
2-2
222
2
2
222
222
(0)
(1)
(2)
(0)
(I)
(2)
(3)
2'3
Z
23
(0)
(1)
(2)
(0)
(3)
(2)
(8)
4
2
2
224
2
244
Figure 7: Illustration of the computational graph for algorithms such as (cumulative) sum/minimum (on the left) and sorting (on the right) for 4. Circles indicate the machines, each indexed by the subscript. The superscript indicates each round of the algorithm. At round O, machine holds the i-th element of the input. The cumulative version of the sum/minimum algorithm includes all arrows (black and orange), while the non-cumulative version is represented only by the orange arrows.
These tasks are examples of those presented in Section 7. Note that these illustrations do not indicate the computational graphs derived by our architecture, as there are multiple ways to achieve the same goal, and they do not necessarily involve the neat graph structures shown in Figure 7. For a more in-depth analysis of the results obtained by our architecture, we refer the reader to Appendix D.
In these computational graphs, we represent each machine by a circle, distinguished by a subscript (from I to 4, since 4). Furthermore, we use superscripts to denote the rounds of the algorithm, with superscript 0 representing the initial stage where no computation or communication is performed. Note that no specific values are provided in these examples. This indicates that the correct results can be achieved by following the computational graph for any set of values. In the subsequent rounds, each machine receives information from other machines (including itself) and performs some computation. For each algorithm in Figure 7, we will briefly describe the computations involved, noting that this is not the main focus of the paper and serves only motivating examples.
For the computation of the minimum and the summing function, each machine applies the minimum (or sum) operator to all incoming values from the previous round. By iteratively applying this operator over multiple rounds, machine 4 ultimately obtains the global minimum value (or total sum) across all n entries, while the other machines hold cumulative minima (or sums) up to their positions. For the non- non-cumulative versions of these algorithms, the local computations are the same as the cumulative versions, and the communication paths are denoted in orange and form a binary tree pattern.
For sorting, the graph on the right of Figure 7 represents the odd-even transposition sort algorithm (Habermann, 1972). This algorithm works by alternating communication between adjacent machines, starting with the odd-indexed pairs (machines with indices 1 and 2, 3 and 4, etc.), then switching to even-indexed pairs (machines with indices 2 and 3, in this example). In stages where two machines communicate, the machine with the lower index picks the minimum value among the two, while the machine with the higher index picks the maximum. The procedure runs for a total of n 1 rounds, after which the values are sorted in ascending order.
A.4. Additional definitions
In this section, we provide important definitions that are used throughout our theoretical results. For n E N : (1...... }. we use [n] to refer to the set {1...., n). For two nonnegative integers m, n E U {0} we use m Y n to refer to the exclusive or of m and n. Next, we define the notion of Hamming neighbors:
15

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Definition A.3 (1hh Hamming neighbors). Two nonnegative integers x; [n] U (0) are HhM-mmmning neighbors, denoted j y, if:
j y = (bin(z); bin(y); Vi + and (bin(x)) A bin(y),),
where bin(x) is the binary representation of æ with [log2 n bits ordered from most least significant bit.
Finally, give the definition of Beneš network (Beneš, 1965):
Definition A.4 (Benek network). Let n = _², k E N. A Beneš network is layered directed graph consisting of L = 2k layers with each layer having n nodes and 2n directed edges between two sequential layers. Nodes in each layer are labeled with integers from 0 to n We identify each node in the network by tuple (f, i) where l E [2k] is the layer, and i {0, 1, ,.....,,.. - } is the node label within layer L. Two nodes (L, and (e) are connected by an edge iff
l' = f 1 and (i (i=3) j or ---1k- j)
An example of a Beneš network with n 2 nodes each layer is shown in Figure
0
2
2
2
3
3
3
4
5
5
5
6
6
6
7
7
Figure 8: A Beneš network with 23 nodes in each layer.
B. Expressivity results
B.1. Proof of Theorem 5.1 (Positional Transformers can simulate MPC)
In this section, we prove our main result that establishes the simulation of Massively Parallel Computation (MPC) model by Positional Transformers. Towards this goal, we present two lemmas that bridge positional Transformers and MPC. More specificalll,, we utilize proxy model that has static communication network. We call this proxy model Parallel Computation with Oracle Communication (PCOC), established in Definition A.2. We note that PCOC is only used to establish our expressivity result. It is not meant model to compete with established parallel models, such MPC. We begin by outlining its main features, followed by the two lemmas that lead to our final result.
The PCOC model operates on n machines labeled from 0 to n 1 over rounds, with each machine having a local memory of size s. The model consists of two fundamental steps per round: computation and communication. During computation, machines perform local computations on data stored their memory. In the communication step, machines exchange data through an oracle that provides destination information for each machine and message, executing communication pattern that is agnostic to the processed data and depends only the input length. Critically, the oracle's communication pattern remains fixed for given input length and task, meaning that if the values at each machine change, the communication pattern between machines remains the same, property that is observed in many parallel algorithms such as computing basic statistics or sorting, where communication is determined by positions rather than values.
This contrasts with other parallel models, such as MPC, where the communication pattern more flexible and directly influenced by the data being processed. the MPC model, each machine has a limited memory size of 8, and the input data
16

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
is distributed among the machines. The computation proceeds in rounds: during each round, machine processes the data in its memory, exchanges messages with other machines (up to total size of s), and retains only the messages it received for the next round. After all rounds, the final output is stored in the memory of one or more machines. Unlike PCOC, the communication in MPC is not predetermined by the imput length or task alone; instead, it can dynamically depend on the data being processed, provided the total communication per machine per round does not exceed the memory limit 8.
Even with limited communication capabilities, PCOC can perform the same functions as MPC, incurring only a logarithmic cost in rounds relative to the number of machines. N. The proof of this lemma presented in Appendix B.2.
Lemma B.1. Consider an instance of MPC T with R rounds, N machines with local memory s. For any instance T, there exists an instance of PCOC with 2R [log N1 rounds, N machines with local memory 2s that simulates T.
An analogy to illustrate the distinction between the two approaches can be drawn from a mailing system. In the MPC model, each processor acts as an individual delivering letters directly to their recipients. Communication flexible and determined by the specific needs of the task. In contrast, PCOC resembles structured mailing system with predefined routes and sorting stations. Each individual sends their letters to the nearest sorting station, where the letters are organized into bins and passed through successive stages of sorting and delivery. After several rounds, all letters reach their destinations, following a fixed communication pattern regardless of the specific contents.
Having established the PCOC model, we now show that the positional Transformer model in Equation (1) can simulate it. More specifically, our results show that R-round PCOC protocol can be simulated by a positional Transformer with layers.
Lemma B.2. Consider a PCOC instance P with R rounds, N machines with local memory 8, and data type T be a model following the architecture in Equation (1) with n N + 1 nodes, R layers and s attention heads. Then, for any instance P with Borel measurable local functions, there exists a configuration of M that approximates P to any desired degree of accuracy.
R. Let M
The proof of the above lemma be found in Appendix B.4, where we show that each round of computation and communication performed by PCOC can be approximated by one layer of positional Transformers. Having established the two necessary lemmas Lemma B.1 and Lemma B.2, we are ready to re-state our main result.
Theorem 5.1. Consider an instance of MPC T with R rounds, N machines with local memory 8. Let M be model following the architecture in Equation (1) with n = N + 1 nodes, 2R[log N1 layers and 2s attention heads. Then, for any instance T with Borel measurable local functions, there exists a configuration of M that approximates T to any desired degree of accuracy.
The proof of this result follows directly from the combination of the two preceding results: MPC can be simulated by PCOC, which can be simulated by positional Transformers. Notably, the structure of MPC enforces that local memory can only grow sublinearly relative to the input size. This constraint prevents trivial simulation configurations in which all data sent to single node, which would undermine the purpose of parallelism.
B. 1.1.1. DISCUSSION OF REMARK 5.2
In Sanford et al. (2024), the authors present not only a simulation of MPC using standard Transformers but also the reverse, i.e., a simulation of Transformers by MPC. From the latter, we derive the following remark that establishes an equivalence between standard and positional Transformers.
Remark 5.2. Consider standard Transformer M' with N nodes, L layers, and dy H sublinear in N. Then, there exists positional Transformer M with O(N2) nodes, O log N) layers, and dy H sublinear in N that can simulate the computation of M'.
An important observation from Sanford et al. (2024) that the embedding dimension, denoted as dy, also governs other intermediate dimensions, such as the embedding dimension d,, of the matrices WK and Wo- However, in positional Transformers, these quantities are decoupled. Specifically, the embedding dimension dm is directly influenced by the positional encodings and scales linearly with the number of nodes.
is worth noting that the observed quadratic dependency on the number of nodes in positional Transformers is primarily a consequence of the specific simulation strategy employed in Sanford et al. (2024). These dependencies should not be taken too literally, as the results do not provide a sharp characterization of the two computational models, an aspect acknowledged
17

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
by Sanford et al. (2024). Moreover, alternative and more efficient simulation strategies may exist. Additionally, this quadratic cost arises only when simulating the same sequence of operations across architectures. In practice, the model may converge to parameter configurations that do not correspond to interpretable algorithms.
B.2. Proof of Lemma B.1 (PCOC can simulate MPC)
The goal of this section is to show that MPC can be efficiently simulated via PCOC (that is, in data-agnostic way). For completeness, we re-state the lemma below:
Lemma B.1. Consider an instance of MPC T with R rounds, N machines with local memory s. For any instance T, there exists an instance of PCOC P with 2R [logg N rounds, N machines with local memory 2s that simulates T.
We break the proof in two steps: first, we show that MPC routing can be simulated through fixed communication pattern given by the connectivity of Beneš network. Subsequently, we show how this routing can be accomplished in data-agnostic way, consistent with the PCOC definition.
Consider an MPC instance with machines, each with memory size s. Assume messages are encoded in the following way: machine i holds list SEND; containing tuples of the form (i, b, pos. where b E A, pos E N and E [p], where b is the pos-th character the message machine sends to machine j. Here, A. denotes the alphabet used to encode the messages (for example, A can be {0, } if we are representing messages by their bits). As a concrete example, if machine sends 010110 to machine 1 then
SEND, = (0, (1.1),1))) 1,,1,
By the assumptions of the MPC model, |SEND;| S for all i {0, p,. 1}. Similarly, denote by RCVD, the list of source-characcerppiitoooddiiiiid
tuples that machine receives. That is:
RCVD, = {{ (1s, b, pos.
U, SENDy |j =1} for {0.1...... 1
Again, by the assumptions of the MPC model, we have |RCVD,| S for all i E [p]. We further impose the following simplifying assumptions:
1. The number of machines is a power of two, i.. p= 2 for some & e N, and
2. For each pair of machines i, j. the number of tuples sent from machine to machine j is either 0 or at least 2 where 1t as above.
Note that the above assumptions do not limit the MPC model since power of two and pad shorter messages with special token (potentially having scale s). The second assumption be lifted altogether at the expense of more complicated proof. We chose to work with both assumptions in order to highlight the essence of the proposed routing algorithm. As we mentioned previously, to highlight the routing strategy we split the analysis into two parts: first we show how we can simulate one round of MPC using a Beneš network routing without the need for extra memory, assuming access to destination information. Building on that, we then show how to extend the Beneš routing protocol to the case where message routing independent of destinations, at the expense of doubling the memory of each machine. The latter shows that PCOC can simulate one round of MPC with the same number of machines and double the memory. The result is then trivially extended to R rounds of MPC by iterating the routing algorithm R times.
always add couple dummy machines reach
Data-dependent Beneš routing We begin by providing a data-dependen routing algorithm that simulates one round of MPC using a Beneš network. Under the two assumptions made above, the routing algorithm executes in 2 log, (p) = 2k: steps and is as follows:
Initially, the memory M, of each machine holds the list SEND, (i.e. the messages needs to send in the source- Characce-postioondddsttatoo format described above).
step I. (1,..., , k- - 1}. machine partitions the contents of its memory as
P
U
Ma
j=1
18

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
where M, {(k; b, pos,j") E I -j.. Let (L 1, denote the neighbor of (L, i) at level L. + with a label different from i (so + i). Machine partitions each Ma.; into two equal parts (this is possible since every message is power of two) and sends one part to machine i and the other part to machine m. Steps through A: 1 are collectively referred to as the "fragmentation phase" of the algorithm.
steps L E (k...., , 2k- - 1). denote by RCVD.") (4) the after L - 1 steps. Also, let (L + 1, m) denote the neighbor of (L, i) at level L + with a label different from i. At step L, machine sends all tuples (ki b, pos, E RCVD such that there is a path between nodes (L 1, and (2k, j) in the Beneš network to machine i, and the rest to machine m. Notice that by the construction of the Beneš network, exactly one of following holds:
tuples received by machine
There is a path between nodes (L + 1,) and (2k: j). or
There is path between nodes (L + 1, m) and (2k,
Steps k through 2k- 1 are collectively referred to as the "collection phase" of the algorithm.
At step 2k: the nodes in the last layer recombine the received messages using the source and position information and the algorithm terminates.
We show correctness of our algorithm using two-step argument. First, we show that during the fragmentation phase (first A 1 steps of the algorithm) the tuples are evenly distributed ensuring that the at every point no machine receives more than tuples. Then we proceed by showing that messages are correctly routed and reach their destination.
Fact 1. After I E (0, , 1} steps, machine i receives exactly 1/24 of each SENDy where i' is such that and i' agree on the first k- L least significant bits. Furthermore, no two machines receive the same part of the same list.
Proof. The proof is by induction. For == 0 the result trivially holds by the initialization step of the algorithm. Assume it holds for some L < k - 1. Let M, be the memory of machine i after L steps. By the connectivity of the Beneš network, after L + 1 steps, each machine will have received messages from itself and machine m := i 2 From the inductive hypothesis, M, contains 1/2" of each SENDy where i' is such that and i' agree on the first h least significant bits, and similarly, Mn, contains 1/24 of each SENDs where i' is such that m and i' agree on the first h: I. least significant bits. Now i and i' agree the first k L 1 least significant bits if and only if they agree the first k L least significant bits or m and i' agree on the first k- - L least significant bits. By construction, machine i will receive half the contents of M, and Mm and it will hence hold exactly 1/2-+1 of each SEND, where i' is such that i and i' agree on the first k (L - 1) least significant bits. Since there is at most one path between any two nodes in the Beneš network, the second part of the statement is trivially satisfied. This concludes the proof.
Fact immediately shows that for the first k - 1 steps of the algorithm, the contents the memory of each machine never exceed 8 (since the length of each SEND; is bounded by s). Similarly, we can show the following fact:
Fact 2. After L E {ki, l...., 2k 1} steps of the algorithm, machine receives exactly 1/22k- 1.-1 of all RCVD, where j is such that i' and i agree on the first L k least significant bits. Furthermore, no two machines receive the same part of the same list.
Proof. We shall only prove the base case L = 1 as the inductive step is completely symmetric to the proof of Fact 1. After the k-th step of the algorithm, machine i will receive all tuples for which the destination j is such that (2k: is reachable from (L+ 1, i) in the Beneš network from itself and from machine m == 2 Recall from Fact I that before the k-th step executed, the memory of machine i, M,, contains 1 of every SEND, where i' is such that i and i' agree on the least significant bit. Similarly the memory of machine m, Mn. contains 1 1/2b-1 of every SEND, where i' is such that m and i' agree on the least significant bit. But since the least significant bit of m and differ, the memories of machines and i combined, M, U Mm, contain exactly 1 of Up SEND, (i.c. a 1 1/26-1 fraction of all messages). Combining the above we see that after the k-th step is executed, machine will receive 1 1/26-1 = 1/22h of every RCVD, where j is such that and j agree on the first 6-K k + least significant bits. As in Fact 1, the uniqueness of paths in the Beneš network guarantees that no two machines ever receive the same part of the same list.
19

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
A direct application of Fact 2 shows that after 2k- steps of the algorithm, machine each machine will have in memory exactly RCVD; showing correctness. Finally, the fact that |RCV;;| S for all E [p] shows that after every step L E (k, k + 1,... 2k 1) of the algorithm each machine receives no more than tuples. This construction shows that even if we restrict the communication of MPC to some fixed routing pattem, we can still achieve any and all possible routing patterns by using the Beneš network efficiently, paying the cost of logarithmically more rounds. We now proceed to show how we can further extend this algorithm to be consistent with our PCOC model.
From data-dependent to data-agnostic Beneš routing The algorithm above brings us one step closer our ultimate goal of data-agnostic simulation since at every step, each machine will send data at most two other machines, respecting the connectivity of the Beneš network. However, both the fragmentation and collection phases of the algorithm (steps 0 to k - 1 and & to 2k- 1) implicitly require access to destination information for each message. In terms of PCOC, assuming the oracle mimics the algorithm execution, it would require access to destination information, which is not allowed in the PCOC model. This is because of the following reasons:
1. During the fragmentation phase it needs to know exactly which messages destined for each machine j to evenly divide them into two parts and route them to the two receiving machines.
2. During the collection phase, destination information is directly used to route messages.
While, at first, these obstacles might seem hard to circumvent, we can do so by using a few simple techniques that enforce a specific structure on the output of the local computation:
1. Observe that we can eliminate the need for destination information during the fragmentation phase if we assume that the tuples in the memory of each machine are always sorted based on their destination. If this is the case each machine can simply send the tuples the even and odd indices on its memory to the two receiving machines, respectively. This achieves the desired outcome: for each j. exactly half of the messages with destination will get routed one machine and the rest to the other.
2. As previously discussed, during the collection phase, the oracle would need access to the destination to decide where to route each tuple. However, if we assume that the memory size of each machine is 2s (instead of s) we can still overcome this issue. Suppose the oracle needs to simulate step L E [k, k + 1, 2 1} of the algorithm. For each machine i, let im # i be such that (L, i) and (L + 1, im) are connected in the Beneš network. Recall that at step k, machine will send data to itself and machine im. depending on the final destination. Given the fixed structure of the Beneš network, we can require the local computation to calculate where each tuple needs to be routed (i or im) and place either the first half of the memory slots (if needs to be routed to i) or to the second half (if it needs to be routed to ) The oracle will then route the first half of the memory to i and the second half to im. Empty spaces on each half can be padded with special tokens.
B.3. Hardmax patterns using positional attention
In this section, we show that the positional attention architecture in Equation (3) can approximate any unique hardmax pattern, a concept we define later in this section. This result used support the simulation results of Lemma B.2. We begin by stating the definition of the row-wise hardmax transformation for a p x q matrix X from Section 3:
hardmax(X)) = {
if Xx, = max X
for i E [p]. E [d].
(6)
otherwise
where we implicitly extend the definition for vectors in Ip" by viewing them as x n matrices.
We use the term hardmax pattern to refer any matrix in the image of hardmax (i.e. a binary matrix with at least one non-zero element in every row). Furthermore, we use the term unique hardmax pattern to refer to hardmax patterns with exactly one non-zero element in every row. Unique hardmax patterns occur when the input matrix has unique maximum value in every row.
We further define key concepts that will be used for the more formal re-statement of Lemma B.3. Let the input have n rows, and the binary positional encodings be defined by the positional encoding matrix P = In. therefore having dp = n. Finally, let T be a positive scalar that represents a temperature parameter that controls the approximation of softmax.
20

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
unique hardmax pallern A, there exists configuration of node positional attention parameters in Equation (3) and temperature parameter T such that the resulting softmax pattern A approximates À to any desired level of accuracy. Formally, for any unique hardmax pattern A and any & > 0, there exists some T = TT() > O such that the inequality [As. Ai.i S & holds for all i.j E [n).
Lemma B.3. For any given
Proof. Without loss of generality, we may assume E < 1. We start by setting the node positional attention parameters to be Wx = n and We = T(2A- - 1). where, T > 0 is our temperature scalar parameter and À is the target pattern. Since, in this construction, node positional encodings are set to be the identity, the inner-product PWq W W.T. PT reduces to Wo. where each entry (i
A. = I,, and -T' otherwise.
This inner product is passed to the softmax operator, resulting in the attention matrix A. For each i, [n/] we separately analyze the following cases:
1. Case A. 1: In that case, the only non-zero element on the i-th row of A A. J. so we can express the difference as
exp((WQ)).))
exp(T)
An, Ass =1
Σx= exp((WQ)(.)) =1 exp(T) + Ekez exp(-I)
exp(T)
I exp(T" + n exx(--T)
= 1 exp(((1 27")
exp(In n 2T)
1 1 exp(lnn - 27) S Exp((nnn - 2T)
2. Case Ar. = 0: Let jo # j be the unique index for which Ai.jo = 1, and we can express the difference as:
exp((WQ):,)
cxp(-T)
A Ars Σi=1 exp((WQ)1.) exp(T) + Ei exp(-T)
exp(2T) - n exp(2T - - +1 S exp((nnn 27)
In any case, we have that [A.,- - Ai.il exp(In n n-27). Therefore, by taking T 2 1/21n(%/c), we have Ais-Ais Atjl Se. LI
B.4. Proof of Lemma B.2 (Positional Transformers simulate PCOC)
We begin this section by outlining the key concepts utilized in the routing protocol employed in our constructions. First, we describe the general structure the input matrix X.
B.4.1. ENCODING
Input matrix: In alignment with the PCOC model, the input matrix X represents N machines, where each machine is denoted by a row in the matrix, and its local memory, MEM, Ts, is represented by the corresponding columns. The maximum size of data that any machine can send or receive is bits, with each bit corresponding
column in X.
However, the actual number of rows and columns in X differs from the number of machines and the local memory size for two reasons:
1. Sink node: A dummy node is introduced to facilitate all possible communication patterns in PCOC using positional attention. This is necessary because PCOC allows for the possibility of information not being sent to any receiving machine. This scenario incompatible with the softmax formulation, which requires at least one non-zero entry. The dummy node serves as a sink, collecting all messages that do not have a destination. Consequently, the number of rows in X is N +
2. Unique node identifier: Each machine also requires a unique identifier to enable element-wise local computations. To achieve this, we encode a unique scalar for each node in the last column of X, resulting in a feature dimension of dx = 8 + for the input matrix.
21

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
As discussed in Section 5, in PCOC, routing is set by an oracle that decides how packets of data should be routed each round. Under this framework, routing must be performed prevent multiple data from being sent to the same destination. Since our construction relies on matrix operations, this leads to the following assumption:
Assumption B.4. (No-collision). For any layer & E [L] no two different machines i 12 E [N] should route data to the same destination i3 E [M] for the same column j E [s]. where each column represents a bit of local memory across the
nodes.
Note that this assumption does not limit the generality of our PCOC model. only defines how data should be stored the memory of each receiving machine, and any valid PCOC routing has corresponding no-collision configuration of bits that realizes it due to the restriction on the total size of received data. As demonstrated in the constructive proof, this directly influences the sparsity pattern generated by each attention head.
Positional encodings: As previously mentioned, although the connectivity at each layer may vary, the positional encodings remain consistent across all layers. Our architecture simulates MPC using node positional encodings with dimension dp by setting P = In. with each positional encoding functioning as a standard basis vector.
B.4.2. SIMULATION RESULTS
We now demonstrate that, with the established encoding, the architecture provided in Section 4 can simulate the execution of any PCOC instance. Each round of such PCOC instance can be decomposed into two stages: communication and computation. Our objective is to provide existential results for both stages.
In the communication stage, routing assigns destination machines for each message. In our architecture, this assignment is analogously captured by the attention weights, which determine whether a message should be received by a node using binary values.
The no-collision assumption ensures that all routing patterns can be represented by unique hardmax patterns. As expressed in Lemma B.3, since any unique hardmax pattern can be approximated by our attention layer using softmax, for simplicity, the subsequent proofs use hardmax instead of softmax. With all details now established, we re-state main simulation result:
Lemma B.2. Consider a PCOC instance P with R rounds, N machines with local memory s, and data type T R. Let M be a model following the architecture in Equation (1) with n N + 1 nodes, R layers and attention heads. Then, for any instance P with Borel measurable local functions, there exists configuration of M that approximates P to any desired degree of accuracy.
Proof. Despite the desired degree of accuracy being influenced by the number of rounds performed, it suffices to show that one layer of our architecture can simulate one round of PCOC. The same constructive arguments can be extended to more rounds, ensuring the overall degree of approximation is respected'. To this end, we begin the proof with the communication stage.
Communication: In PCOC, communication is encoded routing patterns determined by the oracle. At round E E [R]. we denote by (f) = {((i,)), K) j.j E [NV, K E P([s)) the set of valid routing patterns provided by the oracle. This set specifies that the data at positions K in the local memory of machine i must be sent to machine j at the same position. A valid routing pattern requires that no collisions occur (i.c., no two triplets in (4) should have the same destination j
=) ((, K) E H(0, = E K} the subset of routing and memory position is E K). We further denote by FT! patterns corresponding to position a in local memory.
The first part of our result constructively demonstrates that positional attention can reproduce any valid routing set by the oracle that adheres to the PCOC model. We construct & attention heads indexed by E [s). which handle routing for the corresponding subset Hy.
For clarity in the construction phase, we introduce an augmented set to simplify notation. We begin by extracting the set of source nodes for each set ! denoted as ! = {i (()))) E H j E [N], E [s]}. Next, we create a complement
TTo completely align the simulation results with PCOC execution, positional Transformers require an extra initial layer where the attention layer acts the identity. This is to account for the fact that PCOC first starts with computation stage rather than communication as Transformers.
22

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
set s which routes all unused sources (i.e., those not in 7! to the sink node labeled n. N + 1. We denote this complement set by H = {(t, n, 2) i E [x] \ 7. Finally, we define the union of these sets as H H!!) u
The attention parameters are then set as follows:
(wrk))))
ifi=j, j, 1.j 0 otherwise,
1
and (w((t)) in {
if (i,i, E 1/ otherwise,
In this construction, we first observe that both the node positional encodings and the key matrix are identity matrices, reducing the inner product in attention to be solely defined by the query matrix. The query matrix is then designed to encode the source node i as standard basis vector in the row corresponding to the destination node j. This effectively represents
the routing set ( as a binary matrix, which is also preserved after applying hardmax. Additionally, the no-collision assumption, combined with the sink node strategy, ensures exactly one non-zero entry in the first N rows of the attention weights matrix for each attention head h.
For the value and output transformation, we set all value matrices (E,))) to be the identity Lax and define the output matrix w! E ( as follows:
(w:)) is -
{
(n 1)8, j
otherwise.
Here, the output matrix O ensures that only the correct memory position receives the information and places it in the corresponding column. Note that since the outputs of the attention heads are concatenated before being processed by W(() the values along the rows of the output matrix also depend on the attention head.
We now focus on the computation stage for the second part of the proof.
Computation: At round l [R] of PCOC model, let (Ø) ())(,)))_ be the local functions applied by each machine E [N] and let ((,z) correspond to the function of the augmented sink node, which effectively erases all data received. Each
function of (f,z) operates on received data in each machine's local memory. which corresponds to the output of the attention layer, denoted by and outputs vector of the same dimension that is, of
(f,x) R + R
Furthermore, let ((f,z) _dx + [dx be a function common to all nodes, which operates solely on the residual connection x. This function outputs a vector where all entries are zero except the last entry. The value this last entry corresponds to the unique node identifier extracted from the residual input x.
We aim to approximate both σf (L.z) and glf.z) using neural networks. To this end, we define the combined function dxx++ + R.dx by:
o (z r.) = ((t.0) (x:) + 0. (f.2) (z,) 00,
(L,a) with where 2, @ Iy denotes the concatenation of output of 2 R and x; e Rdx, We further augment the output of of zero scalar to match the dimension dx= = + 1.
Let ( 1 correspond to multilayer perceptrons (MLPs), each applied to each input @ xi- By invoking universal approximation results such as those by (Cybenko, 1989; Hornik et al., 1989b), we assert that as long as the local functions s are Borel measurable, there exist neural networks ((0) that can approximate the functions (0) to any desired degree
of accuracy. Additionally, note that the function (f.=) of the sink node, as well as the function (t.xr that operates on the residual connection, are both linear and therefore Borel measurable.
The final step in this argument is to relate these approximations the proposed architecture in Section 4. Specifically, use the MLP ((() in Equation (1) and leverage the aforementioned universality results to approximate all the element-wise functions o!
23

---

Positional Attention:
Expressivity and Learnability of Algorithmic Computation
A crucial aspect of this step is the need for the input of each machine to be uniquely identifiable. This ensures that single model can injectively encode multiple functions. Intuitively, it guarantees that each approximation of the local function can identify that it is processing the right row. The unique identification of each machine is guaranteed by the scalar encodings of every node, which, regardless of the contents in local memory, ensure that the input rows are unique, Therefore, the function that 4' has to approximate is a piecewise Borel function with each branch being one of the of based on the unique machine identifier. Such function is Borel measurable, and so the universal approximation results of (Hornik et al., 1989b) hold, guaranteeing the existence of the desired MIPP ((().
This demonstrates that our neural network architecture can emulate the computations performed by the local functions (f(..) acting the output of the attention layer (with their outputs zero-padded match the required dimension) and the function g(fxx) acting on the residual connection, even though they act on distinct parts of the input.
Therefore, we establish that our proposed architecture can approximate the computations in each round of the PCOC model.
An important observation is that the computational model and expressive results for the proposed architecture are specific to a fixed input length n. Furthermore, one could extend such results to a model with communication and local computations that also consider the input length as an input. For local computations, proof in Lemma B.2 can also cover such cases, provided that the information about the length is also encoded in the input. For communication, we present the following remark.
Remark B.5. For any collection of unique hardmax patterns )==) where A(k) is l x k; there exists a configuration of node positional attention parameters in Equation (3) and a temperature parameter T such that the resulting softmax patterns {A(0) m
k= approximate each (A) to any desired level of accuracy. Formally, for any collection of unique hardmax patterns {A(K) In 1==1 and any 0, there exists temperature parameter T = T(e) and corresponding attention parameters such
that for all i E [n]] and for all k E [mi]. the following inequality holds: A10
(R) S.
The proof of this remark relies on a slight modification of the proof of Lemma B.3. However. to cover all possible patterns, the embedding dimension of positional encodings should also encode the input length and be of the order of O(n"). Although this embedding dimension is theoretically large, in practice, one does not need many dimensions for positional encodings, as demonstrated in the variable length experiments in Section 7.
B.5. Softmax patterns using positional attention
We conclude the discussion expressivity by showing a final, standalone, result, namely that the positional attention architecture in Equation (3) can represent any softmax pattern. This result is further discussed in Appendix F regarding reasons for poor -fitof-diitribution performance by standard Transformers.
We begin by stating the definition of the row-wise softmax transformation for a matrix X E RPX9:
exp((..)) softmax(X)ss Σn exp()...)
for E [p].j e [4]
(8)
As with hardmax, the definition is implicitly extended to vectors in R by viewing them as 1 x n matrices. The image of the softmax function is the set of row-stochastic matrices with entries in (0, 1). Indeed, it is easy to see that when softmax is applied to a matrix, the resulting matrix satisfies the above property. On the other hand, for a matrix B = (biz (be)) with biy E (0, and Ejetgl by = 1 for all E [p] we have softmax(X) = B where X In(byy). ). We use the term softmax pattern to refer to any matrix in the image of softmax.
Consider attention weights ((EA)) that are defined by positional encodings in Equation (3). Let B E (0, 1) be a softmax pattern. We would like to find parameters w-(t.h) and ((...h) that induce B, that is ((C B. From the properties of softmax described above, it suffices solve the matrix equation (PW.(."))) (PW)(A)) ß where B, = In ln(B.)). This equation always has solution when dp and P is invertible. We summarize the above observation in the following expressivity remark:
Remark B.6 (Positional attention is expressive). Positional attention can realize all softmax patterns at every layer provided that dp = n and P is invertible, This is not necessarily true in the case of standard attention where, in subsequent layers,
24

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
positional encodings are modified and, therefore, guaranteed be linearly independent.
C. Proof of Generalization Bound (Theorem 6.1)
This section provides a complete proof of Theorem 6.1. The high level strategy which inductively builds a cover for the class of multi-layer architectures is similar to that of Edelman et al. (2022). We made couple of changes to account for positional attention. For the reader's convenience we re-state all relevant definitions below.
C.1. Preliminaries
C.I.1. NOTATIONS AND THE NETWORK ARCHITECTURE
For a matrix we denote by Z, and Z the i-th row and column of Z, respectively. || ll2 denotes the spectral norm for matrices, and lp.q denotes the (p, q) matrix norm where the p-norm is over columns and q-norm over rows. For vectors, I1p denotes the Ep norm. For example, 22 == max, Z |l2 and || Σ, |2,,
We interested in the positional transformer architecture of L layers, which obtained by composing Equation (1) times. In addition, parameterize 4(() to be 2-layer MLP for each L. Given input X Ra and positional encodings P E Raxdr, the resulting architecture has L layers, each layer has H attention heads, and it is written as
F(0) (X;WWo)))) = X, Vl = L:
FrD (X: (x: = A (E.h) p(t-1)(X; w1:--1 ( F(l- (x: ) w//) w/o
(9)
with ((eM) == solt PT), =h= = 1,2...,, H.
In the above, an La Lipschitz activation function with o00) 0. We write
w (w( for the set weight matrices at layer e, and
(1.1) w(le2) wh( (w (222 w (2.B) w(el)
(2.1) w.(e2)
W (e.H) .:( (00,
w.(0)
wie = (W(1), w(2. w())
for the set of weight matrices up layer f. Note that may assume without loss of generality that the output projection matrix w-(O) is absorbed into the weight matrix w of the MLP, and therefore we omit w in Equation (9). Denote max (dx, dp, dout). Since we are to derive an upper bound which increases with dx dp, dv, dout, we sometimes replace dx, dp, dv, dout with their maximum d in the derivations to simplify the notation.
Given the final output Xout (T) (X; W R"x*, extract scalar output from the architecture w X [n] with trainable weights E Rdd This defines class of scalar-output architectures as in Edelman et al. (2022). The index i E [n] on which we apply the linear function æ + w I can be arbitrary, and here we simply set to the last index, n. Our generalization bound applies to the following class of scalar-output architectures with bounded-norm parameters:
F:
X + T (F())(X;
||w||2 S Bw.
/
S B2, (4), Vh E [H].
(10)
2
2
W
We (ea)
|
(
S B2,, Ve [4] th e [M]
2.1
2,1
2,1
2,
25

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
C.1.2. GENERALIZATION BOUND, RADEMACHER COMPLEXITY, AND COVERING NUMBER
Definition (Risk), Let D be a distribution over IR, and let é: IR R + I be the loss function. For a given class of functions F (f: X R) and f F, the generalization error, the risk, and the empirical risk, are defined
risk(J;D) = BB (((((r),)),
risk(/: (1
(1,))))
1=1
Definition C.1 (Empirical Rademacher complexity). For given class of functions J {f :2 + R} and fo(0) R C X the empirical Rademacher complexity R(F (ax00 (=m)) is defined
R(F; ( )= m
1
172
E (
sup fEF
where is a vector of i.idd Rademacher random variables, that is, IP(e 1) IP(:; = =
The following theorem bounds the generalization gap Irisk risk using the Rademacher complexity of a function class. Theorem C.2 (Bartlett & Mendelson (2003)). Let D be distribution over X lK and let é IR XRR - R be b-bounded loss function that is L-Lipschitz in its first argument. Let J = {f X + Ry be given class of functions. Then for any 5 > O, with probability at least 1 - 8, simutaneously for all f E F, it holds that
log(1/0)
Di
)0) < ALR(F; (x)) ) + 261
2m
We may bound the Rademacher complexity of function class F using covering number. Definition C.3 (Covering number). For given class of vsluld or vector-valued) functions F, the covering number
No. eF;;; (a(0)m Meiil (I) is the smallest size of collection (a cover) C c F such that for all f E F there is f E C satisfying
| - f(z(0) S..
max sE(m]
Lemma C.4 (Rademacher complexity via covering number; Corollary A.3 in Edelman et al. (2022)). For a class of functions F = {f X R} such that [fl consstant > 0:
A for all f E F, suppose that log Noc (F; {=()) S Cy /E?, then for some
CF (1 (1+1og log (AVm/Cs;
R(F;
Ji=l Se. (1)m
Combining Theorem C.2 and Lemma C.4 we may bound the generalization gap rrisk risk] using the covering number of function class, stated in the following lemma.
Lemma C.5. For class of functions FF = {f : X + R} such that f is bounded for all f F and that (200)m I Cz/e² for all {m C X. Then for any > O, with probability at least - 5, si-
multaneously for all f E F, it holds that
Cr (V
log(1/5)
risk(f:D) risk( (z(1), y )=))
Sō
m
m
1.3. SOME USEFUL LEMMAS
Lemma C.6 (Corollary A.7 in Edelman et al. (2022)). For vectors E I² we have that
[|softmax(x) (x) (oltmaxx(y)ll < 2 || ylloo
Lemma C.7 (Lemma 4.6 in Edelman et al. (2022)). Let W (W Rdxdd
W 11,,1 S Bw and consider the
function class Fw := {x - Wa: W E W). For any & > 0 and (1), (2) (m) E Idd satisfying || S Bx for all i E [m), we have
By BB B e² log(dm).
Noc (FW:e;
I (m) (1)
26

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Lemma C.8 (Lemma A.8 in Edelman al. (2022)), For 01, 2 O, the solution the the following optimization problem
min Flomoe
s.t
Σ biE =E
1² and is achieved at E; = e(a1/b:)3/3// where v == Ei=1" 1//0//1
C.2. Lipschitz Property of the Architecture
Denote
( f(z; (W) J (Wkk))H (W/-a)) )m=z)
A zw!i1) 78
h=1
Al:2)
Al.
ZWL 22)
where A(-h) = softmax(( softmax(PW? (PW () w (.h)) PT).
Lemma C.9. For any Z, Z, (W( sh=1. (Wb Jh=1 (W 1 ( M (W( | (W/
(W(-k))
1h=
(z: (
1n H 1==1
(W 11
1h=1 (W )1 - f (2: (w | (w
(i | W
|
2,00
[|W((ah) |12
< max
max km(
( (w.f:) - (x p
12 12,000 + 2212 22100112,001122 1|200
112.001112
max hE(m]
2,00
(whid)_ - w(-hh) T zT
VH max bE(m]
2,00
Proof. Denote AGM) - PT) and A(.h) softmax (Wil W ( PT):
Q
(W/-1)) / (w6 11 ) (WG-)) J1 (WT-)))11) 1/=1) / (2: (wb 1h= (W 1h=1: (W
Ih=1)
wv(.1) (A()))2)
W. (A)
WV(mH) (A(-))2) TT
Wybm (A(.m2))
2
2,00
Wr (A(:1) A())))
(w(-1)) - w (1):)))T
1
:
((GH) - A(.H))) T
w) (AGH))))
-T-
0
2,00
2,00
(A(-1) 2T 2(-1))))
(w - m7)-))) (A
+
W
(A(-H))) T A(-H)
- T (7(
2T 2T
1
2,00
2,00
27

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Wy
(A(-1)2)T= 2)
(wih") - m7b1)) (A
+
(A - A(-H)
(w/ mn T (A(-))2)
1 12
2 2
2,00
2,00
(A(1)))) 2TT
(nr/t) - mm T (2(:1)
T
max
max kE[H]
|
112
+
V max hE|H]
(Al.H) 2) T - 2).
2,00
2 - 2
2,00
||W---A0 112.1 11 ) (1/2T [AC.0T - -A(11)
ACMIT T -(( 7 2,00
< max
max [E(m
+ 1z - T)T AGU1T AC.HT .111 2,00
-
+
H
max he
(wr/: (19::)_W96 -mpp)))))m))))) T (A(M2 2
(1I)
2,00
Using
||oo|22 I || 11ll,,0l1l11 (and therefore
||Q0|2.000
2.00 I
(z T) ( ACHMT 11]
- (2- ((T
[GGD ACHM 1
=
=||2 - -TT
2,00
2,00
1,20
2,00
(12)
and
-
T
(wiith) -
zT
[C-i))
1,,00
2,00
1,20
= -I( - mp):
zT
(13)
2,00
Using || $ |||12,.011Q11, || max,, Prt. S max; ||P.T max, ||Q.,lı = ||PT 2, 1111,,0
and Lemma C.6,
L A ( -1:11) Al.! - -2h11 -|
A
2,00
TZT (A(-A) - - ")l1 2,00
)11.00
max NE(H
AA(-N) [0M T
max ke(H]
||
2,,00
11,00
AAC- -AC-1)
71
max Z KE(H]
1|2,00
00,000
max 2/|2 kEIH
12,00 |P (W>xCA)
-
(M))) Wo PT|
20.00
|| (WWolA)WW <oo--Woo (o
(.N) ( pT
max H
2|
||2,00
| ||2.00
(14)
2,00
Combining Equations (11) to (14) gives the result.
We introduce one more term to simplify the notation. For é = 1,2,.... L denote
000 (X; Wl:)=
A (k.b) p(e- (X; wilt (W: (wtte)) @ F(e- w1:--1 (x;Weee--
(15)
where (((-1) (x;-))(x; (X;: 11 and ((A)) are the same as in Equation (9). That is, (X 1:t is the input to the MLP in the l-th layer of the positional transformer architecture.
28

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Lemma C.10. For any Wid and W
p (x: - F(O (x:
2,00
(w250 -12 m7500
(x: (
+ +Loll (, - m1 g00) (x; W1:e) T
2.00
+L < max bii3
|| 112
||
Whe-l)- - Fre (x; (x:
2,00
+ 2L, | |2 |ell (x; Wr: 1) 112.0011 I|PT
||
2.00
max
max
he[H]
( w(an) - -wani ix ((.)) (CNDT) PT
max kE|
2,00
[ ll2VH
(wiKe mmeen) T ple- (X; Wle--1gT
(X;
max he(m]
2,00
Proof. Using the definition of F(0 (X:Witt) in Equation (9) and the definition of (X; in Equation (15):
(( (x: - FTO)
(x; (x; (x:W.)) T
2,00
o ( (1131111))))110) T
W3 To
2,00
(Using triangle inequality):
s [1/300 - W) (x:
(16)
2,00
+ | (xx (x: )
(17)
2,00
The term in (16) gives the first part of the required bound. We focus on the term in (17) below and show that it is bounded by the second to the last parts of the required bound.
W (o (x(00xx (x) a (x:
2,00
(Using |Q]l2, s $|||||2|||12. and ||P||2 = ||P ||T)
< [|w!!0 le (o (x:W))))))))) -o (g Wi: (9)
(x))))
2,00
(Using Lipschitzness of a and (0)= = -):
< L. L.llw/"'l
|| (o (x w w (0
(X: w.: (
(Using triangle inequality):
(w110 - -10) T (x:
< L. || 12
2.00
+ + 1Lo 112
w, T (x; - (x; Wl:6)
T
2,50
(Using 1Q I ||||2|||||2, and |l|| = |PT|22)
29

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
- m T 900 (x; < Lo (w/
2,00
+ Lo 122 l| | (ox)
whit) (x:
(x: 2,00
(Using Lemma C.9 and definition of (00) W):)))
s Lo 2 112
(w/" - w (x; Whoe)
(w.1
2,20
+ Lo 112lla max
| (Eti)) 12.11.
(1 -
(X; (x; 1olle-1)) 2,00
max
hE[H]
+ 211 (X: 12.12220011 112,00 max he(H]
(wie - Wx (e.h) (
Wile
PT
2,00
+ La 1 1la max he[R]
-iven---ie (x;
2,00
This gives the required result.
C.3. Covering Number Upper Bound
Let x()), ( x ,..,,' (mm) E Rexd be a collection of inputs. Let FF the class of functions Raxd + as defined in Equation (10). Given > 0, let F be such that, for all f F there is f E C such that
- <e
max km(m)
Our goal is to bound the size of the cover set C. To do this we will inductively construct C. Let us start with simple bound on the norm of the output.
For W1:1 satisfying the norm bounds in Equation (10), using ||Pu||a < ||||a||o||a repeatedly, and the fact that o is Lg Lipsschhi with (0) = 0,
[x; 2,00 S Lo (1 12 (1
Σ
I T.g B3 1 + HB3 X 112,00
hEH
and inductively for each E [L] and lid
satisfying the norm bounds in Equation (10) we get
(x:
< <B T|,,,000 where B L, B2 (1 + HB3)
2,00
Let us now build a cover (()) for the first layer of the architecture. For any x(1, xx(2 (m) E Rxad satisfying
(xx 1 S Bx for all i, and W, W E W = (W e Rdxd', :IITTla, 112,, S Bw) because
(ww [ W)
max =E(m]
(W
W) T x
max (e[m]]E[aa
2,00
we may apply Lemma C.7 and get that there is an ew-cover W c pdxd' such that
log W
<
B B
log(dmn), and EW
max min max (W (W) T xx(1) WeW WeW We|m
|| 2,00 SEw.
Therefore, given inputs (1). (2)
(m) Raxd such that |x(' 12,00 I Bx for all i, using Lemma C.7 we get for
each h E [A] there is an (2) cover c(1.h) c Rddd with
B3, B
log (( S
log(dmn)
(1)
Vv
30

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
such that for all h E [M] and w:(l))) satisfying I|W 12,,
(IW(LA) (l,, S 12,1 there is W(
E c(1.)) such that
(W:(aA)- - w x(i
S /
max ve(mi
2,00
follows that there an Ev cover CV
(1)
1 Raxd with
HB2, B
S
log
c S
log(dmn)
h=1
such that for any w(1.1) w
wh( H)
satisfying ||
11 S B2,1 for all h, there is (WLL.))
,...,,.....) W(. E EC
such that
- x()
max max kE iE ml
S EV
2,00
(1) Similarly, there is an EQK cover
(1))
QK C II Rdoxdr with
log
((1) CQK
2
log(dpn)
(1)
2
EQK
such that for any W ((1.1) W (1.1) W Wn (1.2)
wi..H (1.8) satisfying [[
( (x (1,1))122 S B2,1 for
all h, there is (W (1.11)
Wx (1,1) ...., (ill.H Wx W. (m )
c(1)
Ckk such that
max kE(H]
(wilib (1,1)
-e
W W k (M)T
PT
SEO
(I)
2,,0
Given (W WR (1.1) Will Wx
(l) CQk and v
w E c.
we have that
(x; 11: B()BB Therefore, using Lemma C.7 there is an ew cover Ww C CddH++1xx with
(1)
cll)
B(12 201"B,,14 Bx
log c/ x
(1) log(dmmm)
eW,
such that for any w()) satisfying | 12,1 B2,1 there is wt (I)
Cw: such that
-w(" "0 (XI:
(1)
S
S
2,00
Similarly, given (WL Wy, ...., ". Wi (1HH)T) cllk (W)
cl
and w(
E c())
there is an Ww -cover
(1) (1) Cw2 E Rdddd with
B(1)² B2,
B3
log
c(1) W2
log(dmm)
(1)
EW2
such that for any ( satisfying 11,,1 B2,1 there is w (1) Cw3 such that
[[
- (w:")
(x" (,"
|
(1)
S
Ews
2,00
Therefore, using Lemma C.10, an cover (()) for the first layer of the architecture, where
(1) =
Ww + L.g B2 (1) + 2I, B3 Break (1)
+ VHL.B ())
31

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
has size
|c S |c
--@k max { [WW (Wi We (I.A)
.wh)|: (TT " - E c { willi (1.) (1a) H
(we
clll
S h=1
E QQk
max fow (WT-i)) We Willbi (1,A), Wh w : [wil-i)}
1
E cta (iil (1)
"
E coki Wl E cm
J 4=1
h=1
This means that
log (Ca)) S HB2, 21)
B
HB2, B3.
B(1)² B3, 1
B(I)² B3,11 B3
log(H dp)mm)).
(1)
(1)2
(1)
EQK
EW1
FW22
We now inductively construct a cover for deeper layers as follows. For each element W 1: E Cliff where Cl denotes a cover for the architecture up to (and including) layer on inputs _ X (m), we construct an 1 -cover C(++1) (((+1)(W:e) (Wi:t) on inputs { (x()) W
for the (E + 1) layer by following the same steps above, where ic(m)
+ +L,BB Be (f+1) + L.HLL (1+1) (++1) = Lo Br (1+1) Sw
It follows that from Lemma C.10 that
clif+1 := { W(0) Wht E c
w ,,,, E (T
is an ((f++)-covve of the architecture up to (and including) layer + 1, where
(e+1)) Lo B3e000 Ew,"
(e+1) L.g Br + 2L ( BpE (e+1) + v VH BH
(e+1)
((+1)
and
cLet1 S 1c2:9.
max (c((+1) (W.l))] WLiEECl
(
By induction, we obtain an (()) -cvvrr Cl.L of the function (X W1:L) on inputs X(1) ...., X(m) with
( = Σ ( (cW) g00 + +1000 (0) + + 79)
1=1
where
((t) L2B2, 110= 2L, Bl n(0) = VHL.
The size of the cover satisfies
HB B
1) B2,
B3.
B(0) B2, Bx (8) 2
B(O²
B0,, Bx
log 1l
log dpppdmmm mp)m)).
202
(0)
(e) 2
EV
EKK
EW
W2
To build the final cover C for the class of scalar-output architectures F, note that for any W WLL w. W
TO (x: (ml TT (F(D (x; (1bh))))) (X; < TO (p( -xX (X: [mi) + ( w) (x:
Therefore, we may apply Lemma C.7 and build C inductively for every element in CloL, and get that C is an e-cover for F with
E == Ew + B Σ (0) (=i0) + +/000 (2)
Ew + +
(6) EQK 17 (0(())
(18)
32

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
and that the size of C is bounded
B(L)²
B3 B
log |C|
x
log m log |cl.
(19)
To determine good upper bound on log C). we pick EV QQK E W, W, ae [4], and Ew that minimize the right-hand size of Equation (19) while satisfying Equation (18). Invoking Lemma C.8, this gives
(0)
(e)
(6) (0)
(BBW) B3 Bº, log logm)
1/3
log |C| x EΦ
+ BB B2,
dp)
Σ H (1) (
+ 11 B(t- / 1) 2/3
B a 20)2/33
+ B10)/3 2 + B(02/3 (0)2/3 g(0)2/3
HB, B3, 3 (1) (B(L)2 + +z008 + 17(19)
log H(d d + dr)mm)
1 B2 ( (1) 2
H H B3, B&B; BB
log (d dp)mm)
B B3, B B$
= (H OL))
log( (dd (d + mn).
(20)
2²
Letting ||w|2
B2, || 2,,,0
PPT
2,00 and using d 2 dr gives
(2 ((()) B2,
log S
log Hdmn).
(21)
22
Combining Equation (21) and Lemma C.5 proves Theorem 6.1.
Remark. Note that if we additionally apply layer normalization by projecting the output of each layer to the & unit ball as in Edelman et al. (2022), such that
(0) (X; (.W:i))Tll........ |1:00 T
112.00 B 1,
then the bound on log |C] in Equation (21) becomes
[og|c] (z) B2 log (H dmn).
(22)
D. Experiments
This section presents detailed results for the experiments reported in Section as well as additional results not present the main paper. Most reported results combine in-distribution (ID) and out-of-distribution (OOD) in one plot. This is shown by the scale factor on the x-axis, which indicates the level of out-of-distrriution. scale of one indicates in-distribution, whereas scale greater than one shows an increasing level of dut-o--dist
D.1 - Numeric tasks
D.I.I Experimental setting
D.1.2 length generalization experiments
D.1.3 Sample complexity experiments
D.1.4 Performance over variable input lengths
D.1.5 Performance of compact Transformers
33

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
D.1.6 Performance with alternative positional encodings
D.1.7 Ablation study on architectural design choices
D.1.8 Performance on OOD data with fine-tuned V projection matrices
D.1.9 Empirical results using Graph Neural Networks (GNNs)
D.1.10 Accuracy measures
D.2 k-hop induction heads
D.2.1 -Eperimmntal setting
D.2.2 Sample complexity experiment
D.2.3 - Exxressive power experiment
D.3 Mixeddede input task
D.3.1 Task description
D.3.2 Experimental setting
D.3.3 Sample complexity experiments
D.3.4 - Experiments with GPT2
D.1. Numeric tasks
D.1.1. EXPERIMENTAL SETTING
All numeric tasks employ the same model configuration. The model uses the structure of Equation (1), augmented with linear encoding and decoding layers.
We compare the standard Transformer, which utilizes the attention mechanism in Equation (2), and the positional Transformer, which employs the attention defined in Equation (3). Both variants share the same number of layers and dimensional configurations, with any specific differences explicitly noted.
In all configurations, the total number layers is set to [log2 n. + 1, where n denotes the maximum input length, and each layer uses 2 attention heads. Along with each input sequence, we also append an empty scratchpad entry. This extra entry does not count toward the total number of layers and not used to compute the loss. included solely to aid in the computation of the tasks. For the function ((e), we employ 2-layer MLP with ReLU activation functions. The embedding dimension of the encoder and the hidden dimensions of the MLP are both set to 64.
We use one-hot encoded vectors of dimension for positional encodings, where the non-zero entry corresponds to the node position. Consequently, the embedding dimensions of We and WK are set to n. A key difference between the models is that standard Transformers concatenate positional encodings to the input, whereas positional Transformers supply positional information exclusively through the matrix P. Therefore, positional Transformers, input values are solely encoded the input, and positional information is exclusively encoded in the positional encoding matrix.
Both models trained end-to-end using the squared loss between the predicted and target vectors of size n, with no intermediate supervision. We train models with Adam, starting with a learning rate of 5 10 and learning rate scheduler for total of 2000 epochs.
Our training data consists of samples from the range [-2, To ensure diversity in the data, for each input sample, we first select lower and upper bounds "/t and Yu uniformly in [-2, 2, 2]. and then for each the n elements of the input sample, we select its value uniformly from (11, 7a]-
1.2. INPUT LENGTH VS. GENERALIZATION EXPERIMENTS
In this section, validate that our results hold for multiple values of n. We train models for each fixed length n (2, 8, 16, 32} using 30,000 samples across all settings. The model depth varies with the input length with the number of layers set to [log, n + 1. We report the performance in-distribution and out-of-distribttion using 1,000 test samples. As illustrated across Figures 9 to 13, positional Transformers exhibit comparable in-distribution loss to standard Transformers across various sequence lengths. Naturally, for a fixed number of samples, the in-distribution loss slightly increases as the sequence length grows, indicating the need for more samples. For out-o--distributtin performance, we observe that the
34

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
positional Transformer maintains stable performance even with inputs of length -2. In contrast, the standard Transformer's performance rapidly degrades with the input length.
Train
In-distribation
Out-of-distribution
10-3
10-4
10
Loss (MSE)
10-4
10-6
10-3
10-1
10-6
10-1
10-3
10-7
10-7
10
20
30
10
20
30
10
20
Input Length
Positional
Standard
Figure 9: Training, in-distribution, and out-of-distributtion test performance for the summing task is shown for standard Transformers (red) and positional Transformers (blue) across different input lengths. The x-axis indicates the fixed input length on which the model was trained. Models are trained on the range [-2, 2] with 30,000 samples, validated on the same domain, and tested on an extended domain, [-6, 6]. each with 1,000 samples.
Train
In-distribation
Out-of-distribbttion
109
10-6
10-5
10
Loss (MSE)
10
10-6
10-6
10-3
10-7
10-4
10-7
10
20
30
10
20
30
10
20
30
Input Length
Positional
Standard
Figure 10: Training, in-distribution, and out-of out-of-distribution test performance for the minimum task is shown for standard Transformers (red) and positional Transformers (blue) across different input lengths. The x-axis indicates the fixed input length on which the model was trained. Models are trained on the range [-2, 2] with 30,000 samples, validated on the same domain, and tested on an extended domain, [-6, 6]. each with 1,000 samples.
Train
In-distribation
Out-of-distribution
10a
10-4
10
10-1
Loss (MSE)
10-0
10-6
10-2
10-0
10-6
10
10-7
10-7
10-4
10-8
10-8
10
20
30
10
20
30
10
20
30
Imput Length
Positional
Standard
Figure 11: Training, in-distribution, and out-of-distributian test performance for the median task is shown for standard Transformers (red) and positional Transformers (blue) across different input lengths. The x-axis indicates the fixed input length on which the model was trained. Models are trained on the range [-2, 2] with 30,000 samples, validated on the same domain, and tested on an extended domain, [-6, 6), each with 1,000 samples.
35

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Train
In-distribution
Out-of-distribution
10
10
10
10a
10
10
Loss (MSE)
10-1
10
10
10-2
10
10-9
10-3
10
10-7
10
20
30
10
20
30
10
20
30
Imput Length
Positional
Standard
Figure 12: Training, in-distribution, and out-of-distribution test performance for the sorting task is shown for standard Transformers (red) and positional Transformers (blue) across different input lengths. The x-axis indicates the fixed input length on which the model was trained. Models are trained on the range [-2, 2] with 30,000 samples, validated on the same domain, and tested on an extended domain, [-6, 6]. each with 1,000 samples.
Train
In-distribution
10a
10-3
10-3
Loss (MSE)
10
10
10a
10
10
10
10
10-6
10-2
10-3
10
10
10
20
30
10
20
30
10
20
30
Input Length
Positional
Standard
Figure 13: Training, in-distribution, and out-of-distribution test performance for the maximum sum subarray task is shown for standard Transformers (red) and positional Transformers (blue) across different input lengths. The x-axis indicates the fixed input length on which the model was trained. Models are trained on range [-2, 2] with 30,000 samples, validated on the same domain, and tested on an extended domain, [-6, 6]. each with 1,000 samples.
D. 1.3.3. SAMPLE SIZE VS. GENERALIZZTION EXPERIMENTS
In this section, we provide detailed results showcasing the training, in-distribution, and out-of-distribution test performance for each of the five numerical tasks as function of the number of training samples used. From the results, we can draw two conclusions about the behavior of the models as the number of samples increases. First, both models achieve good in-distribution performance. Second, only the positional Transformer achieves better OOD performance. The results for this experiment are presented in Figures 14 18.
36

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Train
In-distribution
Out-of-distribution
10
10-3
10"
Loss (MSE)
10
10-4
10-1
10-5
10-5
10-2
10-3
10
10-9
20000 40000
20000
40000
20000
40000
Number of Samples
Positional
Standard
Figure 14: Training, in-distribution, and out-of-distribution test performance for the summing task is shown for standard Transformers (red) and positional Transformers (blue) function of the number of training samples (indicated on the x-axis). Models are trained on the range [-2, 2] with varying training set sizes. In-distribution testing is performed on the same domain, and testing is conducted on an extended domain, [-6, 6]. each using 1,000 samples.
Train
In-distribution
Out-of-distribution
10a
10
10
10-1
Loss (MSE)
10
10-5
10-2
10-9
10-3
10-6
20000
40000
20000
40000
20000
40000
Number of Samples
Positional
Standard
Figure 15: Training, in-distribution, and out-of-diittibution test performance for the minimum task is shown for standard Transformers (red) and positional Transformers (blue) as a function the number of training samples (indicated on the x-axis). Models are trained on the range [-2, with varying training set sizes. In-distribution testing is performed on the same domain, and testing is conducted extended domain, [-6, 6], each using 1,000 samples.
Train
In-distribution
Ont-of-distrribttion
10a
10-3
Loss (MSE)
10-1
10
10
10-2
10-5
10-3
10-5
20000
40000
20000
40000
20000
40000
Number of Samples
Positional
Standard
Figure 16: Training, in-distribution, and ditt---diitributionn test performance for the median task is shown for standard Transformers (red) and positional Transformers (blue) as a function the number of training samples (indicated on the x-axis). Models are trained on the range [-2, 2] with varying training set sizes. In-distribution testing is performed on the same domain, and testing is conducted extended domain, [-6, 6), each using 1,000 samples.
37

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Train
In-distribution
Out-of-distribution
10"
10-3
10-3
Loss (MSE)
10-1
10-2
10
10-4
10-3
20000
40000
20000
40000
20000
40000
Number of Samples
Positional
Standard
Figure 17: Training, in-distribution, and out-of-distribttion test performance for the sorting task is shown for standard Transformers (red) and positional Transformers (blue) as function of the number of training samples (indicated on the x-axis). Models are trained on the range [-2, 2] with varying training set sizes. In-distribution testing is performed on the same domain, and out-of-distribution testing is conducted on an extended domain, [-6, 6]. each using 1,000 samples.
Train
In-distribution
10
10-3
10
10
Loss (MSE)
10-1
10-1
10-4
10-2
10-3
20000
40000
20000
40000
20000
40000
Number of Samples
Positional
Standard
Figure 18: Training, in-distribution, and out-of-distribution test performance for the maximum sum subarray task is shown for standard Transformers (red) and positional Transformers (blue) as a function of the number training samples (indicated on the x-axis). Models are trained on the range [-2, 2] with varying training set sizes. In-distribution testing is performed on the same domain, and out-of-distribution testing is conducted an extended domain, [-6, 6). each using 1,000 samples.
1.1.4. PERFORMANCE OVER VARIABLE INPUT LENGTHS
In this section, we present generalization results for models operating on variable-length inputs. This setting aims to verify the models' ability to generalize across different scales while maintaining the flexibility to handle inputs of varying lengths.
We evaluate the models' ability to process sequences of varying lengths up to a maximum size of n = 8 Specifically, the Due to the more challenging setting,
model is required to perform tasks on input sequences with lengths ranging from we train models with 500,000 samples and ensure that all input lengths are equally represented.
We then evaluate the in-distribution (ID) and out-of-distribune Note that when 1, the setting actually corresponds to ID generalization. The losses reported are calculated using 3,000 samples for each scale. As shown in Figure 19, positional Transformers maintain the same in-distribution performance while consistently outperforming standard Transformers across all distrribibt scales and tasks. Additionally, our architecture maintains robust OOD performance even in tasks where the output can exceed the input magnitude (c.g., sum and maximum sum subarray).
(OOD) loss across different scale factors E 21, 10).
38

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Sum
Min
Modian
Max smit subaray
10
11P
10a
10
10F
10P
10
Loss (MSE)
100
10
10
11F
10
10
10
10F
10
10
10
10
10
10
10
10
10
2
9 10
4
5
18
6 Scale
10
8
10
3
5
Positional
Standard
Figure 19: ID/OOD loss (measured mean squared error, MSE) for standard Transformers (red) and positional Transformers (blue) across all five tasks for variable lengths (up to 8). The x-axis represents the ID/OOD scale factor. The solid line and shaded area denote the median and the region between the 10th and 90th percentiles over ten trials, respectively.
1.1.5. PERFORMANCE OF COMPACT TRANSFORMERS
This section presents additional generalization results for simpler models, aiming to rule out potential overfitting caused by the excessive complexity of the standard Transformer. We examine two configuration variants: one with log + layers (4 layers) and another with single layer. The plots also illustrate the outcomes for different hidden dimensions in the MLP. We report the generalization results for the cumulative sum (Figure 20) and cumulative minimum (Figure 21) tasks. As observed, in both cases, the OOD performance deteriorates as the network size decreases. Although single-layer networks exhibit slightly better performance, they remain inferior to the performance of positional attention reported in the main paper.
10a
10a
10
Loss (MSE)
Loss (MSE)
109
10-1
10
neurons in MLP
10-3
neurons in MLP
neurons in MLP
neurons in MLP
neuuoos in MLP
10-5
8 peurons in MLP
16 neurons in MIP
16 neurons in MLP
3
10
2
3
9
10
Scale
Scale
Figure 20: ID/OOD loss for various standard Transformer models the cumulative sum task with fixed length (n = 8). The left plot displays results for models with 4 layers, while the right plot shows results for single-layer models, both featuring varying hidden dimensions in the MIPs. The x-axis represents the out-of-distrihie from the training distribution. The solid lines and shaded areas denote the median and the regions between the 1h and 90th percentiles across ten trials, respectively.
scale factor, indicating the distance
D.1.6. PERFORMANCE WITH ALTERNATIVE POSITIONAL ENCODINGS
We compare the performance of positional Transformers and standard Transformers using alternative positional encodings, such as binary and sinusoidal encodings (Vaswani et al., 2017) and Rotary Positional Embedding (RoPE) (Su et al., 2024a), a widely adopted technique natural language processing contexts that has also been applied algorithmic tasks (Bounsi al., 2024).
For binary positional encodings, use [log. n dimensions, where each entry represents the binary encoding of its index in [log, n.] bits, with zeros encoded as -1. The result for binary encodings is shown in Figure 22. For sinusoidal positional encodings, we follow the strategy outlined in (Vaswani et al., 2017), with the encoding dimension set to [>/2]. The result for sinusoidal encodings is shown in Figure 23. From an expressivity perspective, while these encodings are less expressive than one-hot positional encodings. they maintain consistent out-o--diittibbutin (OOD) performance across all ranges. Furthermore, positional Transformers outperform standard Transformers in every task tested.
39

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
10H
10
Loss (MSE)
Loss (MSE)
10
10-2
neurons in MLP
neurons in MI.P
neurons in MLP
neurons in MLP
10-5
neurons in MLP
neurons in MLP
16 neurons in MLP
16 neurons in MLP
3
10
2
3
9 10
Scale
Scale
Figure 21: ID/OOD loss for various standard Transformer models on the cumulative minimum task with fixed length (n 8). The left plot displays results for models with 4 layers, while the right plot shows results for single-layer models, both featuring varying hidden dimensions in the MLPs. The x-axis represents the the distance from the training distribution. The solid lines and shaded areas denote the median and the regions between the 10th and thth percentiles across ten trials, respectively.
scale factor, indicating
As for RoPE, while manages decrease the OOD test loss, this improvement is far from what is observed in positional Transformers across every task. The results of this experiment are presented in Figure 24.
Sum
Min
Median
Sort
Max sum su subarray
Loss (MSE)
10
10
10
107
10F
109
10
107
109
102
11
10
10-2
10
10
10
10
10
10m1
10
10
Scale
Positional
(Bin. PE)
Standard (Bin. PE)
Figure 22: ID/OOD loss (measured as mean squared error, MSE) for standard Transformers (red) and positional Transformers (blue) using binary positional encodings across all five tasks for fixed length (n = 8). The x-axis represents the OOD scale factor, indicating the distance from the training distribution. The solid line and shaded area denote the median and the region between the 1hth and 9hth percentiles over ten trials, respectively.
Sum
Min
Median
Sort
Max sum su subarra
10"r
10:
10"
10a
Loss (MSE)
10
10
11k
10
10P)
11F
10
10-
20
10
10
10
10-2
10-2
10
10-3
10
10
33"
Scale
Positional (Sim. PE)
Standard (Sin. PE)
Figure 23: ID/OOD loss (measured as mean squared error, MSE) for standard Transformers (red) and positional Transformers (blue) using sinusoidal positional encodings across all five tasks for fixed length (n = 8). The x-axis represents the OOD scale factor, indicating the distance from the training distribution. The solid line and shaded area denote the median and the region between the 1h and th percentiles over ten trials, respectively.
40

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Sum
Min
Modian
Max smit subaray
10F
10
10
10P
108
Loss (MSE)
10
10
10
10
10
10
10
10
10
11
10
10
2
10
10
10
10
2
9 10
4
5
88
9
10
6 Seale
8
16
3
4
5
6
8
3 5
10
Pasitional
Standard (RaPE)
Figure 24: ID/OOD loss (measured mean squared error, MSE) for standard Transformers with RoPE (red) and positional Transformers (blue) across all five tasks for fixed length (n 8). The x-axis represents the OOD scale factor, indicating the distance from the training distribution. The solid line and shaded area denote the median and the region between the 10th and 90th percentiles over ten trials, respectively.
D.1.7. ABLATION STUDY ON ARCHITECTURAL DESIGN CHOICES
We carried out additional experiments to explore variations in the input data format and the placement of positional encodings. This leads the following seven architectural choices (including two from the main paper and five additional variations):
1. Standard Transformers: Input numbers and positional encodings fed MLPs, value, query, and key matrices.
2. Standard Transformers with positional input (but no positional encodings): Input numbers are placed in one-hot positions, that is. the input is Diag(X) where X is the list of input numbers we give to other architectures. No additional positional encodings are used.
3. Positional Transformers: Input numbers are fed to the MLPs and value matrix; positional encodings are fed to query and key matrices.
4. Misaligned Positional Transformers: Input numbers are fed the MLPs and value matrix; positional encodings are fed to the MLPs, value, query, and key matrices. That is, compared with Positional Transformers, add positional encodings to the input.
5. regulazegul Standard Transformers: Input numbers are fed to the MLPs, value, query, and key matrices; positional encodings are only used in query and key matrices.
6. No Positional Encodings: input numbers are fed to the MLPs, value, query, and key matrices. No positional encodings are used.
7. Using RoPE Only: Input numbers are fed to MLPs, value, query, and key matrices, removing absolute positional encodings, and using only RoPE in standard transformers.
We test the architectures for the cumulative task and the sorting task representative tasks, for fixed input length 8. We keep the train/valid/test setup the same as before. The results are shown in Figure 25 and Figure 26. We note that our proposed positional Transformer architecture has all of the following three important factors that enabled OOD generalization: (1) use positional encodings, (2) do not use positional encodings in the value matrix of the attention, (3) use only fixed positional encodings in the key and query matrices. This design principle aligns with algorithms that are typically used solve algorithmic tasks.
1. EXPERIMENTS ON HINE-TUNING V PROJECTION MATRICES ONLY
Fine-tuning pre- trained models is very common in practice. In this case, takes model that has been trained different (and usually larger) dataset, and re-trains small subset of layers new dataset, while keeping all other model parameters frozen. To further study the potential advantage of positional attention, we consider the following fine-tuning experiment. We take both positional Transformer and standard Transformer models that have been trained over in-distribution data (cf. Section 7.2 where c = 1, and Appendix 1. 1) for the numeric algorithmic tasks, fix all their model parameters except the weight parameters the projection matrices, and fine-tune the V projection matrices over OOD data with 10x scale factor

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
standard, input X, self-attention
10a
standard, input X, self-attention + RoPE standard, input X, hybrid-attention XIP standard, input X|P, self-attention standard, input (iag(X), self-attention positional, input X/P positional, input X
Loss (MSE)
10a
10
10-3
10-5
2
3
4
6
7
9
10
Scale
Figure 25: ID/OOD loss for the cumulative sum task under various architectural choices. We fix input length = 8. The x-axis represents the OOD scale factor. The solid line and shaded area denote the median and the region between the 10th and 90th percentiles for 10 trials, respectively.
10a
standard, input X, self-attention
standard, input X, self-attention RoPE standard, input X, hybrid-attention X/P standard, input X/P, self-attention standard, input (X), self-attention positional, input X/P positional, input X
102
Loss (MSE)
10
10-2
10-4
2 3
4
5
6
7
8
9
10
Scale
Figure 26: ID/OOD loss for the sorting task under various architectural choices. We fix input length n = 8. The x-axis represents the OOD scale factor. The solid line and shaded area denote the median and region between 10 and 90th percentiles for 10 trials, respectively.
(cf. Section 7.2 where c 10). We consider three algorithmic tasks: cumulative sum, cumulative minimum, and sorting. For these tasks, find that fine-tuning only the V projection matrices in the OOD allows the positional Transformer achieve the same level of generalization as in the in-distribution case, while, the other hand, the same does not hold for the standard Transformer.
Table reports the median test error (MSE) for both positional and standard transformer models when (i) they are trained and tested on in-distribution data (in-dist), (ii) they are trained on in-distribution data but tested on numbers 10x larger than the original in-distribution data (10x OOD), (iii) they are trained on in-distribution data, fine-tuned V projection matrices on numbers 10x larger than the original in in-distribution data, and then tested on numbers 10x larger than the original in-distribution data (10x OOD after fine-tuning). We use the same empirical setting as before, where we train both models for 2000 epochs 20 epochs. We normalize the MSE by the scale factor so that the results indicate relative accuracy. The results in Table indicate that fine-tuning only the V projection matrices brings the accuracy of the positional Transformer to the same level as the in-distribution case. However, for the standard Transformer, even after fine-tuning the projection matrices, the accuracy on 10x test data is still very high compared to that the positional Transformer. The comparison highlights another potential advantage of the positional Transformer for executing algorithmic tasks: one might efficiently fine-tune
in-distribution data. For fine-tuning,
adopt the common practice and train for only
42

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
pre-trained positional Transformer on OOD data and achieve good performance.
Table I: Performance (MSE test loss) of fine-tuned positional and standard transformer models
Task
Architecture
in-dist
10x OOD
10x OOD after fine-tuning
Positional
6.07e-06
1.31e-03
1.32e-05
Cumulative Sum
Standard
5.20e-06
1.09e+00
1.23e-01
Positional
1.03e-05
3.92e-04
2.19e-05
Cumulative Minimum
Standard
1.74e-06
1.39c-01
4.00e-02
Positional
1.20e-04
7.37c-04
1.23e-04
Sorting
Standard
9.05e-04
5.18c-01
6.90e-02
1.1.9. EMPIRICAL RESULTS USING GRAPH NEURAL NETWORKS (GNNs)
Graph Neural Networks are popular choice for solving algorithmic tasks on graphs. We tested the performance of Graph Convolutional Network (GCN) and Graph Attention Network (GAT) in solving the cumulative sum and sorting tasks. Since the tasks tested have no underlying native graph, we tested these models on complete and star graphs. Notably, the original GAT' architecture on a complete graph is similar to a standard transformer but differs in that the value, key, and query weights are shared in GAT. In Table we present results reporting the median MSE loss. As the results indicate, neither GCN nor GAT works very well even for in-distribution data (OOD scaling factor = 1). let alone achieving OOD generalization. We believe this is because these tasks inherently unsuitable for standard message-passing architectures.
Table 2: ID/OOD loss of Graph Convolutional Network (GCN) and Graph Attention Network (GAT)
OOD scale factor
Task
Graph
Architecture
GCN
3.9c+0
2,1c+1
4.3e+1
7.4c+1
1.2c+2
4.8c+2
Complete
GAT
3.7e+0
2.0e+1
4.3e+1
7,5e+1
1.2e+2
5.4e+2
Cumulative Sum
GCN
4.0e+0
2,0e+1
4.2e+1
7.4e+1
1.2e+2
4.6e+2
Star
GAT
3.8e+0
2.0e+1
4.1e+1
7.5e+1
1.2e+2
5.0e+2
GCN
1.9e-1
9.8e-1
2.2e40
3.9e+0
6.1e+0
2.7e+1
Complete
GAT
1.9e-1
9.6e-1
2.1e+0
3.7e+0
5.6e+0
2.4e+1
Sorting
GCN
1.9e-1
1.0e40
2.1e+0
3.7e+0
6.1e+0
2.6e+1
Star
GAT
1.9e-1
9.9c-1
2.0e+0
3.5e+0
5.6e+0
2.3e+1
D.1.10. ACCURACY MEASURES
Given the regressive nature of the tasks considered in this work, the notion of model accuracy is not properly defined. However, we propose the following transformation strategies that assign binary labels ("cooredd"/inncoreeet)) to the models" outputs allowing us to measure their accuracy (with respect to these transformations):
Rounding transformation: We evaluate the model on lists containing integers (while training is still done using real numbers) and round the model's 's output to the nearest integer (or nearest 0.5 for the median task). A prediction is considered "correct" if the rounded and ground truth lists are the same. The generalization accuracies for all arithmetic tasks in the main paper using this transformation strategy are presented Figure 27.
Closeness transformation: We evaluate the model on lists of real numbers, considering a prediction "correct" if each entry in the predicted list is within an absolute precision of 0.05 and a relative precision of 5% compared to the corresponding entry of the ground truth list. The generalization accuracies for all arithmetic tasks in the main paper using this transformation strategy are presented in Figure 28.
is important to note that the above metrics are quite unforgiving, as even a single element in the predicted list failing to meet the corresponding criterion results in the entire list being classified as "incorrect". It is therefore expected that
43

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
increasing the OOD scale factor causes the model's accuracy decrease rapidly. Nevertheless, the positional Transformer significantly outperforms the standard Transformer in terms of accuracy in all five tasks.
Sum
Min
Median
Sort
Max sum subarray
10053
10006
1005
10053
50%
8053
Accurey
60%
80%
6076
6095
10%6
4058
20%6
27%
2076
2055
0996
070
0%
098
Scale
Positional
Standard
Figure 27: OOD accuracy when using "rounding transformation" across all five tasks for standard Transformers (red) and positional Transformers (blue) as a function of the scale factor. The solid line and shaded area denote the median and the region between the 10th and 90th percentiles over ten trials, respectively.
Sum
Min
Median
Sort
Max sum subarray
10053
10006
100%
10055
8095
8076
80%
805
Accuracy
60%
80%
01%
6056
4098
20%
40%
4050
2015
20%
20%
2055
099
070
050
Scale
Positional
Standard
Figure 28: OOD accuracy when using "closeness transformation" across all five tasks for standard Transformers (red) and positional Transformers (blue) as a function of the scale factor. The solid line and shaded area denote the median and the region between the 10th and 90th percentiles over ten trials, respectively.
D.2. k-hop induction heads
D.2.1. EXPERIMENTAL SETTING
We adopt a slight modification of the configuration used in the numeric task for the k-hop induction heads task of Sanford et al. (2024). Specifically, since the input consists of tokens, we use an embedding layer as the encoding layer, with an embedding dimension of 60. The decoding step remains a linear layer to produce the logits, whose output dimension is set to the vocabulary size. Additionally, due to the sequential nature of the task, we apply a causal mask over the attention coefficients in both architectures.
The dataset creation is adopted from Sanford et al. (2024) and consists of the following types of characters: numbers from 0 to h: which indicate the number of hops required for each sample (as mentioned in Section 7, we set h: = for our experiments). The sequences themselves are constructed from specific character sets: {a, b, c, d, e} for in-distribution and (f, , g, h, i, j} for -itoof--iitribution.. An additional character (~) denotes cases where a k-hop character cannot be found. To ensure that most output tokens can be reached within the specified number of hops (thus avoiding ~ as an output), we set the sequence length 30.
D.2.2. SAMPLE COMPLEXITY EXPERIMENT
In this section, provide detailed results showcasing the training, in-distribution, and cutt---diitribution test performance for the k-hop induction heads task function of the number of training samples used. From the results on Figure 29, we can observe sharp contrast between the Icarnability of the standard and positional Transformers, indicating that, for this particular task, the sample complexity of positional Transformers is higher than that of standard Transformers. As discussed in Section 7, this is likely a result of the nature of the induction heads task, which is purposely designed to showease the relational capabilities of self-attention. As a result, positional Transformers require more data to replicate this same behavior.
44

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Train
In In-distribttion
Out of dddditrrbution
0.125
Loss (Cruss Eniropy)
1.5
0.100
0.075
40
0.050
0.5
0.025
20
0.000
0.0
10000
200001
30000
40000
50000
30000
Samples
Positional
Standard
Figure 29: Training, in-distribution, and out-of-distribution test performance for the k-hop induction heads task are shown for standard Transformers (red) and positional Transformers (blue) as a function of the number training samples (indicated on the x-axis).
D.2.3. EXPRESSIVE POWER EXPERIMENT
In this section, we provide detailed results showcasing the training, in-distribution, and out-of-distribution test performance for the h-hop induction heads task as function of the expressive power of the architectures, measured by the number of layers. From the results presented in Figure 30, we can conclude that, for this particular task, positional Transformers require many more layers to achieve a similar performance as standard Transformers, which empirically validates our findings established in Remark 5.2.
Train
In-distribution
Out-of-distribution
0.6
40
Loss (Cross Emtropy)
0.4
30
0.2
0.2
0.0
0.0
Layers
Positional
Standard
Figure 30: Training, in-distribution, and out-of-distribution test performance for the k-hop induction heads task are shown for standard Transformers (red) and positional Transformers (blue) as function of the number of layers (indicated on the x-axis).
D.3. Mixed-type input task
In this section, we provide details on the task Section 7.3. Specifically, we give formal description of the task, provide details experimental setting, and present additional experiments using fine-tuned version of GPT2-large from HuggingFace (Radford et al., 2019). Lastly, we present additional experiments with varying training sample size for both standard and positional Transformer.
D.3.1. TASK DESCRIPTION
Let n, kt m E with n 2 le and n 2 m. The input to the model is a list consisting of text and numerical values of the following form:
45

---

augmented by adding more alphanumeric elements at random positions in the list which we call irrelevant structure. Here, i,,,,, E N, ........ Vi, E R2o. ....... Jk E {ir,.... i, and (type) } is one of 'min', 'max', ' 'sum'. The answer to the query is either the minimum, maximum, or sum of the set [. U/e } depending on [type]. The irrelevant Positional Attention: Expressivity and Learnability of Algorithmic Computation
I'Catii', , u, 'Catig', Vizr CCatin', Vaar 'Find ftype) of catji, Catfa, Catjk-1 and Catjs'
structure is generated by concatenating the string Cat with one of the characters in {' category identifier. Notice how this setting allows for one model to potentially process multiple query types. In fact, we present one such experiment later.
' 'x') } and a
OOD generalization For this task, we measure OOD generalization in three ways (simultaneously):
1. We test using category identifiers (i.e i ,,) that the models haven't 't encountered during training, and . Un used for testing is larger than the range used for training.
2. The range of values v .2.....
3. The type of irrelevant structure used for testing is different than the type used for training. We detail this difference below.
D.3.2. EXPERIMENTAL SETTING
We fix n = 8, k and m = 2. Both standard and positional Transformers consist of 3 layers, two attention heads per layer, and an embedding dimension of 32. All characters of the non-numeric parts of the input are tokenized and passed through embedding layer, while the numeric ones are passed through a linear layer. The category identifiers i,, 2...... is sampled randomly from {1, 2, 20} for training and {21, 22, ... 40} for testing. The query category identifiers then sampled randomly from {81, 111,,,,...,,4 is). The values : ..... Big are sampled using the technique of Section 7 from [0, 5] for training and from [0, 5c] where e E {1, 2, ......,00.. 10} is the scaling factor for testing. We also apply the rejection step of Section 7.2 when generating testing samples. Furthermore, when sampling a training list we form irrelevant structure by concatenating the string Cat' with either one of +' or (chosen at random) and a category identifier that is present in the actual input. For testing, irrelevant structure is formed similarly by concatenating the string Cat x with either one of or (chosen at random) and category identifier that is present in the actual (test) input. In both cases, irrelevant structure is injected into the "truc" list at random positions. We run experiments for the following three variants of the task:
1. One where training and testing prompts consist exclusively of prompts where {type} is 'min'.
2. One where training and testing prompts consist exclusively of prompts where [type} is ' maxx.
3. One where training and testing prompts consist exclusively of prompts where [type} is either 'min' or max'. We refer to this experiment "multitasking" since it allows single trained model to process different query types. In fact, the choice of minimum and maximum as the two query types is, in some sense, an "extreme" case, given the "opposite" nature of minimum and maximum operations.
Both positional and standard Transformer models are trained using Adam for 2000 epochs with a sample size of 50,000, a batch size of 1024, and initial learning rate of 5 10-4 which is controlled by learning rate scheduler.
D.3.3. SAMPLE COMPLEXITY EXPERIMENTS
In Figures 31 to 33 we present the in-distribution and out-of-distribution performance of standard and positional Transformer on the three variations of the mixed-type input task. The conclusions are similar to those of Appendix D.1.3, namely, both models fit the training data but positional Transformer achieves significantly better out-of-distribution performance.
D.3.4. EXPERIMENTS WITH GPT2
For completeness, and given the nature of the task, we also perform experiments with a fined-tuned version of GPT2-large from HuggingFace (Radford al., 2019) on the above task. For all three variations (min, max, and multitasking), we fix the
"The ellipses are not part the prompt
46

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Train
In-distribution
Out-of-distribution
2 10
10
2.25 x 10
10
2x 1
Loss (MSE)
1 1.
10
6
1 10-9;
1.25 x 10
10
4x 10-
10-2
10a
25000 50000
75000
100000
25000
50000
75000
100000
25000
50000
75000
100000
Number of Samples
Positional
Standurd
Figure 31: Training, in-distribution, and out-of-distribution test performance for the minimum variation of the mixed-type input task is shown for standard Transformers (red) and positional Transformers (blue) function of the number of training samples (indicated on the x-axis). Models are trained on category values in the range [0, 5] with varying training set sizes. In-distribution testing is performed on the same domain, and -uttt---dittiuttion testing is conducted on an extended domain of category values, (0, 50), each using 1,000 samples.
Train
In-distributtio
Out-of-distribattio
10-1
16x 10-1
10
Loss (MSE)
10
14x 10-1
12 10
10a
8x 10-4
10-1
10
25000 50000
75000
100000
25000
50000
75000
100000
25000
50000
75000
100000
Number of Samples
Positional
Standard
Figure 32: Training, in-distribution, and nut-of-distribution test performance for the sum variation of the mixed-type input task is shown for standard Transformers (red) and positional Transformers (blue) as function of the number of training samples (indicated on the x-axis). Models are trained on category values in the range [0, 5] with varying training set sizes. In-distribution testing is performed on the same domain, and out-of-distribttion testing is conducted on an extended domain of category values, [0, 50), each using 1,000 samples.
precision of the input numbers digits (2 digits for the whole part and 2 digits after the decimal point) and the precision of the output 4 digits for min and max (2 digits for the whole part and 2 digits after the decimal point) and digits for sum (3 digits for the whole part and 2 digits after the decimal point). This covers all possible numbers that can be sampled. We then use byte pair encoding for tokenization and fine-tuned using 50,000 training samples for 3 epochs. Since this setting is autoregressive (next-token prediction), there were cases where the model's output did not correspond to a real number. As this did not occur frequently enough to be considered a problem, we report the median OOD losses (MSE and MAPE) ignoring samples for which no numeric value could be extracted from the model's output. The results of this experiment are presented in Figure 34. For easier comparison, we also include the losses of the positional Transformer on the same tasks.
E. Probability of generating OOD test data in our empirical setting
Recall that we sample the training and test data in the following way. The training data consists of i.i.d samples whose values are drawn from the range 2-2, 2]. To ensure diversity, for each training sample, we first select lower and upper bounds n and Yo uniformly in 2-2, 2]. and then for each of the n elements of the training sample, we select its value uniformly from the interval [7a, T. We employ a similar sampling strategy for testing but extend the value range to [-2c, 2c). where
47

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Train
In-distribution
Out-of-distribttion
10a
10-2
x 10-
Loss (MSE)
10
6 10
10
10a
25000
50000
75000
100000
25000
50000
75000
100000
25000
50000
75000
100000
Number of Samples
Positional
Standurd
Figure 33: Training, in-distribution,, and out-of-distribution test performance for the multitask variation of the mixed-type input task is shown for standard Transformers (red) and positional Transformers (blue) function of the number of training samples (indicated on the x-axis). Models are trained on category values in the range (0, 5] with varying training set sizes. In-distribution is performed on the same domain, and out- of-distbibutiin testing is conducted on an extended domain of category values, [0, 50), each using 1,000 samples.
<> 1 is the OOD scale factor. Additionally, during the test sampling process, we apply a rejection step to ensure that either
< - or > 2, while maintaining -c S 73 S Yv < 2c.
We will compute the probability that a randomly sampled test instance x E IR" lies in the domain of the training distribution, i.e., we will compute (P)) (xr E [-2, ). In particular, we will show that this probability is proportional to 1 1/nc². Consequently, in our experiments, the majority of the test data lie outside of the domain of the training distribution.
Without loss of generality, let us assume that the training data are sampled within the interval [-1, 1] and the test data sampled within the interval [-c, c), where e is the OOD scale factor. Note that this does not affect the probability that we want to compute. In the test sampling process, when we sample two uniform numbers 't and To from [-c, c|. exactly one of the following 3 disjoint events can happen.
Event A. Exactly one of "e and 7h lies in [-1, This happens with probability 2(==))
Event B. Neither nor is inside the interval [-1, 1]. This happens with probability (z=1)²
Event C. Both "/ and %o are inside the interval [-1, 1]. This happens with probability
Our rejection step rejects the samples generated under Event C. Therefore, in our setting when we sample pair of and Tu in order to generate single instance of the test list, we have that
((vvent A/Rejecting Event C) 2(c 1) -
P(vvent B/Rejecting Event C) (c 1)2 1
(23)
We analyze the probability of generating OOD test data under each event. First, let us suppose that Event A happens when we sample "/ and Tu. This means that either "/ E [-c, -1) Tu E (1,c) (but not both). More precisely, the following two sub-events partition Event A:
Event A.I. N/ E [[1, and 7a E (1,c). Given that Event A happens, this sub-event happens with probability 1/2.
Event A.2. E [-c, -1) and 7a [-1,1]. 1]. Given that Event A happens, this sub-event happens with probability 1/2.
By symmetry of the probability distributions, the probability that we wish to compute remains the same under both of the above sub-events. Therefore, let us focus on Event A. Suppose that Event A.1 happens, i.e., 74 E [[1, 1] and 7a E (1,c).
48

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Min
Sum
Multitask
10
10
10
Loss (MSE)
10
10
10
10
10-
10-h
10
10
2
Scale
Positional
GPT2 GPT2-large lange
Min
Sum
Multitask
25%
Loss (MAPE)
40%
70%
61%
30%
50%
13%
20%
40%
10%
10%
20%
5%
1006
0%
096
0%
0 Scale
3
Posicional
GPT2-large
Figure 34: This experiment shows ID/OOD losses (measured as mean squared error, MSE, and mean absolute percentage error, MAPE) for fine-tuned GPT2-large (red) and positional Transformers (blue) across all three variations of the mixed-type input task. The x-axis represents the OOD scale factor. The solid line and shaded area denote the median and the region between the 10th and 90th percentiles over ten trials, respectively.
Conditioning on this event, we know that is uniform on [-1, 1] and 7u is uniform on (1, The probability density functions for "ie and are
(s) - 1 I[-I58s11, (.(1) c-1 <(<()5)).
Let X RR be random vector whose ith coordinate X₁ is independently and uniformly sampled from the interval [Ye, Yul. then the conditional probability density function for X given Ye, is
(2 Tn 8,7u = 1 " I(s 5 S.
Therefore, the joint density function is
1 1
fxcrette (Tt, Ini 8,
(
((-15:51,11155 viV)).
2 (c - 1)
follows that
Xi E 11, vi e [m] Event A)
PA,in = Px
1
dx dt ds
se([-1.1] (rE(L.)) Jne(s.1]n
(c 1)
1 1 8
dt ds
2 (c -
-0[--1,1] 2E(1,)
t-s
1 1 1
(c 1) (n (n-1) /
1
(1= 1
ds
1
(1-8)
ds
(c 1) (n 11 1
49

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
n-1
8
+
(1 (1-8)
1
ds
SE|0,1]
n-1
(1-s)
ds + sE|0,11
(1 - s)ds
2 ( - 1) 1)
sE|-1,00
- 11 +1 4(n ()(cc_
(24)
This is the probability that, under Event A, a randomly generated test sample lies within the domain of the training distribution. Again, recall that by scaling down the domain of the test distribution to [-c, c] accordingly, we have assumed that the domain of the training distribution is [-1, 1]" without loss of generalization.
Now suppose that Event B happens. this case both T/e and Yo are uniformly distributed over [-c, -1) (1,c). The following two sub-events partition Event B:
Event B.I. Either both 7o or both ne, -1. Given that Event B happens, this sub-event happens with probability 1/2.
Event B.2. Ni < -1 and 7a > 1. Given that Event B happens, this sub-event happens with probability 1/2.
Let X E IK" be a random vector whose ith coordinate X; is independently and uniformly sampled from the interval Yu]. Note that under Event I.I, one always has that X & [-1,,]",,.. i.e.
18. == Px (x E [-1, vi
[n] Event B.1) = 0.
(25)
Therefore let us consider Event B.2. Conditioning this event, we know that ne is uniform on [-c, -1) and 7a is uniform on c1. The joint density function (conditional on Event B.2) for X, 7e, Tu is
1 fxcerr (I1...., Tn. S,t) (c- (c-1)² 1)2 (t
n II I
Therefore, have that
PB.2, PX
E [-1, E [n]] Event B.2)
= -El--.1) =/8--a
/ SSCCLO JDEla1ja
1 n
da dt ds
(c- (c-1)2 1)2
t t-s
1
1
2"
dt ds
(c 1)2
xc([-c,1 [E(1I
1-8
2
(1 - s)n-1 (c- s)n-1 ds
(c 1)2(n 1) oE(-c,1)
4 8( ) + 4( m
if 2 3
(c- 1)2(n 1)(n 2)
(26)
2 2log(c + 1) - log - 2 log 2
if 2.
(c 1)2
Combining Equation (23), Equation (24), Equation (25), Equation (26), we get that, if 2 3,
3(1 1/c" +1
2 2-4( 2 - )n2 + 2
Pin := P ( [-1,1]]") <
+ (n-1)(m-2](22-1) 1)
(27)
2(c² 1)(n -
2)
2
2
< 22 ((n-))
( (n-1)(n 2)
no²
50

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
and if 2,
3(1 (11-1/c) +1 2ll-oo - log - 2 log 3 3(1- (11-1/c)+9) 2(c² - 1)
Pin 2
(28)
2(c2
In the above, Pin is the probability that a randomly sampled test list x E IR" has all its elements lie within [-1, 1], that is, the probability that x lies within the domain of the training distribution. This probability is at most 011/n²)). Suppose that we generate N test instances, then straightforward application of the multiplicative Chernoff bound yields that with probability (²,)) samples will lie in the domain of the training distribution.
at least N-C for some constant C O, at most O n
-2-03
56-03
03
10
9
2c-02
10-02
8
7
6
02-
5
4
3
2
2e-01
5e-01
4
16
32
64
Figure 35: Contour plot of (2) (r E supp(Prrain (X))). i.c., the probability (upper bound in Equation (27) and Equation (28)) that a randomly sampled test instance x E I2" lies the domain of the training distribution.
For E (2, 8, 16,32) and e E 21. 10} which we consider in our experiments, Figure 35 shows contour plot of the probability upper bound Equation (27) and Equation (28). This probability is sufficiently small such that the majority of test instances in our test data does not belong to the domain of the training distribution.
For small and c, to determine the fraction of sampled test instances that will be within the domain of the training distribution, it more informative to directly invoke the additive Chernoff bound with Pin- Let N denote the total number of test instances that we sample, and further let Nin denote the number of sampled instances that lie in the domain of the training distribution. Then by the additive Chernoff bound we have that
PPN 2 N(pnn + e)) S exp(-2Ne3).
(29)
For example, suppose that sample N 1000 test instances from the test distribution. Suppose that we generate the test data using list length n = 2 and OOD scale factor c 2. Then in this case Pin S 0.4375. Take 0.0625. Then (29) says that with probability at least 0.9995, at least N/2 samples do not lie in the domain of the training distribution. For another example with slightly larger n and c, suppose that we generate the test data using list length n = and OOD scale factor e 10. Then Pin S 0.0034. Take = 0.0466. Then (29) says that with probability at least 0.98, more than 95% of test instances do not lie in the domain of the train distribution.
F. Potential reasons for failure in self-attention for out-of-distribution
In this section, we discuss potential reasons for the shortcomings in out-of-distribution generalization of standard Transform- ers for algorithmic tasks. While it is more straightforward to elicit reasons for the success positional attention motivated by the algorithmic alignment (Xu et al., 2020) between positional Transformers and parallel computational models - is considerably more challenging to pinpoint the causes of failure in standard Transformers.
51

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Finstly, Transformers simulate parallel algorithms, demonstrated by (Sanford et al., 2024). Intuitively, a single layer of self-attention should be more powerful than positional attention, as it leverages attention beyond positional encodings and allows for a more flexible structure in response input variations. However, as discussed in Section 5, executing parallel algorithms does not require using anything beyond positional information in attention.
Assuming that standard Transformers should adopt positional information to effectively execute parallel algorithms, the operations required by standard Transformers become increasingly difficult than positional Transformers for two main reasons:
1. Self-attention layers must learn to ignore input values and exploit positional information.
2. Transformer layers must preserve positional encodings for subsequent layers.
Namely, these desirable properties of positional Transformers present two significant challenges for standard Transformers. The first challenge arises naturally from the differences between standard and positional attention mechanisms. The second challenge highlights the compositional structure of attention layers, which can be detrimental during training. Specifically, the operations performed by each attention and MLP layer degrade the inputs of subsequent layers.
This issue is further emphasized in Remark B.6, where we state that while positional Transformers can represent any softmax pattern at any layer, standard Transformers may fail to do so due to potential degradation of the attention inputs. Although residual connections can mitigate this issue by preserving input information, they must ensure that no overlaps hinder the use of positional encodings in subsequent layers. Moreover, this problem compounds across layers, making training more difficult as errors in earlier layers adversely affect subsequent computations.
Nevertheless, these remain speculative reasons for the observed failure of standard Transformers. Determining the exact causes and the difficulty of achieving the two aforementioned goals through training requires thorough analysis of the training dynamics, which is inherently challenging. Future in-depth work within the mechanistic interpretability framework (Nanda et al., 2023) can potentially shed light on these issues by inspecting network parameters at convergence, thereby uncovering the underlying reasons for the failure standard Transformers.
Along this direction, we present some empirical evidence, in Figure 36 and Figure 37, that self-attention layers in a trained Transformer model can be highly sensitive to input values. In particular, the attention weights change dramatically when the input values of the test data do not necessarily lie in the domain of the training data. This suggests that self-attention potentially overfits the training data and, therefore, offers plausible explanation for why the standard Transformers exhibit such poor value generalization in our experiments.
52

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
c=
== 2
= 3
c == 4
c= 5
C
=
C
Layer 0
Layer
Layer 2
Layer 3
Figure 36: Visualization of learned attention weights in the standard Transformer model trained to solve the sorting task in experiments. The input list the model is c.X where X [1.75, 1.25, 0.75, 0.25, -0.25, -0.75, -1225, -1.75] and = 1,2,.... 8 is scaling factor. The model is trained on data whose input values range from -2 to 2. Therefore gives in-distribution data and larger yields OOD data. For each layer in the architecture, we plot of the 2 attention heads for illustration purposes. The trend for the other head is similar. Observe that the attention weights change dramatically as we increase the scaling factor of input values, with deeper layers suffering from more radical changes in the attention pattern e 2). This behavior potentially explains why the standard
1
under even small change in the scale (e.g. going from Transformer model performs poorly OOD test data in experiments.
2
3
4
5
Layer
Layer 1
Layer
Layer 3
Figure 37: Another visualization of learned attention weights in the standard Transformer model. We use the same setting as described in Figure 36, except that the input list is cX where X = [2, 22 -22 2, Again, we observe that the attention weights are highly sensitive to scaling factor e, especially those at deeper layers.
53

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
G. Informal Discussion on Out-of-distribution Generalization Bounds and future work
The topic of OOD generalization has been studied extensively from a theoretical perspective. Researchers often focus on bounding the risk on the test distribution using a mixture of terms that go to zero as the number of samples increases, plus a divergence term that does not necessarily go to zero depending on the distributions and the hypothesis class. For an extensive survey, we refer the reader to (Redko et al., 2022). Although, in general, such bounds offer valuable intuition, they might not be tight for our particular setting. In particular, we examine a popular type of bound found in (Mansour et al., 2009) which can be large even if the difference in the support of the train and test distributions is the smallest it can be. Note that there are more types of bounds in (Redko et al., 2022) than the one found in (Mansour et al., 2009). Although we have not conducted an in-depth analysis of all cases, we note that all of them depend in a worst-case on the hypothesis class. We believe that to improve upon such generie bounds, one must consider the dynamics of the training procedure for deep positional attention architectures. We find this topic extremely interesting, but we leave it for future work, as it is highly non-trivial.
In what follows, use a popular example of one of these bounds (Mansour et al., 2009) and illustrate why it is not tight for simple task of interest. Briefly, the main issue with this particular OOD bound is that it depends in worst-case manner on the hypothesis class. We demonstrate this issue using the task of computing the minimum over an array of length n. We assume that n is even. For simplicity, we do not work with the cumulative version of the minimum problem, as we did in the main paper. Therefore, the ground truth is simply the minimum of the input array. Durain is the train distribution over arrays of length n, where each component of the array is sampled independently and uniformly at random from integers in the range 0 to Dimin where D/raan is a constant. Drest is the test distribution over arrays of length n, where each component of the array is sampled independently and uniformly at random from integers in the range 0 to D 2, where 0 is constant integer. We use an equal number of samples from the train and test distributions, denoted by The loss function is (h((z). y) = [hxx)- - yl. where h is a hypothesis, x is an input array of length n. and y = min(x), which is the minimum function over the array sc.
The hypothesis class If is the architecture in Equation (1) with log, n layers, 2 heads per layer, and Wo and Wv as the identity matrices for all layers. For positional vectors in general position with dimension n, Lemma B.3 implies that there exist key and query matrices of size consists of 2 layers with a hidden dimension equal to 4. We use the ReLU activation function in all MLPs. This allows the MLP to represent the minimum and maximum functions on two input values exactly. This is because the minimum and maximum functions can be written using ReLUs and linear operations:
that can represent any attention pattern at each layer. The MLP at each layer
min(11, 42) 2 (ReLU(a) + x2) ReLU(-21 - -x2) ReLU - 22) ReLUU(x x11)
and
max (2), (ReLU(z) + x2) ReLU(-=x - -22) + ReLU - x2) + ReLU ( x1))-
Note that the MLP's ability represent the minimum function for two inputs exactly also the reason it can represent the maximum function. In other words, the exact representation of the minimum function comes with the consequence that the MLP can also represent the maximum function for two inputs. This observation is crucial later when we show that an existing popular bound from (Mansour et al., 2009) is not tight for this particular task.
Furthermore, we assume that [b(z)] is constant, which further implies that the magnitude of the loss is bounded above by a constant. Observe that this hypothesis class can represent a binary tree reduction algorithm for the minimum and maximum functions. This is possible because positional attention can represent any attention pattern, and the MLPs can represent the minimum and maximum functions for two inputs exactly. Specifically, the first layer of the positional attention architecture can represent the connections between layers 0 (leaf nodes) and 1 in the binary tree computational graph, the second layer can represent the connections between layers 1 and 2, and so on, up to the log, n-th layer. The MLPs are used locally at each node of the binary tree to compute the minimum between two input values. Therefore, the minimum and maximum functions over an array of n elements are in the hypothesis class.
Let us now discuss one of the most popular OOD generalization bound results for regression. We will use the third case of Theorem in (Mansour et al., 2009), which states the following.
Theorem G.1 (Theorem 8 in (Mansour et al., 2009), repurposed for the minimum function task). Assume that the loss function& is symmetric, it obeys the triangle inequality, and it bounded above by a constant. the minimum function is in
54

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
the hypothesis class H, then, for any hypothesis h E H, the following holds:
Reun (h) S RD (h) + disce (Drests Divaia)
where
disce (Diess, Derain) =
max E.
[R h²) DD (1,11)
and
RD (h, =
and RD (h, is defined similarty.
The above theorem states that the difference in risk between the test and train distributions is bounded only by the discrepancy term between the two distributions. It is important to note that the discrepancy term depends on the hypothesis class and measures the worst-case difference in the train and test risks within that class. The fact that the discrepancy considers the worst-case difference is why this bound is not tight for our task of computing the minimum function.
Let us now focus on the discrepancy term. Corollary 7 in (Mansour et al., 2009) states that
disce(Prett, Drain) S discy(Prest Drain) + 4(Rss + R. (11) + O
m
where Dest and Duain are the empirical versions of the test and train distributions, repsectively. We will use the common assumption that these empirical distributions are uniform over the samples. Sies and Stain are the sample sets for the train and test distributions, respectively. Moreover, R.S.. (f) and R..... (H) are the empirical Rademacher complexities for the train and test distributions, respectively. It is well-known that the part of the bound corresponding to the empirical Rademacher complexities goes to zero as the number of samples increases. The same holds for the square-root term the bound. Therefore, the only term left to understand is the discrepancy between the empirical distributions. Let us try to understand how this term behaves. Its definition is:
disce (Diest, D = max
Σ ((h()),(())) SE
Σ (((()) S
- m
K. H
A lower bound of discæ (Desa, Dunin) ) is given by setting h to be the minimum function and h to the maximum function:
disce (Drest Drrai))
Σ (max(x)- min(x)) S
1 Σ (max(x) min(x))
m
m7 SES
We claim that for polynomial number of samples m, e.g., m = T2", where is a positive integer, there exists no. such that for all n 2 no we have that min(x:) = 0 and max(x) = LD... for all x Suain and min(z) = 0 and max(z) = L + for all Syst with probability at least 0.8. The proof of this claim is trivial and provide it below. For now, let us discuss the implications of this claim. We have that discr Drain) 2 I. + I.D... 2 with probability least 0.8. Therefore, even if is the smallest it can be such that there is difference between the train and test distributions in this particular setting, i.e., 2 = 1, then the empirical discrepancy is going to be at least 1. This means that the upper bound of Theorem G. is at least one, and it is not going to zero as the number of samples increases. This is because the discrepancy definition considers the worst-case scenario without considering the training procedure. In practice, the training procedure may help to discover a hypothesis that is close to the minimum function since the minimum function is part of the hypothesis class. If the hypothesis discovered by the training procedure is close enough to the minimum function, the OOD generalization error may be much smaller than 1. Therefore, depending on the learning task and the hypothesis class, the bound in Theorem G. I can be loose.
Consider the following example, Lo.... = 3 and 2 = 1. Therefore, the largest value in the test distribution is 4. For m and n as noted above, the bound implies that the loss might be up to for any hypothesis h. This further implies that for any hypothesis h the relative error might be up to 25% with probability at least 0.8, despite the fact that the hypothesis class includes the true function and the training procedure could converge to a good approximation of it.
55

---

Positional Attention: Expressivity and Learnability of Algorithmic Computation
Let us prove the above probability claim. For Suain have
P(2; # O for all i E [m])
II P(x, +0)
;=1
II P(a E 21, D }
D
D
I.Dain
Duiia
Similarly, we have that
D
Pa, + Dani for all E [n])
I +
Furthermore, we have that
P(3ijj [n]] x, and 2, L.D.m)
1 P(z, + 0400 [n]or & L ViE [n])
2 1 - P(x: + 144 E [n])
P(r: + L 4 i E [n])
L 1 D +
n
Therefore, conclude that
T.Dyan + )))
D
P(all samples x E Stain have at least one 0 or L 2 1 1-2
and, similarly, we conclude that
L L + +z)) Pall samples E Stest have at least one 0 or L + 2 11-2 I.Daan + ))
For polynomial number of samples m, e.g., m = n², where e is a positive integer, there exists some Ro E N, such that for all n 2 no we have that the latter two probabilities are at least 0.9 (since for m = n both lower bounds tend to 1 as n tends to infinity). Therefore, our claim about the minimum and maximum over the sampled arrays holds with probability at least 0.81.
56


Reviews:
[{"review_id": "d1RPz0cdbg", "reviewer": "Reviewer_DrFP", "summary": "This paper proposes the Positional Transformer architecture for learning algorithmic problems over abstract data structures. In Positional Transformers, the attention maps are computed merely based on the positional embeddings and therefore the learned interaction patterns between different positions in the input sequence are independent of the values at those positions. This property makes Positional Transformers a better candidate compared to regular Transformers for certain class of algorithmic problems, especially in terms of OOD generalization as empirically demonstrated in the paper. The paper has also provided an elegant theoretical analysis proving the expressivity and learnability of Positional Transformers, most notably the generalization bound for Positional Transformers.", "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 the reviewer’s time in evaluating our manuscript. We are grateful for their recognition of the novelty and significance of our work. Below, we address their insightful comments.\n\n> There’s a quite extensive literature around GNNs for algorithmic problems and in particular NP problems (e.g. SAT solving, CSPs, etc.) that is missing from this paper. \n\nWe thank the reviewer for pointing out this gap in our literature review section. We have identified key references (original works and highly-cited) like [1,2,3,4,5,6,7] as well as a comprehensive survey [8], which we will include in the camera-ready version. If the reviewer believes specific works need to be mentioned explicitly, we would be happy to add them.\n\n> For the OOD case, do you think if we keep the Q and K projection matrices frozen in Positional Transformers and merely fine-tune the V projection matrices, we would be able to achieve the same level of generalization as the original in-distribution case? \n\nThis is a great question! The answer is yes: fine-tuning only the V projection matrix in the OOD case allows the positional Transformer to achieve the same level of generalization as in the in-distribution case. On the other hand, the same does not hold for the standard Transformer. We appreciate the reviewer’s suggestion of this scenario, which closely aligns with the typical approach to “fine-tune a pre-trained foundation model with unseen data” when dealing with new and potentially OOD data in practice.\n\nIn the tables below, we consider three simple algorithmic tasks and report the median test error (MSE) for both positional and standard transformer models when \n- they are trained and tested on in-distribution data (in-dist), \n- they are trained on in-distribution data but tested on numbers 10x larger than the original in-distribution data (10x OOD),\n- they are trained on in-distribution data, fine-tuned V projection matrices on numbers 10x larger than the original in-distribution data, and then tested on numbers 10x larger than the original in-distribution data (10x OOD with fine-tuning).\n\nWe find that fine-tuning only the V projection matrices brings the accuracy of the positional Transformer to the same level as the in-distribution case. However, for the standard Transformer, even after fine-tuning the V projection matrices, the accuracy on 10x test data is still very high compared to that of the positional Transformer. The comparison highlights another potential advantage of the positional Transformer for executing algorithmic tasks: one might efficiently fine-tune a pre-trained positional Transformer on OOD data and achieve good performance.\n\n| Task  | Positional (in-dist) | Positional (10x OOD) | Positional (10x OOD with fine-tuning) |\n| :--- | ---: | ---: | ---: |\n| sort | 1.20e-04 | 7.37e-04 | 1.23e-04 |\n| min | 1.03e-05 | 3.92e-04 | 2.19e-05 |\n| sum | 6.07e-06 | 1.31e-03 | 1.32e-05 | \n\n| Task | Standard (in-dist) | Standard (10x OOD) | Standard (10x OOD with fine-tuning) |\n| :--- | ---: | ---: | ---: |\n| sort | 9.05e-04 | 5.18e-01 | 6.90e-02 |\n| min | 1.74e-06 | 1.39e-01 | 4.00e-02 |\n| sum | 5.20e-06 | 1.09e+00 | 1.23e-01 | \n\n*We keep the same empirical setting where we train both models for 2000 epochs. For fine-tuning, we adopt the common practice and only train for 20 epochs. We normalize the MSE by the scale factor so that the table results indicate relative accuracy. \n\n---\n\nAgain, we thank the reviewer for carefully reading our manuscript and for their thoughtful feedback. We welcome any further questions.\n\nReferences:\n\n[1] Vinyals et al. Pointer Networks, NeurIPS 2015\n\n[2] Prates et al. Learning to Solve NP-Complete Problems: A Graph Neural Network for Decision TSP, AAAI 2019\n\n[3] Joshi et al. An Efficient Graph Convolutional Network Technique for the Travelling Salesman Problem, 2019\n\n[4] Bai et al. Learning-Based Efficient Graph Similarity Computation via Multi-Scale Convolutional Set Matching, AAAI 2020\n\n[5] Selsam et al. Guiding High-Performance SAT Solvers with Unsat-Core Predictions, SAT 2019\n\n[6] Karalias and Loukas Erdos Goes Neural: an Unsupervised Learning Framework for Combinatorial Optimization on Graphs, NeurIPS 2020\n\n[7] Gasse et al. Exact Combinatorial Optimization with Graph Convolutional Neural Networks, NeurIPS 2019\n\n[8] Cappart et al. Combinatorial Optimization and Reasoning with Graph Neural Networks, JMLR", "claims_and_evidence": "The paper has done an excellent job of precisely stating the claims as well as providing both theoretical and empirical evidence for those claims.", "methods_and_evaluation_criteria": "Regarding evaluating both expressivity and learnability, the authors have provided relevant settings and metrics to empirically evaluate the theoretical results.", "theoretical_claims": "I did not check the proofs.", "experimental_designs_or_analyses": "The paper does a good job of setting up sufficient experiments to evaluate both the merits and the shortcomings of Positional transformers in solving algorithmic problems. I especially liked the fact that different classes of problems were used to pinpoint the scenarios in which Positional Transformers are indeed useful.", "relation_to_broader_scientific_literature": "I believe the findings of this paper can be further relevant useful for other areas of Deep Learning such as neuro-symbolic methods. In particular, it'd be interesting to see how symbolic computation can be further injected as inductive biases into the positional transformers and more generally purely neural models.", "essential_references_not_discussed": "There's a quite extensive literature around GNNs for algorithmic problems and in particular NP problems (e.g. SAT solving, CSPs, etc.) that is missing from this paper.", "other_strengths_and_weaknesses": "Strengths:\n\nThis paper is very well-writte and well-motivated. \nThe claims are carefully stated and adequately supported. \nThe authors have done a great job with their empirical study covering both the strength and weaknesses areas for Positional Transformers.\n\nWeaknesses:\n\nIn terms of practicality of the proposed models, it'd be interesting to have experimental results on more realistic datasets with much larger problem sizes.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- For the OOD case, do you think if we keep the Q and K projection matrices frozen in Positional Transformers and merely fine-tune the V projection matrices, we would be able to achieve the same level of generalization as the original in-distribution case? I believe that'd be an interesting ablation study which can further shows the importance of Positional Attention for certain algorithmic problems.", "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": "Xy37QTgpEA", "reviewer": "Reviewer_U2No", "summary": "This paper presents transformer with positional attention (PT), a mechanism that implements data-free query and key inputs for computing attention scores. From the empirical side, this mechanism aims to emulate the massively parallel computation (MPC) model, which the authors show expressivity and learnability results. Experimental results are presented that shows on select algorithmic tasks positional attention can generalize better than standard attention.", "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 thank the reviewer for their comments and appreciate their recognition of the soundness of our paper. We appreciate the connections made to other related works, which we will incorporate into the main paper. Below, we address their insightful comments, specifically those written in Methods And Evaluation Criteria:\n\n> [...] it'd be great to discuss its limitation (beyond those discussed in Section 8) in modeling natural language data discuss how this type of algorithmic transformer could be useful in real world settings.\n\nWe appreciate the reviewer’s suggestion. Along this line, one relevant computational task we considered in our work is the k-hop inductive heads task discussed in Section 7, which is inspired by the in-context learning abilities of language models and captures some of the pattern matching and higher-order dependencies found in NLP tasks [1]. While we agree that experiments in Natural Language are an interesting direction for studying the effectiveness of positional Transformers (PT), the manuscript is already quite extensive. To ensure a thorough analysis, we believe it is best to reserve these explorations for future work, which we intend to pursue as a follow-up. Regarding its usefulness in real-world settings, one possible direction is applying PT to unstructured tabular data in the form of text, where prompts require algorithmic computation to be answered. In Section 7.3, we present experiments in a simplified setting where PT proves useful and outperforms standard Transformers (ST). Building on this, we are actively preparing further empirical studies in this direction as part of our future work.\n\n> [...] (though of course, one could argue that this is precisely PT's benefit). I do believe that additional experiments on real datasets, even training small language models could further strengthen this paper.\n\nWe appreciate the reviewer’s insight and agree that additional real-world experiments could further illustrate PT’s potential. Interestingly, our ablation study in Appendix D 1.7 already captures this intuition but also reveals more: even when positional encodings are incorporated at every layer of ST, its performance remains suboptimal. This suggests that PT’s advantage is not merely due to its handling of positional information but rather to the presence of two distinct computational streams—one for positions and one for input values—which we hypothesize is a key factor behind its effectiveness. Furthermore, the new theoretical learnability trade-off for the in-distribution setting in Section 6 comes from this decoupling property (see comments in line 176).\n\nAgain, we thank the reviewer for carefully reading our manuscript and for their thoughtful feedback. We welcome any further questions.\n\nReferences\n\n[1] Olsson et al. In-context learning and induction heads. 2022", "claims_and_evidence": "Yes. I believe the claims are well articulated. The authors clearly established their aim in establishing equivalence between PT and MPC in Section 5, and discussed potential limitations (in terms of quadratic complexity) from ln 207 onwards. The generalization bound given in Section 6 appears to be correct.", "methods_and_evaluation_criteria": "I believe the experiments are designed reasonably well to establish that PT can achieve comparable (in-distribution) or better (out-of-distribution) performance on proposed algorithmic tasks, which are commonly used for evaluating theoretical works that establish transformers as computation models. While I understand that real-world applications is not a focal point of this paper, and I certainly does not lower my review because of it, it'd be great to discuss its limitation (beyond those discussed in Section 8) in modeling natural language data, or discuss how this type of algorithmic transformer could be useful in real world settings.", "theoretical_claims": "Due to time constraint, I was only able to check the expressivity results (Appendix B), which seems correct. I was not able to check generalization result in Appendix C carefully but it seems to be a an application of Theorem 5 of Bartlett et. al. (2003), coupled with Lipschitz results of the PT model. And I'm conservatively optimistic about the results.", "experimental_designs_or_analyses": "The experiments are designed to assess the generalization performance of PT on algorithmic tasks. The generalization and OOD experiments are done properly.\n\nOne of my core concern for this comparison is that the authors directly supplement positional encoding in the form of QK for PT, but it is only fed into the standard transformer at the input layer. One could imagine that for deeper networks, positional encoding may vanish for standard transformers but not the case for PT (though of course, one could argue that this is precisely PT's benefit). I do believe that additional experiments on real datasets, even training small language models could further strengthen this paper.", "relation_to_broader_scientific_literature": "This work falls into a growing body of work that assess the algorithmic capability of transformers [1][2], investigations of optimal encodings of positional information [3], graph neural nets for algorithmic execution [4]. \n\n\n[1] Neural Networks and the Chomsky Hierarchy, https://arxiv.org/abs/2207.02098\n\n[2] Training Neural Networks as Recognizers of Formal Languages, https://arxiv.org/abs/2411.07107\n\n[3] Round and Round We Go! What makes Rotary Positional Encodings useful?, https://arxiv.org/abs/2410.06205\n\n[4] Everything is Connected: Graph Neural Networks, https://arxiv.org/abs/2301.08210", "essential_references_not_discussed": "I'm not aware of any undiscussed literature.", "other_strengths_and_weaknesses": "I believe this is a solid paper with a valid hypothesis, correct theoretical analysis, and reasonable experimental design. OOD (or length) generalization is a challenging task in this line of literature and it's nice to see some positive results from PT.", "other_comments_or_suggestions": "NA", "questions_for_authors": "See comments about discussions on natural language datasets and broader impacts in the **Methods And Evaluation Criteria** section.", "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": "P8lXqbDgN6", "reviewer": "Reviewer_USLR", "summary": "This paper introduces and analyzes positional attention in Transformers where attention weights depend exclusively on positional encodings rather than on input data. The authors prove that Transformer with positional attention with logarithmic depth has the same expressive power as MPC, and demonstrate that positional Transformers has a trade-off in-distribution and has certain benefit in OOD tasks when the task relies on positional information.\n\n### Update after rebuttal\n\nThe response addresses most concerns so I will keep my positive rating.", "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 sincerely thank the reviewer for taking the time to read and review our manuscript. In what follows, we provide a comprehensive response to the weaknesses highlighted by the reviewer.\n\n> The proposed method seems to be heavily dependent on the positional encodings used in the Transformer model and the experiment only uses one type of the positional encodings. How to learn useful positional encodings and how the positional Transformer will behave using other positional encodings remain uncertain\n\nAs noted in lines 289-293 (second column) of the main paper, we have experimented with various configurations of positional encodings, including sinusoidal, binary, and rotary (RoPE) positional encodings. In all cases, we observe the effects reported in the main paper. We refer the reviewer to Appendix D 1.6 for more details on those cases.\n\n> As mentioned in Experimental Designs Or Analyses, length generalization seems more natual to Transformers as OOD and the paper does not investigate it. It also poses extra difficulity since the depth of the model needs to be updated when the length of the input changes. How to solve this problem remains unknown from the paper.\n\n*Depth of the model:* We’d like to emphasize that a fixed-depth Positional Transformer can indeed work with inputs of different lengths (up to a predetermined length). This is demonstrated by our experiments, where a PT model is trained to perform algorithmic tasks on multiple input lengths, a setting reminiscent of the notion of context length in Transformers. Appendix D.1.4 presents a single model capable of processing different input lengths up to a fixed upper bound, demonstrating the effectiveness of our approach in processing variable-length inputs while achieving good generalization on unseen input values. We would also like to point out that our in-distribution generalization bound does not require fixing the number of layers; it assumes a fixed upper bound on the length instead. \n\n*Types of OOD generalization:* Algorithmic execution using neural networks involves multiple kinds of OOD generalization. Length generalization refers to executing the same algorithm on input lengths unseen during training, while value generalization applies the algorithm to new input values not encountered during training. Both are crucial, as an ideal model should perform well on unseen lengths and values. In this work, our primary focus is on value generalization, as it ensures that, for a fixed upper bound on the length, the model has learned an algorithmic execution applicable to any combination of numbers.\n\n> The benefit of the OOD case in the paper seems to be highly related with the nature of the tested task. Is there any case where the standard Transformer behave better than positional Transformer in OOD? Such results would also help to understand the particular inductive bias that both models favor.\n\nWe thank the reviewer for the insightful question. As mentioned in Section 7.2, we hypothesize that the alignment between the architecture and the underlying problem promotes OOD generalization. While Figure 4 shows instances where Positional Transformers outperform standard Transformers, we also present a case where the opposite occurs. In the k-hop induction heads task, which involves dynamic communication, our hypothesis of a fixed communication graph is violated. As a result, our architecture underperforms compared to standard Transformers, which converge more quickly and with fewer layers, an advantage we attribute to algorithmic alignment. Here, the dynamic nature of the problem aligns better with the flexibility of self-attention compared to the rigidity of positional attention.\n\n---\n\nOnce again, we thank the reviewer and welcome any further comments. We hope the clarifications are sufficient for them to consider raising their score.", "claims_and_evidence": "Yes the claims made in the submission supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes the proposed methods and/or evaluation criteria (e.g., benchmark datasets) make sense for the problem or application at hand.", "theoretical_claims": "I briefly checked the proofs in Appendix B.", "experimental_designs_or_analyses": "I checked the experimental setups for both in-distribution regime and out-of-distribution regime. I found the out-of-distribution regime rather limited since in Transformers people also consider longer lengths as OOD.", "relation_to_broader_scientific_literature": "1. The paper proposed to use positional only data to get K, Q matrices in attention as a noval architecture design as opposed to traditional attention mechanism.\n2. The paper shows the expressiveness of the positional Transformer respect to MPC, which is new to my best knowledge.\n3. The paper honestly presents the trade-offs between positional and standard Transformers, showing where each excels and explaining the underlying reasons.", "essential_references_not_discussed": "Not to my knowledge.", "other_strengths_and_weaknesses": "Strengths: see contributions.\n\nWeaknesses:\n1. The proposed method seems to be heavily dependent on the positional encodings used in the Transformer model and the experiment only uses one type of the positional encodings. How to learn useful positional encodings and how the positional Transformer will behave using other positional encodings remain uncertain.\n2. As mentioned in Experimental Designs Or Analyses, length generalization seems more natual to Transformers as OOD and the paper does not investigate it. It also poses extra difficulity since the depth of the model needs to be updated when the length of the input changes. How to solve this problem remains unknown from the paper.", "other_comments_or_suggestions": "For Weaknesses 1&2, the paper would benefit from further investigation in such directions.", "questions_for_authors": "1. The benifit of the OOD case in the paper seems to be highly related with the nature of the tested task. Is there any case where the standard Transformer behave better than positional Transformer in OOD? Such results would also help to understand the particular inductive bias that both models favor.", "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": "0LIf4PyQC1", "reviewer": "Reviewer_WR9k", "summary": "This paper is a theoretical study on positional transformers, where the attention is determined solely by positional encoding regardless of the content. The paper includes a representation theory that position transformers can simulate MPC, followed by a generalization bound. Empirical study is conducted on several synthetic tasks, showing that positional transformers can indeed complete the tasks, though the model needs more layer and data to achieve so.", "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 thank the reviewer for allowing us to clarify the motivation of our paper. We will address their questions individually and provide general remarks below.\n\n*High-level motivation:* Generally, there is growing interest in the relationship between neural networks and computational models [1,2,3,4]. Computational models execute algorithms on data. For example, Figure 7 in our paper shows the communication trace of an algorithm for computing a parallel cumulative sum.\n\n*Motivation for positional constraints:* As in Figure 7 and many other parallel algorithms, data does not influence communication among processors within computational models. Instead, communication relies solely on processor IDs (or positions). For standard Transformers (ST), this implies that the computation of attention weights can be independent of the input values for certain algorithms, and that is how we arrive at positional Transformers (PT), which is our object of study. In c., we discuss the trade-offs of these constraints.\n\nWe now address specific concerns raised by the reviewer.\n\n> a. The novelty is not clear. I think the MPC representation result is also proven for standard transformers.\n\nPT poses an algorithmic computation-inspired departure from the original attention mechanism, and the theoretical results of one do not directly translate to the other. This fact alone justifies the need for an MPC simulation result. The proof techniques we used are significantly different from the ones used for the ST and cannot be directly borrowed from [1], which we further discuss in response b.\n\n> b. [...] the simplification of attention make the theoretical analysis easier and we can obtain more theoretical insights to PT, and more importantly, to ST.\n\nThe reviewer seems to suggest that the goal of introducing PT is to simplify the theoretical analysis of ST. We would like to emphasize that this is not the case. We argue that the expressivity analysis of PT is even more challenging. Simulating MPC communication with ST is facilitated by self-attention’s ability to depend on the input values. In contrast, PTs rely on fixed communication via positional encodings, making simulation more challenging. Consequently, MPC simulation with PT must show dynamic communication over a static network, which we accomplish by using Benes networks, a novel strategy unseen in previous related works.\n\n> c. However, I don’t quite see the gain of the positional constraints. It seems [...] PT is worse (requires more layers). \n\nThe comparison between PT and ST is more nuanced than suggested. PT offers a different learnability trade-off, illustrated in theory and practice. While PT requires more layers to simulate MPC, this is a worst-case analysis that accounts for all types of algorithms MPC can execute. However, as shown in Section 6, PT has an advantage in learnability: it removes a dependency on the norms of the query and key parameters when all other factors are fixed. Moreover, our empirical results show the competitiveness of PT: Figure 2 showcases some algorithmic tasks where PT outperforms ST in-distribution (ID) with the same number of layers. Furthermore, PT consistently outperforms ST in OOD scenarios for those tasks (Figure 4). As discussed in Section 7.2, we hypothesize that PT exhibits an OOD advantage in data-independent algorithmic tasks due to its architectural alignment with the target function. Our experiments are designed to test this hypothesis.\n\n> d. Why [...] study PT then? To show [...] significant representation power loss with only position? \n\nThe conclusion that PT lacks representation power is incorrect. Our theory demonstrates that with logarithmically more layers, similar to ST, PT can simulate any computation within the MPC protocol. Empirically, we observe that for many algorithmic tasks, PT has comparable ID performance and significantly better OOD performance than ST, with the same number of layers. We refer the reviewer to our response in c. for more context.\n\n> e. Similarly, I do not understand what insight we gain by studying PT rather than standard transformers from the generalization result.\n\nThis is highlighted in Sections 1 and 6 of the paper. The interesting takeaway is the tradeoff between a better generalization bound (independent of parameter norms) and the (potential) need for logarithmically more layers needed to maintain expressivity. We refer the reviewer to response c. for more details.\n\n---\nOnce again, we thank the reviewer and welcome any further comments. We hope the clarifications are sufficient for them to consider raising their score.\n\nReferences:\n\n[1] Sanford et al. Transformers, parallel computation, and logarithmic depth. ICML 2024\n\n[2] Pérez et al. Attention is Turing-Complete. JMLR 2021\n\n[3] Loukas, A. What graph neural networks cannot learn: depth vs width. ICLR 2020\n\n[4] Wei et al. Statistically meaningful approximation: a case study on approximating Turing machines with transformers. NeurIPS 2022", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "I checked the soundness of the theorem, but not the proof.", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "The contribution of the paper seems to be extending the theoretical analysis of standard transformer to positional transformers.", "essential_references_not_discussed": "no", "other_strengths_and_weaknesses": "Strength:\n- The MPC representation theory is interesting, which can shed light on how transformers solve problems;\n- It seems that a better generalization bound is derived thanks to the positional restriction.\n- Positional transformers can be indeed of interest for some tasks, as they are at least as powerful as CNNs.\n\nWeaknesses:\n- The novelty is not clear. I think the MPC representation result is also proven for standard transformers.\n- The biggest concern I have is, I are not sure *why* we need to study positional transformers. I think the motivation is might the simplification of attention make the theoretical analysis easier and we can obtain more theoretical insights to positional transformers, and more importantly, to standard transformers. However, I don’t quite see the gain of the positional constraints. It seems that (Sanford et al., 2024) already have the MPC simulation result, and the positional version is worse (requires more layers). Why do we want to study  positional transformers then? To show that the content is indeed important, and you will suffer significant representation power loss with only position? Similarly, I do not understand what insight we gain by studying positional transformers rather than standard transformers from the generalization result.", "other_comments_or_suggestions": "If the authors can make the novelty and motivation more convincing, I can raise my score.", "questions_for_authors": "see weaknesses\n\nUpdate:\n=\n\nthe authors have addressed my concerns on the motivations of studying PT, and the OOD generalization result is interesting. I raised my score.", "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: This paper introduces Transformer with Positional Attention (PT), a mechanism that computes attention scores using only positional information rather than content. The authors establish the equivalence between PT and the MapReduce-like distributed computing model (aka MPC), provide generalization bounds, and demonstrate through experiments that PTs can generalize better than standard Transformers on select algorithmic tasks, particularly in out-of-distribution scenarios. Even though the near-equivalence between MapReduce and TF has been widely known in the theory community, I believe it is a solid contribution to make this statement more rigorous by proposing a variant of attention. 

The reviewers unanimously recommend acceptance based on several strengths. First, the paper presents clear theoretical contributions with expressivity and learnability results. Second, the experimental design effectively demonstrates both the merits and limitations of PTs across different algorithmic problems. Third, the paper presents the trade-offs between positional and standard Transformers, showing where each excels. Fourth, the OOD generalization results seem valuable, as this is a challenging area in transformer research. Overall, the paper is well-written and well-motivated with claims carefully stated and supported.
Decision: Accept (poster)