# Explicitly Encoding Structural Symmetry is Key to Length Generalization in
Arithmetic Tasks

Mahdi Sabbaghi  University of Pennsylvania
{smahdi, pappasg, hassani<EMAIL_ADDRESS>George Pappas  University
of Pennsylvania
{smahdi, pappasg, hassani<EMAIL_ADDRESS>Hamed Hassani  University
of Pennsylvania
{smahdi, pappasg, hassani<EMAIL_ADDRESS>Surbhi Goel  University of
Pennsylvania
{smahdi, pappasg, hassani<EMAIL_ADDRESS>
###### Abstract

Despite the success of Transformers on language understanding, code
generation, and logical reasoning, they still fail to generalize over length
on basic arithmetic tasks such as addition and multiplication. A major reason
behind this failure is the vast difference in structure between numbers and
text; For example, the numbers are typically parsed from right to left, and
there is a correspondence between digits at the same position across different
numbers. In contrast, for text, such symmetries are quite unnatural. In this
work, we propose to encode these semantics explicitly into the model via
modified number formatting and custom positional encodings. Empirically, our
method allows a Transformer trained on numbers with at most 5-digits for
addition and multiplication to generalize up to 50-digit numbers, without
using additional data for longer sequences. We further demonstrate that
traditional absolute positional encodings (APE) fail to generalize to longer
sequences, even when trained with augmented data that captures task
symmetries. To elucidate the importance of explicitly encoding structure, we
prove that explicit incorporation of structure via positional encodings is
necessary for out-of-distribution generalization. Finally, we pinpoint other
challenges inherent to length generalization beyond capturing symmetries, in
particular complexity of the underlying task, and propose changes in the
training distribution to address them.

##  1 Introduction

Large language models (LLMs) [BMR+20, Ope23], powered by transformer
architectures [VSP+17], and extensive computational resources, have
demonstrated emergent abilities in language understanding, code generation,
and logical reasoning [WTB+22]. Despite these capabilities, LLMs still
struggle to generalize when handling unseen data in highly structured tasks
like arithmetic [DLS+23, LAG+23a]. Specifically, they fail to extrapolate from
shorter instances to longer ones, a challenge known as _length
generalization_. This issue is particularly evident in arithmetic, where
models fail to generalize from smaller numbers to larger ones. [NAGA+21a,
NJL21, LSL+23].

It is now well understood that the major reasons for the failure due to
structural issues are the use of positional encodings that do not scale with
sequence length [KPR+23, SBE+23], and format of the data itself [SBE+23].
Though the Transformer can learn to solve this task for a fixed length of
numbers given enough samples, the inductive bias of the architecture is
insufficient to capture such structure across unseen positions, in order to
generalize to longer lengths.

We posit that the challenges of length generalization in arithmetic tasks,
such as addition and multiplication, can be broken down into two factors, (1)
increase in complexity of the task, (2) inability to capture the positional
structure in the tasks. For example, when we add two numbers, the positional
structure corresponds to aligning each digit from right to left to compute the
digit-wise addition, and the complexity of the task corresponds to the longest
sequence of carry-overs. In contrast, when we perform multiplication of two
numbers (say 1 digit number multiplied with $k$ digit number), the positional
structure corresponds to aligning the first number with all of the digits of
the second number to compute digit-wise multiplication, and complexity
corresponds similarly to longest sequence of induced carries that need to be
resolved.

We focus on rectifying the inability of Transformers to capture positional
structure and provide strong evidence for the need to explicitly encode
structural biases into the architecture for achieving length generalization.
We show that with the correct data format and choice of positional encodings,
a Transformer trained on only data with numbers of a bounded size (e.g., $\leq
5$) can generalize to numbers of much longer length (e.g., $\geq 50$) on both
addition and multiplication. We further show, empirically and theoretically,
that standard tricks of implicitly inducing this positional structure, such as
data augmentation with shifting contexts [BMR+20], are not sufficient to
guarantee this generalization. Our key findings are summarized below:

(F1) Positional encodings that encapsulate the structure enable length
generalization for addition and multiplication. For addition, relative
position encoding (RPE) allows the Transformer to naturally capture digit-wise
alignment in order to achieve length generalization from adding two numbers
with up to 5-digit to two numbers with up to 50-digit (similar generalization
was observed in [JdDE+23b]). For multiplication, our proposed new uniform
positional encoding (UPE) that assigns each digit of a number the same
position encoding, coupled with RPE, allows the Transformer to length
generalize from multiplication of 3-digit numbers with 5-digit numbers to
multiplication of a 3-digit number with 20-digit numbers. To our knowledge,
this is the first positive result on length generalization for multiplication
without using chain-of-thought [WWS+22] or data priming [JdDE+23b].

(F2) Data augmentation is insufficient for length generalization. The most
natural data augmentation of shifting the numbers by adding zeros to the right
to implicitly induce the necessary alignment needed for both addition and
multiplication helps with length generalization; However, the performance is
significantly worse than the crafted positional encodings. To further
elucidate this, we propose a simple linear setting that captures the relative
position alignment of addition, and prove that a one-layer linear Transformer
trained with gradient descent with RPE generalizes, and shifting with zeros to
encode a similar bias does not help.

(F3) Incorporating higher complexity shorter sequences improves generalization
to longer sequences. While our proposed positional encodings address the
structural aspect, we also isolate the other major reason for failure on
longer samples to be related to the increase in complexity of the task. We
show that the increase in coverage of higher complexity examples, which lie in
the tails of the distribution over shorter sequences, leads to improvements in
performance on longer sequences. For the task of addition, we show the
required length of shorter sequences needed to adequately cover complex
examples for larger sequences under the uniform distribution on numbers is not
very large. This allows us to cover sufficient complexity with 5-digit numbers
to generalize to 50-digit numbers.

#### Related Work.

Despite significant advances, Transformer models such as GPT-4 [Ope23]
struggle with high-complexity arithmetic tasks [DLS+23, YDL+23].
Investigations into these limitations have ranged from ineffective
representation learned for numbers [WSC+16, SHB16, WWL+19, TPSI21], to the
inability to achieve numeracy via unsupervised learning [RaGS22, LM22, KHK+21,
WZNZ21, MSK+23]. Attempts to mitigate these deficiencies such as fine-tuning
pre-trained models[GGB20, LAD+22], scratch-padding[NAGA+21b], and leveraging
chain-of thoughts reasoning [WWS+22, OWJ+22] only provide marginal out-of-
distribution (OOD) benefits [AWA+22]. Specifically, [AWA+22] shows that
scratchpad fine-tuning fails to generalize to longer problems, and even with
chain-of-thoughts, model are prone to error-propagation, learning shortcut
solutions and insufficient features, etc [DLS+23]. Recent efforts have
proposed modification to data and model structures to achieve length
generalization [JdDE+23a, SBE+23]. This includes introducing new structures to
positional vectors or leveraging positional information like [SUV18, HVU+18,
SLP+23, RSR+23]. Some studies suggest removing positional vectors or using
randomized encodings [KPR+23, RDG+23]. Despite these efforts, arithmetic tasks
remain difficult. In this work, we assert that without explicitly encoding
symmetries, generalization remains unattainable even with implicitly encoding
symmetries via data augmentation (see Section 3). Building on prior work, we
employ RPE [SUV18] motivated by the results of [JdDE+23a] for addition, and
design our custom embedding for multiplication in Section 3, providing a new
framework to leverage symmetries effectively and give the first theoretical
explanation for the effectiveness of this approach in Sections 3.1 and 5.

Beyond architectural changes, modifying the training distribution has also
been a prominent direction for improving performance [LSL+23, SBE+23],
particularly through curriculum learning [WCZ21, YDL+23], and priming
[JdDE+23b]. Our observations in Section 4 suggest that exposing models to more
complex in-distribution samples can enhance OOD performance. We provide
additional related works in Appendix A.

##  2 Preliminaries: Positional Encodings

To ground our design of positional encodings, we first describe a general
formalism for positional encodings in this section.

#### Transformer Architecture.

A transformer model takes as input a sequence of $d_{x}$-dimensional vectors,
denoted as $\bm{x}=(x_{1},\dots,x_{n})$ where $x_{i}\in\mathbb{R}^{d_{x}}$.
These vectors are embedded using an embedding matrix
$\bm{E}:\mathbb{R}^{d_{x}}\to\mathbb{R}^{d_{z}}$, producing
$\bm{z}^{(0)}=(z^{(0)}_{1},z^{(0)}_{2},\dots,z^{(0)}_{n})=(\bm{E}x_{1},\dots,\bm{E}x_{n})$,
where $d_{z}$ is the dimension of the embedded space. Each Transformer block
consists of attention layers operating on $\bm{z}^{(l-1)}$ to produce
$\bm{z}^{(l)}$ using $H$ heads, followed by a shared MLP layer. Each head
computes attention scores:

$A^{(l)}_{i,j}={\rm
Softmax}_{j}\big{(}(W_{Q}z^{(l-1)}_{i})^{T}(W_{K}z^{(l-1)}_{j})\big{)}$

The output of each attention head is computed as follows:

$b^{(l)}_{i}=\sum_{j}A^{(l)}_{i,j}(W_{V}z^{(l-1)}_{j})\in\mathbb{R}^{d_{z}/H}$

where $W_{Q},W_{K},W_{V}\in\mathbb{R}^{d_{z}/H\times d_{z}}$, are three
matrices that together form the core of each head. The concatenated head
outputs are processed by a feed-forward network to obtain $z^{(l)}_{i}$.

In the formulation given above, the output at position $i$ remains invariant
to the permutation of the rest of the sequence
($\\{x_{1},\dots,x_{i-1},x_{i+1},\dots,x_{n}\\}$), which is a limitation for
our tasks where the sequence order is indispensable. As a result, positional
encodings are employed.

#### Absolute Positional Encoding (APE).

One approach to overcome this limitation is to encode positional information
by adding a collection of unique and fixed vectors to the input $x_{i}$’s:

$z^{(0)}_{i}=\bm{E}x_{i}+p_{i}$

This method is known as absolute positional encoding (APE), where the
positional vectors can either be predefined and fixed or learned during
training, adapting to specific task requirements.

#### Pairwise Positional Encoding.

More generally, since every layer of an attention model only considers
pairwise interactions between elements, we can move the positional vectors
inside the calculation of attention scores either by adding the pairwise
vectors ($p_{ij}$) to the embedded keys, or to the embedded queries, or to the
computed attention between a query and a key. In all experiments, we adhere to
the method of adding them to the key, as suggested in [SUV18]. That is,

$A_{i,j}={\rm Softmax}((W_{Q}z_{i})^{T}(W_{K}z_{j}+p_{i,j})).$ (2.1)

We propose tailoring positional vectors $p_{i,j}$ to the specific
relationships between positions in a task. If the connection from position
$i_{1}$ to $j_{1}$, denoted by $i_{1}\to j_{1}$, has the same functionality as
$i_{2}\to j_{2}$, we can set $p_{j_{1},i_{1}}$ equal to $p_{j_{2},i_{2}}$.
This ensures $i_{2}$ influences $j_{2}$ similarly to how $i_{1}$ influences
$j_{1}$. In general, if any positions have equivalent functionality, we will
set their corresponding pairwise positional vectors to be the same.

#### Relative Positional Encoding (RPE).

For tasks with translational symmetries, we have the following property:
$\big{(}i\to j\big{)}\sim\big{(}(i+t)\to(j+t)\big{)}\quad\forall t$, in which
"$\sim$" stands for being equivalent in our notation. This symmetry implies
that: $p_{i,j}=p_{i+t,j+t}$ for all $t$. Therefore, we can express positional
vectors using the difference between $i$ and $j$:

$p_{i-j}\leftarrow p_{i,j}$

Consequently, computations are based on relative positions, i.e. the output
depends on the relative, not absolute, positions in the sequence, hence the
term relative positional encoding (RPE) [SUV18].

##  3 Length Generalization for Arithmetic tasks

In this paper, we focus on two arithmetic tasks: addition and multiplication.
We first formalize the setup of our tasks and then describe our findings.

#### Data format.

The input consists of two integers represented per-digit as
$x^{(1)}=(x^{(1)}_{l_{1}},\dots,x^{(1)}_{1}),x^{(2)}=(x^{(2)}_{l_{2}},\dots,x^{(2)}_{1})$,
and we ensure both sequences are of the same length ($l_{1}=l_{2}=l$) by
adding "pad" tokens to both integers, as necessary. Consider for illustration
the following example: "123 + 4095 = 4218", and $l=20$:

$\displaystyle\text{{Input}:
}(\underbrace{\text{pad},\dots,\text{pad}}_{\text{17
times}},1,2,3,+,\underbrace{\text{pad},\dots,\text{pad}}_{\text{16
times}},4,0,9,5)$ $\displaystyle\text{{Output}: }\hskip
52.03448pt(\underbrace{\\_,\dots,\\_}_{\text{20
times}},\underbrace{\text{pad},\dots,\text{pad}}_{\text{17 times}},4,2,1,8)$
(3.1)

For multiplication, we use the same input format with the difference that the
multiplier is not padded and the "$+$" token is replaced with "$\times$":

$\displaystyle\text{{Input}:}\ \hskip
31.07556pt(5,6,\times,\underbrace{\text{pad},\dots,\text{pad}}_{\text{16
times}},4,2,9,7)$ $\displaystyle\text{{Output:}}\
(\\_,\dots,\\_,\underbrace{\text{pad},\dots,\text{pad}}_{\text{16
times}},2,4,0,6,3,2)$ (3.2)

Note that in the example above, the maximum length of the output can reach 22
digits, accounting for the worst-case scenario of multiplying a 20-digit
number by a 2-digit number.

Indexing from right to left for simplicity, the first digit of the answer
corresponds to the first position of the input. The "$\\_$" symbol in the
output indicates neglected tokens that are not part of the supervision signal
for the addition task. Note that our formatting approach is distinct from that
of [JdDE+23a] in that we pad each integer from the left to ensure that the
relative coordinates of digits remain consistent regardless of the integers’
lengths, and also from other prior approaches that reverse the digits [SBE+23,
LSL+23]. The fixed length assumption is not unrealistic since numbers are
typically represented with fixed bit-complexity in most computations.
Consequently, the tokenizer can readily impose a fixed-length structure of any
chosen length as sketched in Figure 6.

#### Length generalization task.

For our tasks, the output consists of the integer resulting from the
arithmetic operation on $x^{(1)}$ and $x^{(2)}$. We can model our length
generalization problem as an instance of generalization on the unseen
[ABLR23]. We denote the domain of all pairs of integers with lengths $\leq l$
by $\mathcal{D}:=\Big{\\{}(x^{(1)},x^{(2)})\in\mathbb{Z}^{2}:0\leq
x^{(1)},x^{(2)}<10^{l}\Big{\\}}$. In the problem of length generalization , we
observe only a subset of this domain during training. Specifically, pairs of
integers with lengths at most $l_{s}<l$ are used in training. We denote this
_seen domain_ as
$\mathcal{D}_{s}:=\Big{\\{}(x^{(1)},x^{(2)})\in\mathbb{Z}^{2}:0\leq
x^{(1)},x^{(2)}<10^{l_{s}}\Big{\\}}$, and denote the _unseen domain_ , which
consists of at least one integer with longer length, as
$\mathcal{D}_{u}:=\Big{\\{}(x^{(1)},x^{(2)})\in\mathbb{Z}^{2}:0\leq
x^{(1)},x^{(2)}<10^{l},\ 10^{l_{s}}\leq x^{(1)}\lor 10^{l_{s}}\leq
x^{(2)}\Big{\\}}$. We thus have:
$\mathcal{D}=\mathcal{D}_{s}\cup\mathcal{D}_{u}$. We will focus on the
following question: If we train the model using samples from $\mathcal{D}_{s}$
to near perfect in-distribution accuracy, what can be inferred about the
accuracy on $\mathcal{D}_{u}$?

#### Model.

For all tasks, we utilize a BERT-based encoder-only attention model [DCLT18].
We use an encoder-only architecture because causal maps are not beneficial for
our tasks. For instance, to compute the second digit of a sum, the model must
consider the first two digits of both integers. This entails both looking
forward and backward in the input. We also chose to avoid autoregressive
settings because they often require several tricks to function, such as
reversing the digits order either in the product or the inputs [SBE+23,
LSL+23]. Since arithmetic operations naturally start from the least
significant digit, generating them from left to right in an auto-regressive
manner offers no computational advantage. The number of attention layers that
we employ will vary depending on the difficulty of the task.

#### Training and Evaluation.

During training, samples are generated uniformly from $\mathcal{D}_{s}$. At
test time, we evaluate the generalization performance based on prediction
accuracy over a uniform distribution of numbers up to $l$ digits. An output is
deemed correct only if all digits, including "pad" symbols, are correct;
otherwise, it is considered false. Note that this is worst-case accuracy over
digits but average-case accuracy over $l$-digit numbers. We present results
for different training checkpoints. As noted in [JdDE+23b], out-of-
distribution accuracy can fluctuate even as in-distribution accuracy improves.
Therefore, we show performance across various checkpoints, using a small
validation set with longer sequences to choose the best checkpoint. Further
experimental details are in Appendix B.

###  3.1 Task: Addition

In a transformer model, characterized by parallel information processing, the
common iterative way to perform multi-digit addition is not effective.
Instead, the model must simultaneously identify all relevant positions
contributing to each of output digits.

#### Addition in parallel.

The result of [QB24] showcases that the transformer dissects the calculation
of the sum into parallel streams of computation. Refer to Figure 1, which
presents a parallel method for addition; First, it adds the corresponding
digits across both integers, and then propagates the carries. To compute the
output at position $i$, the model must identify two items: First, the pair of
positions in the input sequence corresponding to the $i$-th digit of the first
and second integers. The second item is a list of all other positions in the
input sequence required to compute the incoming carry-over. Having described
these items, we will now synthesize them into $\sigma(i)$ to denote the list
of all positions required for position $i$. Next, the model learns a function
to compute the output at position $i$ based on the values at these positions
in $\sigma(i)$. For instance, to compute the output at position $i=2$ for the
sample given in Section 3, we have $\sigma(2)=(1,2,22,23)$, pinpointing the
first two digits of both integers. In general, for two $l$-digit integers in
our format, $\sigma$ is written as
$\sigma(i)=\big{(}1,\cdots,i-1,i,l+2,\cdots,i+l,i+l+1\big{)}$. The model then
needs to learn and apply $\sigma(\cdot)$ to any pair of numbers from the
unseen domain $\mathcal{D}_{u}$.

(a)

Algorithm 1 Parallel Carry-Handle Addition

1:$x^{(1)}=(x^{(1)}_{l},\dots,x^{(1)}_{1})$ $\triangleright$ first integer

2:$x^{(2)}=(x^{(2)}_{l},\dots,x^{(2)}_{1})$ $\triangleright$ second integer

3:$s=(s_{l+1},\dots,s_{1})$ $\triangleright$ array of size $l+1$ initialized
to 0

4:for $i\leftarrow 1$ to $l$ do

5: $s_{i}\leftarrow x^{(1)}_{i}+x^{(2)}_{i}$

6:while $\exists i,s_{i}\geq 10$ for some $i$ do

7: for $j\leftarrow 1$ to $l$ do

8: $s^{\prime}_{j+1}\leftarrow s_{j+1}+\lfloor s_{j}/10\rfloor$
$\triangleright$ Adjust one carry

9: for $i\leftarrow 1$ to $l$ do

10: $s_{j}\leftarrow s^{\prime}_{j}\mod 10$ $\triangleright$ Remainder counts

11:return $s$

(b)

Figure 1: Parallel Carry-Handle Addition. (a) An illustration of a parallel
algorithm that handles cascading carries in addition. The algorithm first adds
the corresponding digits of the two integers and then propagates the carries
to the latter positions. (b) Pseudo-code explanation of the process. Note that
the for loop computations can take place in parallel for each position.

#### (A1) Failure of absolute positional encoding.

We first demonstrate that while APE can succeed in generalizing to in-
distribution data, its performance degrades significantly on longer length
sequences. Figure 2(a) shows the performance of our trained model (on the
uniform distribution of numbers with $l_{s}=5$) on a uniform distribution over
domains with lengths varying from 6 to 20. As expected, performance drops
drastically when the length increases from 5 to 6. APE does not capture
translational symmetries in the task of addition. In this regard, prior work
[ABLR23] has shown that it is in general impossible to obtain length
generalization without providing extra information to the model, either by
imposing task symmetries on the model or by using data augmentation in
training. Next, we will see that encoding the translational symmetries through
relative positional encoding (RPE) will lead to generalization on the unseen
positions in the task of addition.

(a)

(b)

(c)

Figure 2: (a) Accuracy of the model using absolute positional encoding while
increasing the length of the sequence. (b) The model with RPE trained with up
to 5-digit sums and tested on up to 50-digit sum. While showing high
fluctuations, the inductive bias of relativity makes generalization possible.
(c) The model with APE when trained with augmented data, i.e., all shifted
versions of samples from the dataset are given in training.

#### (A2) Relative position encoding enables length generalization.

Figure 2(b) confirms that RPE is capable of encoding translational
equivariances for the task of addition. Adding RPE to the Transformers
maintains its high performance at longer lengths. This effectiveness of RPE
was first captured in [JdDE+23b]. Note that the abrupt fall in 50-digit sums
is due to the input format. In all samples seen from $\mathcal{D}_{s}$, the
two integers are separated with a "+" and several "pad" tokens, whereas when
testing on 50-digit sums, there is no "pad" sign that separates them. We
consider this to be a spurious correlation that the model has learned. the
gradual decline in accuracy for longer sequences is a consequence of two
factors. More importantly, the gradual decline in accuracy for longer
sequences is caused by samples with higher complexities that are more likely
among larger integers, according to the definition we give for complexity in
Section 4.

#### (A3) Adding symmetries implicitly via augmented data is not enough.

Knowing that the task has certain symmetries, another potential approach is to
use data augmentation to implicitly encode the symmetries. For the task of
addition, a natural data augmentation scheme consists of shifting a pair of
integers in the training data by adding the same number of zeros from the
right. 111We are leveraging the neutral role of zero in a sum. Note that this
augmentation does not introduce any new nonzero digits, avoiding the need for
additional computations. It also does not append training samples that could
potentially cover the entire unseen domain, which might artificially enhance
performance by unfairly exploiting dependency constraints in the addition task
that are oblivious to relative positional encoding. We trained a BERT model,
again with APE, on this augmented dataset, and the result can be seen in
Figure 2(c). The figure shows that augmentation does not lead to
generalization on the unseen positions. However, we remark that the in-
distribution performance (on shifted test data) is near-perfect (above 0.999),
indicating that it has learned translation by zeros. To understand the failure
of the model trained with augmentation, we examined the model’s output for
10-digit sums. Most errors occurred in the least significant digit, even
though the model performs well on single-digit sums in Figure 2(c). The
detailed explanation of this experiment can be found in Section C.1. Since
augmentation only adds zeros when shifting, causing the model to learn
spurious correlations between the first position and positions larger than
$l_{s}$, which will not hold at test time. We will address this erroneous
learning in second part of Proposition 5.1, where only some parameters are
properly learned during training, while others can grow inadvertently. We show
a similar spurious correlation appearing in our simplified theoretical setup
in Section 5.

###  3.2 Task: Multiplication

Efficient algorithms for multiplication, such as ’shift and add’, cannot
compute multiplication in completely parallel streams. However, a parallel
algorithm is necessary to effectively induce the task structures into a
transformer architecture. Figure 3 proposes a method that, under a limited
length constraint for the multiplier, splits a multi-digit multiplication into
the product of the multiplier with each digit of the multiplicand, and
propagates the carries subsequently. To induce this mechanism, we introduce
_uniform positional encoding_ , where each digit of the multiplier uniformly
affects all digits of the multiplicand.

(d)

Algorithm 2 Parallel Carry-Handle Multiplication

1:$x^{(1)}=(x^{(1)}_{l_{1}},\dots,x^{(1)}_{1})$ $\triangleright$ first integer

2:$x^{(2)}=(x^{(2)}_{l_{2}},\dots,x^{(2)}_{1})$ $\triangleright$ second
integer

3:$m=(m_{l_{1}+l_{2}},\dots,m_{1})$

4:for $i\leftarrow 1$ to $l$ do

5: $m_{i}\leftarrow x^{(1)}*x^{(2)}_{i}$

6:while $\exists i,m_{i}\geq 10$ for some $i$ do

7: for $k\leftarrow 1$ to $l_{1}$ do

8: for $j\leftarrow 1$ to $l_{1}+l_{2}-k$ do

9: $m^{\prime}_{j+k}\leftarrow m_{j+k}+\lfloor m_{j}/10^{k}\rfloor$
$\triangleright$ Handle one carry

10: for $j\leftarrow 1$ to $l_{1}+l_{2}-1$ do

11: $m_{j}\leftarrow m^{\prime}_{j}\mod 10$ $\triangleright$ Remainder counts

12:return $m$

(e)

Figure 3: (a) An illustration of a parallel algorithm that computes the
multiplication, inspired from the uniform functionally of the multiplier in
the absence of carries. Besides, carries still has a relative structure for
propagation to following positions. (b) Pseudo-code explanation of the process
assuming $l_{1}\ll l_{2}$.

#### Uniform Positional Encoding (UPE).

Starting from Equation 2.1, consider a scenario where position $i$ uniformly
affects every other position $j\neq i$. In this case, the positional vector
$p_{j,i}$ must remain constant for every $j$:

$p_{j,i}=p_{j^{\prime},i}=c_{i}\quad\forall j,j^{\prime}\neq i$ (3.3)

Furthermore, if a task involves multiple positions with uniform functionality,
assuming ${i_{1},\dots,i_{m}}$ is the set of such positions, we assign them
the positional vectors $c_{1},\cdots,c_{m}$, respectively. If we denote those
positions in the sequence whose outputs have to be computed as
$\\{j_{1},\cdots,j_{n}\\}$, the computation will need the following collection
of positional vectors: $\\{p_{i,j};i,j\leq n\\}\cup\\{c_{1},\dots,c_{m}\\}$.
Note that additional relationships may exist among ${j_{1},\dots,j_{n}}$, such
as translational properties, which would justify the use of relative
positional encoding. For example, for multiplying a 3-digit number by a
20-digit number, we will use $c_{1},c_{2},c_{3}$ for the multiplier as well as
relative positional vectors $\\{p_{-21},\cdots,p_{21}\\}$ (length of the
answer is at most 22 digits).

Remark. Multiplication is significantly more challenging than addition, even
when the multiplier is a single-digit integer. The model has to deal with
larger stack of carries compared to the addition task. This is because the
probability of a carry occurring is higher at each stage, and unlike in
addition, carries can exceed 1. See Section C.4 for more explanation.
Consequently, we limit our experiments to 3-digit multipliers and focus on the
generalization of the multiplicand. Even in this setting, we require larger
networks compared to the task of addition to generalize.

#### (M1) APE and RPE fail at length generalization.

We evaluated both absolute and relative positional encodings, in a single-
digit multiplier scenario, as provided in Figures 4(a) and 4(b). As
anticipated, APE fails to generalize on the unseen positions, and RPE shows a
decline in accuracy. In the three-digit multiplier scenario, as carries become
more challenging, RPE cannot generalize to longer multiplicands as shown in
Figure 5(a), which is consistent with the results of prior work [JdDE+23b].
RPE is unable to preserve the functionality of the multiplier on those
positions of the multiplicand padded at training.

#### (M2) UPE (along with RPE) is effective length generalization.

In contrast, our UPE imposes this constraint and shows a significant
improvement over RPE in the single-digit multiplier scenario (Figure 4(c)).
This advantage becomes more apparent for the three-digit multiplier (Figure
5(b)); Once we incorporate the three uniform positional vectors from Equation
3.3 to the multiplier’s digits and RPE for the rest, they enable the model to
maintain 90% of its original accuracy at the length of 16. As far as we know,
this is the first positive result for length generalization in multiplication
purely from architectural modifications. For the same reason explained in
Section 3.1 for RPE, the fall in accuracy in 20-digit sums is inevitable
considering our input format.

(a)

(b)

(c)

Figure 4: (a) Accuracy of a BERT model with APE for single-digit $\times$
multi-digit multiplication when trained only up to 5-digit multiplicands. (b)
Same setting as in (a) but with RPE. (c) Same setting as (a) but with our
proposed UPE. Using the uniform symmetry naturally gives advantages over RPE.

#### (M3) Augmentation fails in multiplication.

When some zeros are added to the end of the multiplicand while keeping the
multiplier unchanged, the result is a shifted answer by zeros. This highlights
a symmetry of the problem, aligned with the translational symmetry of RPE for
the multiplicand and the constant functionality of UPE for the multiplier. We
used this to augment the training data and see if it improves results for
models with absolute and relative positional encodings. However, as shown in
Figure 5(c), the model utilizing APE fails similarly to the same attempt in
the addition task. As for the model with RPE, despite extensive training data,
it still falls short of the performance of our UPE. See Figure 5(d).

(a)

(b)

(c)

(d)

Figure 5: (a) Accuracy of a BERT model with RPE for 3-digit multi-digit
multiplication when trained only up to 5-digit multiplicands. (b) Same setting
as in (a) but with our new positional encoding. The difference between UPE and
RPE becomes more apparent as the length of the multiplier increases. (c) The
model using APE with the privilege of augmented data. (d) The model with RPE
and trained with augmented data. Both models with APE and RPE are outperformed
by our UPE.

###  3.3 Extension: Applying to Data with Text + Numbers

Our input format may initially seem limited to numeric experiments. However,
we show how our proposed positional encodings can be effectively applied to a
broader corpus, including text. We propose to set up the pairwise positional
encodings in Equation 2.1 based on the significance of the digits within their
integer; For instance, always the same pairwise positional vector is
incorporated to the first digits of both integers regardless of their
positions in the input. This approach leverages the fact that the structures
in arithmetic tasks like addition and multiplication are grounded on the
number system. As illustrated in Figure 6, we employ separate positional
encodings for text (using APE) and arithmetic calculations (using Pairwise
PE). Since the latter must be implemented at the attention level, specific
attention heads are designated for this purpose. For instance, Figure 6
demonstrates that attention labeled with "Add Headi,j" utilizes RPE. These
heads add RPE only to the pairwise interactions among the numerical parts of
the input, as detailed in Equation 2.1, while the non-numerical parts remain
unaffected by RPE. Note that we do not mask any part of the input sequence,
allowing all attention heads to access the entire input.

Addition and Multiplication in the presence of text. While we include text in
our experiments, attention heads are specifically trained to produce the
correct answers for the arithmetic tasks. Although fully integrating this
method into LLMs requires extensive resources, our current setting is
sufficient to demonstrate its effectiveness and can be easily scaled up in
future work. To test the robustness of our encodings in a more complex
scenario, for both tasks of addition and multiplication, we altered formats in
Sections 3 and 3 to include up 20 random tokens before, between, and following
the two integers, and as previously, the answer is fetched from the output of
the position of the second integer. For the addition task, we modified RPE
such that the same positional vector is always incorporated to the interaction
of the positions of the least significant digits for both integers, and so on
for other positions in spite if their variable distance in the input sequence.
Figure 7(a) shows the method is as effective as the scenario without the text.
During training, attention heads designated for arithmetic operations learn to
ignore the text, focusing solely on arithmetic tasks. In turn, for the
multiplication task, analogous of this approach is applied to equip the model
with UPE, labeled with "Mult Headi,j" in Figure 6, and Figure 7(b) depicts
comparable generalization performance to Figure 5(b).

Figure 6: Extra attention heads utilized with pairwise positional encodings
are integrated into the model to exploit structures across various tasks. The
diagram illustrates how certain attention heads employ relative position
encodings to enable length generalization for the addition task.

(a)

(b)

Figure 7: (a) The generalization performance on the task of addition for a
model that also allows for text with variable length is as good as the
original implementation in Figure 2(b). (b) The trained model with text for
the multiplication task also follows the same generalization trend as in
Figure 5(b).

##  4 Beyond Structural Symmetries: Complexity of Task

So far, we have explored the importance of leveraging the structural
symmetries. However, we may want the model to consistently work for all
examples, not just those that are more probable. Here, we focus on analyzing
the errors of these models, relating these observations to the complexity of
the examples. Specifically, for the task of addition, we will assert that
examples with longer dependencies across the input sequence are considerably
more difficult for the model. Finally, we will show how to counteract this
performance drop by adjusting the training distribution in favor of higher
complexity examples.

Typically, it is not necessary to consider all preceding digits of the two
integers when calculating the result of the current position. Without carry-
overs, the calculation simplifies to repeating single-digit additions across
the sequence in parallel. However, as the chain of carries extends, managing
them grows more challenging in parallel, as the model needs to track more
positions in the input for each output position. We introduce a notion of
complexity based on how far the model should look back to compute the carry-
overs, suggesting that this window may not need to be large for most
instances.

#### Definition (complexity).

For a general seq-to-seq task and a sequence ${\rm S}$ on which the task
operates, we define ’complexity’ as the maximum dependency length, i.e.,
number of positions such that their values are necessary to calculate the
output of the current position, across the entire input.

Consequently, for the task of addition, and within the seen domain that only
contains samples with limited length $l_{s}=5$, the complexity of any sequence
is trivially upper bounded by $10$, the total number of positions. As another
straightforward example, consider a scenario where no two carries occur
consecutively. In this case, the required list of positions, $\sigma(i)$,
would only include $i-1$, $i$, $i+l$, and $i+l+1$.

#### (C1) Consecutive carries are not representative of complexity.

Last example might prompt one to think that the mere consecutive occurrence of
carries qualifies for the long dependency complexity, and will lead to
performance drop for the model. Figure 8(b) shows the performance of the model
for consecutive carries of different lengths. As seen from the figure, the
performance does not drop significantly as the length of the sequence of
consecutive carries increases. This signifies the fact that consecutive
carries do not accurately capture the level of difficulty for this problem.

Now we will quantify the complexity rule for addition. Consider a sample
$(x^{(1)},x^{(2)})$ and the sum at position $i$. If
$x^{(1)}_{i-1}+x^{(2)}_{i-1}\geq 10$, the output at position $i$ is
$x^{(1)}_{i}+x^{(2)}_{i}+1$ regardless of the values of all previous
positions. Similarly, when $x^{(1)}_{i-1}+x^{(2)}_{i-1}\leq 8$, the output at
position $i$ is $x^{(1)}_{i}+x^{(2)}_{i}$. Thus, the output at position $i$
depends on previous positions when the carry cannot be determined without
knowing the prior carries. This is the number of consecutive indices which sum
to 9 that are triggered by a carry, leading to a cascade effect. Figure 1
shows an example where carries cascade for three stages, from the first to the
fourth position.

#### (C2) Cascading carries quantify failure to higher complexity.

In Appendix E, we show that RPE learns a fixed list of relative positions for
computing every digit of the output under a restricted complexity constraint,
which is met during training (within the seen domain that only contains
samples with limited length $l_{s}=5$, the complexity of any sequence is
trivially upper bounded by $10$). Nevertheless, this does not address the
scenario where test samples exceed the constraint. If the model has seen
cascades up to a certain length at training, it may not be able to generalize
to samples with larger cascade sizes. Figure 8(a) demonstrates that an
increase in cascade length causes a significant decline in the performance.
However, as will discuss in the next part and in Section C.4, higher
complexity samples are rare; in fact, most of the cascade lengths are covered
by a small number and thus lower complexity. Consequently, given that our
evaluation are based on the uniform distribution over the unseen domains,
higher complexity samples do not tangibly affect the performance in
expectation.

(a)

(b)

(c)

Figure 8: (a) Cascading carries are noticeably hard samples for the model
trained with RPE. (b) If the cascade effect is not taken into account,
focusing only on a series of carries in a row fails to obtain harder samples.
(c) The performance of the model trained with 100k samples chosen based on
different mixtures of easy and hard examples. Exposure to more complex samples
improves the accuracy on higher complexities.

#### (C3) Models trained with more complex samples generalize better.

In Figure 8(c), in a bid to expose the models to more complex samples during
training, we adjusted the training distributions. Beyond the uniform
distribution across all samples, we trained another model using samples that
are selected equally for their complexity (namely, uniform distribution on the
complexity). We further trained a third model based on a hybrid approach,
where each sample is chosen based on a 50/50 probability from the first two
distributions. These models were then tested on 40-digit sums with varying
cascade lengths. Allowing the model to observe more samples from higher
complexities aids in preserving its accuracy even on cascading lengths above 5
(the maximum length observed in-distribution). Notably, when a sufficient
number of samples from higher complexities are observed (approximately 16k of
samples with a cascade length of 5 in the uniform complexity scenario, and 8k
in the mixed distribution scenario), the uniform complexity loses its
advantage in higher complexity areas. The model trained with the mixed
distribution outperforms the other two at most lengths. We conclude that
smaller cascades are in fact more beneficial for learning the positional
aspect of the task. This aligns with the findings in [SBE+23] about assigning
higher weights to shorter examples. Moreover, samples with higher complexity
are necessary to enhance the performance on higher complexities.

##  5 Theoretical Results

We theoretically validate our findings using a one-layer transformer model
trained on a simple regression task that captures translational symmetries. We
show that in the population regime, i.e., having access to infinite amounts of
data, when the transformer is trained by gradient descent (with weight decay),
encoding these symmetries into the structure of the positional vectors through
RPE will provably lead to length generalization on the unseen positions.
However, in the same setting, the model with APE, which is oblivious to
translational symmetries, will fail to generalize on the unseen positions. We
further show that APE fails even if we use data augmentation to help the model
learn the translational symmetries during training. We defer the proofs to
Appendix D.

#### Task description.

We consider a seq-to-seq regression task consisting of an input sequence
$(x_{1},x_{2},\cdots,x_{n})$ and an output sequence
$(y_{1},y_{2},\cdots,y_{n})$, where for $i\in[n]$:

$y_{i}=\alpha\langle\theta,x_{i}\rangle+\beta\langle\theta,x_{i-1}\rangle+\beta\langle\theta,x_{i+1}\rangle$
(5.1)

Here, $\alpha$ and $\beta$ are some constants, and we assume $\|\theta\|=1$.
The translational symmetries should now be clear from Equation 5.1; if we
apply a (circular) shift to the input, then the output will undergo the same
type of shift.

We consider a length-generalization scenario where, for a given integer
$n_{1}\leq n$, the training set will include sequences
$(x_{1},x_{2},\cdots,x_{n})$ where the values at the first $n_{1}$ positions
are generated according to the Gaussian distribution, i.e. for
$i\in[n_{1}]:x_{i}\stackrel{{\scriptstyle\rm{i.i.d.}}}{{\sim}}\mathcal{N}(0,\mathbb{I}_{d})$,
and the rest are zero-padded, i.e. for $i>n_{1}$
$x_{i}=\bm{0}\in\mathbb{R}^{d}$. At test time, however, the trained model will
be evaluated according to the inputs $(x_{1},x_{2},\cdots,x_{n})$ where the
values at _all_ the positions are generated according to the Gaussian
distribution. We can view the 0 padded positions as the "unseen" positions. To
measure the performance, we consider the MSE loss
$\ell_{i}=(\hat{y}_{i}-y_{i})^{2}$, where $\hat{y}_{i}$ is the model’s
prediction at position $i$. We will train a simplified linear Transformer on
this task:

#### Linear Attention model.

One layer of a transformer architecture, specifically the linear modification
(no Softmax bottleneck), has been deployed as the main tool in some recent
work [vONR+23, MHM23, ACS+24] to analyze the dynamics of training. A Linear
Transformer is defined by the following parameters: three matrices including a
query matrix $W_{Q}\in\mathbb{R}^{d\times d}$, a key matrix
$W_{K}\in\mathbb{R}^{d\times d}$, and a value matrix
$W_{V}\in\mathbb{R}^{d\times d}$. Given the sequence of tokens
$\\{x_{1},\cdots,x_{n}\\}$ the output from the attention head is then passed
through a linear projection $h\in\mathbb{R}^{d}$ to obtain the scalar outputs
of the linear attention:

$\hat{y}_{i}=h^{T}\sum_{r=1}^{n}(W_{V}x_{r})(W_{K}x_{r})^{T}(W_{Q}x_{i})=h^{T}\sum_{j=r}^{n}(W_{V}x_{r})(x_{r}^{T}W_{k}^{T}W_{Q}x_{i})$

To encode the positional information, positional vectors
$\\{p_{1},\cdots,p_{n}\\}$ are added to the tokens: $x_{i}\leftarrow
x_{i}+p_{i}$. Furthermore, for our analysis, we study a simplified one-layer
transformer with factored linear attention, a model commonly studied in prior
work [BTR+21, RGLG23, JSL22]:

$\hat{y}_{i}=h^{T}W_{V}\sum_{r=1}^{n}x_{r}(p_{r}^{T}W_{K}^{T}W_{Q}p_{i})$
(5.2)

In this model, the positional vectors and the input are decoupled. This is to
isolate the attention to only depend on the position, as to be consistent with
the task. This simplification abstracts the learning of position encodings
which is the main focus of our study We also will demonstrate in Section D.3
that a modified version of this model, expressed as:
$\hat{y}_{i}=h^{T}W_{V}\sum_{r=1}^{n}x_{r}(p_{i-r}^{T}W_{K}^{T}W_{Q}u)$, where
$u$ is a constant vector, is the appropriate abstraction for RPE.

For simplicity, we assume that $W_{K}=W_{Q}=W$. Besides, as we aim to study
the structures of the learned positional vectors, we will absorb the attention
weights into the positional vectors without altering the representational
capabilities of these models. This approach is equivalent to fixing the
attention weights in the model, thereby training solely on the positional
encodings. Note that for the model to achieve generalization on all positions,
$h^{T}W_{V}$ must get aligned with $\theta$,
$p_{i}^{T}W^{T}Wp_{i}\propto\alpha$,
$p_{i}^{T}W^{T}Wp_{\\{i-1,i+1\\}}\propto\beta$, and zero otherwise. For RPE,
this similarly requires $p_{0}^{T}W^{T}Wu\propto\alpha$,
$p_{\\{-1,+1\\}}W^{T}Wu\propto\beta$. Having this said, we will prove the
following proposition:

######  Proposition 5.1 (informal).

For the seq-to-seq regression task, the transformer model after training with
GD with infinitesimal weight-decay on the positional vectors in the population
regime:

  * •

(APE) After training we have $p_{k}\approx 0$ for all $k\not\in[n_{1}]$.
Hence, APE fails at generalization.

  * •

(APE with augmentation) For $\beta=0$, after training $|p_{i}^{\top}p_{j}|$ is
large for all $i\neq j$ and $p_{i}^{\top}p_{i}=\alpha$. Hence, APE with
augmentation fails to generalize.

  * •

(RPE) After training $\langle v,\theta\rangle p_{0}^{T}W^{T}Wu=\alpha,\langle
v,\theta\rangle p_{-1}^{T}W^{T}Wu=\langle v,\theta\rangle
p_{1}^{T}W^{T}Wu=\beta$, and $p_{j}\approx 0$ for all $j\not\in\\{-1,0,1\\}$,
thus the model generalizes.

The key idea is that APE does not receive any useful signals on the padded
positions during training, and therefore, does not learn how to predict on
those positions. For APE with augmentation, while the inner products of a
positional vector with its adjacent vectors match the underlying task, those
further apart are unsupervised and thus may behave unpredictably. Finally,
making the model relative through RPE resolves the issue of undesirable
trajectory of inner products occurring for the model with APE, leading to full
generalization. We refer the reader to Appendix D for additional details of
the setup and full proofs.

#### Experimental Validation.

To remain consistent with our assumptions, we conducted an experiment that
positions are on a ring and for the task in Equation 5.1. Here, $n=51$,
$n_{1}=10$, and $d=200$. Figure 9(a) shows that all three settings (APE, APE
with augmentation, and RPE) converge to zero in their training loss. However,
when we compute the test loss at each position in Figure 9(b), that is,
$\operatorname{\mathbb{E}}\\{l_{i}\\}$ according to the normal distribution at
all positions and not only the allowed windows in training, both APE and APE
with augmentation cases cannot generalize.

(a)

(b)

Figure 9: (a) All three scenarios perform perfect in-distribution, when only a
part of the sequence is not padded. At test time, all positions get non-padded
values. (b) The test loss at each position in the sequence. The model with RPE
outperforms the other two. The inherent noise present in APE and APE with
augmentation is a characteristic of their generalization capabilities, The
plots are generated by averaging over 5000 samples at each position.

##  6 Conclusion

Our work provides a general approach for capturing structural symmetries in
arithmetic tasks. This allows us to provide the first length generalization
result for multiplication that relies only on minor architectural
modifications. Exploring how to induce such structural symmetries beyond these
basic arithmetic tasks, and in tandem with text data is a natural next
direction. Future work will explore extending our approach to handle more than
two integers and to combine multiple arithmetic tasks at once. Another
possibility is exploring the compatibility of our methods with text-formatted
arithmetic tasks. Orthogonally, our understanding of the current models is
that, without seeing longer dependencies, the model cannot generalize to more
complex samples. Next steps would be to understand the limitations of
generalization abilities of these models on the complexity axis, and suggest
potential architectural, data, and algorithmic fixes to improve their
abilities.

## References

  * [ABLR23] Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk.  Generalization on the unseen, logic reasoning and degree curriculum, 2023.
  * [ACS+24] Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra.  Linear attention is (maybe) all you need (to understand transformer optimization), 2024.
  * [AWA+22] Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur.  Exploring length generalization in large language models, 2022.
  * [BMR+20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  In Advances in neural information processing systems, 2020.
  * [BTR+21] Nick Bhattacharya, Neil Thomas, Roshan Rao, Justas Daupras, Peter K Koo, David Baker, Yun S. Song, and Sergey Ovchinnikov.  Single layers of attention suffice to predict protein contacts, 2021.
  * [DCLT18] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  arXiv:1810.04805, 2018.
  * [DLS+23] Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, et al.  Faith and fate: Limits of transformers on compositionality.  arXiv preprint arXiv:2305.18654, 2023.
  * [ENO+21] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.  A mathematical framework for transformer circuits.  Transformer Circuits Thread, 2021.
  * [FP88] Jerry A. Fodor and Zenon W. Pylyshyn.  Connectionism and cognitive architecture: A critical analysis.  Cognition, 28(1):3–71, 1988.
  * [GGB20] Mor Geva, Ankit Gupta, and Jonathan Berant.  Injecting numerical reasoning skills into language models, 2020.
  * [HVU+18] Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Ian Simon, Curtis Hawthorne, Andrew M. Dai, Matthew D. Hoffman, Monica Dinculescu, and Douglas Eck.  Music transformer, 2018.
  * [JdDE+23a] Samy Jelassi, Stéphane d’Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton.  Length generalization in arithmetic transformers.  arXiv preprint arXiv:2306.15400, 2023.
  * [JdDE+23b] Samy Jelassi, Stéphane d’Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton.  Length generalization in arithmetic transformers, 2023.
  * [JSL22] Samy Jelassi, Michael E. Sander, and Yuanzhi Li.  Vision transformers provably learn spatial structure, 2022.
  * [KHK+21] Jeonghwan Kim, Giwon Hong, Kyung-min Kim, Junmo Kang, and Sung-Hyon Myaeng.  Have you seen that number? investigating extrapolation in question answering models.  In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 7031–7037, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
  * [KPR+23] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy.  The impact of positional encoding on length generalization in transformers, 2023.
  * [LAD+22] Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al.  Solving quantitative reasoning problems with language models.  In Advances in Neural Information Processing Systems, 2022.
  * [LAG+23a] Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang.  Exposing attention glitches with flip-flop language modeling.  arXiv preprint arXiv:2306.00946, 2023.
  * [LAG+23b] Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang.  Transformers learn shortcuts to automata, 2023.
  * [LH19] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization, 2019.
  * [LM22] Yuxuan Li and James L. McClelland.  Systematic generalization and emergent structures in transformers trained on structured tasks, 2022.
  * [LSL+23] Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos.  Teaching arithmetic to small transformers, 2023.
  * [MHM23] Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma.  One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention, 2023.
  * [MSK+23] Sadegh Mahdavi, Kevin Swersky, Thomas Kipf, Milad Hashemi, Christos Thrampoulidis, and Renjie Liao.  Towards better out-of-distribution generalization of neural algorithmic reasoning tasks, 2023.
  * [NAGA+21a] Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena.  Show your work: Scratchpads for intermediate computation with language models.  arXiv:2112.00114, 2021.
  * [NAGA+21b] Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena.  Show your work: Scratchpads for intermediate computation with language models, 2021.
  * [NJL21] Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin.  Investigating the limitations of transformers with simple arithmetic tasks.  arXiv:2102.13019, 2021.
  * [Ope23] OpenAI.  Gpt-4 technical report.  arXiv preprint arXiv:2303.08774, 2023.
  * [OWJ+22] Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe.  Training language models to follow instructions with human feedback, 2022.
  * [PSL22] Ofir Press, Noah Smith, and Mike Lewis.  Train short, test long: Attention with linear biases enables input length extrapolation.  In International Conference on Learning Representations, 2022.
  * [PTDU16] Ankur P. Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit.  A decomposable attention model for natural language inference, 2016.
  * [QB24] Philip Quirke and Fazl Barez.  Understanding addition in transformers, 2024.
  * [RaGS22] Yasaman Razeghi, Robert L. Logan IV au2, Matt Gardner, and Sameer Singh.  Impact of pretraining term frequencies on few-shot reasoning, 2022.
  * [RDG+23] Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness.  Randomized positional encodings boost length generalization of transformers, 2023.
  * [RGLG23] Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt.  What does self-attention learn from masked language modelling?, 2023.
  * [RSR+23] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer, 2023.
  * [SBE+23] Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang.  Positional description matters for transformers arithmetic, 2023.
  * [SHB16] Rico Sennrich, Barry Haddow, and Alexandra Birch.  Neural machine translation of rare words with subword units.  In Katrin Erk and Noah A. Smith, editors, Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1715–1725, Berlin, Germany, August 2016. Association for Computational Linguistics.
  * [SLP+23] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding, 2023.
  * [SMF+22] Paul Smolensky, R. Thomas McCoy, Roland Fernandez, Matthew Goldrick, and Jianfeng Gao.  Neurocompositional computing: From the central paradox of cognition to a new generation of ai systems, 2022.
  * [SUV18] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani.  Self-attention with relative position representations, 2018.
  * [TPSI21] Avijit Thawani, Jay Pujara, Pedro A. Szekely, and Filip Ilievski.  Representing numbers in nlp: a survey and a vision, 2021.
  * [Ver18] Roman Vershynin.  High-Dimensional Probability: An Introduction with Applications in Data Science.  Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
  * [vONR+23] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.  Transformers learn in-context by gradient descent, 2023.
  * [VSP+17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Advances in neural information processing systems, 2017.
  * [WCZ21] Xin Wang, Yudong Chen, and Wenwu Zhu.  A survey on curriculum learning.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(9):4555–4576, 2021.
  * [WDS+19] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush.  Huggingface’s transformers: State-of-the-art natural language processing.  arXiv:1910.03771, 2019.
  * [WSC+16] Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Lukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean.  Google’s neural machine translation system: Bridging the gap between human and machine translation.  arXiv:1609.08144, 2016.
  * [WTB+22] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al.  Emergent abilities of large language models.  Transactions on Machine Learning Research, 2022.
  * [WWL+19] Eric Wallace, Yizhong Wang, Sujian Li, Sameer Singh, and Matt Gardner.  Do nlp models know numbers? probing numeracy in embeddings, 2019.
  * [WWS+22] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou.  Chain of thought prompting elicits reasoning in large language models.  arXiv:2201.11903, 2022.
  * [WZNZ21] Cunxiang Wang, Boyuan Zheng, Yuchen Niu, and Yue Zhang.  Exploring generalization ability of pretrained language models on arithmetic and logical reasoning, 2021.
  * [YDL+23] Zhen Yang, Ming Ding, Qingsong Lv, Zhihuan Jiang, Zehai He, Yuyi Guo, Jinfeng Bai, and Jie Tang.  Gpt can solve mathematical problems without a calculator, 2023.
  * [ZBL+23] Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran.  What algorithms can transformers learn? a study in length generalization.  arXiv preprint arXiv:2310.16028, 2023.

##  Appendix A Additional Related Work

#### Length generalization.

It has long been known that AI models struggle to fully grasp the underlying
concept from limited experience [FP88, SMF+22, LM22]. For the special case of
transformers, there is ongoing effort to understand to what capacity these
models can learn the simple "algorithms" [LM22, LAG+23b, LAG+23a]. Large
models tend to learn map functions rather than master the underlying tasks,
failing at generalizing to variable lengths [LSL+23]. Even fine-tuning pre-
trained models and scratch-padding only provides marginal benefits [NAGA+21b,
AWA+22]. To mitigate this, recent work has suggested modifications to
architectural components or data formats [JdDE+23a, SBE+23], along with
providing frameworks to argue when generalization may be feasible [ZBL+23] or
infeasible [ABLR23, MSK+23]. In particular, [ZBL+23] show feasibility of
length generalization in Transformers for tasks that use positions only via
certain operations such as comparison, successor, predecessor, etc., and
ignore the role of positional encodings.

#### Solving arithmetic tasks with transformers.

Arithmetic tasks are among the tasks LLMs struggle with. [YDL+23] studies the
performance of GPT-4 [Ope23] at solving arithmetic problems – GPT-4 gets zero
accuracy on 11-digit arithmetic operations, and the best methods have achieved
$40\sim 50\%$ accuracy. In this regard, a growing body of work has focused on
improving the representations of the numbers fed to the model [WSC+16, SHB16,
TPSI21], showing the sub-optimality of the current sub-word methods [WWL+19].
Another line of work has introduced the concept of intermediate steps by using
chain-of thoughts reasoning [WWS+22, OWJ+22], or by gradually increasing the
complexity of tasks [WCZ21, YDL+23]. Despite all these works, transformers
continue to face challenges with high-complexity tasks [DLS+23]. Authors in
[DLS+23] have investigated these problems in detail, and there are many causes
for this including error-propagation, learning shortcut solutions and
insufficient features, etc. It has also been observed that the distribution of
the samples can greatly affect generalization [LSL+23, SBE+23]. For instance,
[SBE+23] states that favoring smaller numbers more that larger numbers leads
to better representations. Moreover, [JdDE+23b] claims that adding a few
samples from the unseen domain can significantly improve generalization for
longer lengths. Here, we maintain the uniform distribution over the data and
contend that this approach requires a large amount of data to better
generalize, but without entailing any longer samples. Our investigation
centers on the two primary arithmetic tasks, addition and multiplication. In
contrary to intermediate step approaches, we tackle this in an end-to-end
manner. Another line of work has focused on the mechanistic interpretability
of transformers [ENO+21] for the task of addition [QB24], trying to understand
the mechanism that a single-layer transformer implements this task in
parallel. In the same vein, we study which input positions are necessary for
each output position when doing the task in parallel, and on this basis, we
build up a new approach that interprets our models.

#### Positional encodings.

In tackling the length generalization problem, various studies have altered
the positional vectors by introducing new structures to them such as
restiveness in [SUV18, HVU+18]. Other work has devised new schemes to leverage
the positional information such as RoPE [SLP+23], and adding scalars to the
attention scores in [RSR+23]. Some works have even suggested removing
positional vectors [KPR+23], or using randomized positional encoding [RDG+23]
for better generalization. Another line of work has focused on adding
structures to attention maps, eliminating the need for positional vectors
[PTDU16, PSL22].

##  Appendix B Experiment Details

For all the experiments we used Huggingface [WDS+19] implementation of Bert
[DCLT18]. Throughout the addition section we used models with the same
specifications: $6$ attention layer, $H=8$ heads, embedding dimension
$d_{z}=768$, and dropout = $0.1$. A model with the same specification was also
employed for the single-digit multiplier setting. In the 3-digit multiplier
scenario, to compensate for the hardness of the task, we increase the number
of attention layers from 6 to 9. We observed that we need more coverage for
this setting , so all the plots in Figure 6 were trained with five times the
amount of data compared to previous settings, i.e. 100k number of samples for
the addition and the single-digit multiplier settings vs 500k number of
samples for the 3-digit multipliers setting. Each element of the input is in
the one-hot format, and the model will print the score for each digit of the
output. All models are trained on the Cross-Entropy loss using AdamW [LH19],
and from the scratch, with batch-size = $64$, learning-rate = $1{\rm e}-4$,
and weight-decays in $\\{3{\rm e}-6,1{\rm e}-5,3{\rm e}-5,1{\rm e}-4\\}$. We
used one NVIDIA RTX A5000 to run our experiments.

##  Appendix C Additional Experiments

###  C.1 Spurious correlations learned with augmentation

As we explained the experiment for Figure 2(c), despite that the model does
well in single-digit sums, it cannot maintain this performance on the first
digit of multi-digit sums. The accuracy on every position of the output is
listed in Section C.1. We conclude that augmentation based on the
translational symmetry leads to spurious correlations between positions, which
is in agreement with our findings in theory described in Section 5.

Test Case | Accuracy
---|---
Shifted integers | 99.7 $\pm$ 0.1
Uniform distribution | 81 $\pm$ 1
Accuracy on the 1st digit | 90 $\pm$ 1
Accuracy on the 2nd digit | 96 $\pm$ 1
Accuracy on the 3th digit | 98.8 $\pm$ 0.1
Accuracy on the 4th digit | 98.0 $\pm$ 0.1
Accuracy on the 5th digit | 98.8 $\pm$ 0.1
Accuracy on 6:11 digits | 96 $\pm$ 1
Table 1: 10-digit sum’s accuracy of the augmented model for in-distribution,
out of distribution (uniform integers in $[1,10^{10}-1]$), and isolating the
performance at every coordinate of the output. Despite the acceptable accuracy
of the model on most digits, the first digit errors hurt the model.

###  C.2 Addition with no carries

One of the experiments we did to establish that the main issue in the addition
task is the existence of carries, was to try a setting with no carry, every
step’s output being independent of former calculations. In this case, addition
is just a digit-wise operation that can be implemented in parallel for a
sequence. Therefore, with our input’s format and relative positional encoding,
the model at each position has to only look at itself and $l+1$ (Addition in
parallel Section), and calculate the sum of those tokens. To this end, Figure
10 approves that this problem does not cause difficulty unless carries are
included.

Figure 10: A relative model easily learns addition where there is no carry.
Here, training samples are up to 5-digit sums.

(a)

(b)

Figure 11: (a) Plotting the pairwise positional vectors for each layer of the
model with RPE in the single-digit multiplier scenario. The main signal is in
the 23rd vector, meaning that the model mainly looks at only one position back
at each layer. (b) The model with UPE in the same setting as (a) exhibits more
interpretable positional maps.

(a)

(b)

Figure 12: (a) Plotting the pairwise positional vectors for each layer of the
model with RPE in the 3-digit multiplier scenario. The model has to look at
more positions in addition to the previous one. (b) The model with UPE in the
same setting as (a) is still more interpretable.

###  C.3 Positional maps

Using weight-decay in training causes unnecessary parts of the network to
vanish over time. Thus, plotting the positional vectors is informative about
what positions the model is looking at. Figures 11(a) and 11(b) are the
positional plots for models with RPE and UPE trained in the single-digit
multiplier scenario (accordingly for models in 4(b) and 4(c)). y-axis accounts
for the index of the plotted vectors: $i-j+(\text{max-length}-1)$ for
$P_{i-j}$ so that in all the plots of 11(a) and 11(b) the 24th is $P_{0}$. And
x represents elements of a positional vector here. Because the attention
dimension is 768 in our setting and the number of heads is 8, positional
vectors are 86-dimensional. We want to highlight one similarity and one major
difference from the plots. In both models and in every layer, the 23rd vector
is the most noticeable, which means looking at the previous position.
Considering our task, this seems very intuitive, as carry comes only from the
previous position in the single-digit multiplier case. The difference is that
attention is less dispersed in UPE, especially in layers 2 to 6, which will
avoid overfitting to in-distribution data. We think this is one of the reasons
for a better out-of-distribution performance of the model with UPE.

We plot the positional vectors for the 3-digit multiplier setting in Figures
12(a) and 12(b). This time max-length = 27, and therefore, $P_{0}$ is the 26th
row of every plot. This time besides $P_{-1}$, $P_{-2}$ and $P_{-3}$ are also
noticeable (this can be seen at layer 2 of 5(b)), which again is reasonable
since carries jump to at most 3 latter positions for a 3-digit multiplier.
Although less conspicuous here, the model with UPE still depicts less
dispersed attention than the RPE.

###  C.4 Distribution of the Complexity

Distribution of cascade carries. Figures 13(a) and 13(b) show the distribution
of the cascade length for randomly sampled 50-digit and 5-digit sums. As can
be seen, the cascade length probability is very small beyond length 4. In
fact, the probability of cascading carries up to length 4 covers 0.998 of the
total probability. Training on 5-digit addition can provide us with cascades
of length up to 5, however, extensive training samples are required for the
model to cover and observe high-complexity samples (the probability of length
4 cascades in random 5 digit numbers is roughly 8.3e-4). This explains why the
performance drops with higher cascade length. Note that despite never seeing
cascades of length higher than 5, the model still gets non-trivial performance
on samples with cascade length more than 5.

#### Complexity of the single-digit multiplier scenario

Unlike addition, there is no direct way to identify samples with higher
complexities in multiplication since the carry’s value is not limited to 1.
However, we can use the fact that length of the maximum dependency in a
sequence is equal to the total number of levels in Figure 3. By the last
level, every position’s value will be a single digit. Figures 14(a) and 14(b)
show the dependency distribution for the 20-digit multiplicand and the 5-digit
multiplicand accordingly, when the multiplier is uniformly sampled from 1 to
9. Similarly to addition, up to length of 4 covers 99.8 percent of the
distribution in Figure 14(a), and the probability of 4 is 1.9e-3 in Figure
14(b). Hence, a sufficient number of training samples from $\mathcal{D}_{s}$
will ensure generalization on the unseen part.

(a)

(b)

Figure 13: (a) Distribution of cascading carries for 50-digit sums. Most of
the mass lies within 1 to 4. (b) Distribution of cascading carries in 5-digit
sums. Increasing the coverage of high-complexity samples of 5-digit sums can
help offset the disparities in distributions.

(a)

(b)

Figure 14: (a) Distribution of dependency for 20-digit multiplicand. (b)
Distribution of dependency for the 5-digit multiplicand scenario.

##  Appendix D Proofs

###  D.1 Proof of Proposition 5.1: APE

After absorbing $W$ into the positional vectors, i.e. $p_{i}\leftarrow
Wp_{i}$, and letting $v:=W_{V}^{T}h$, we obtain:

$\hat{y}_{i}=\sum_{r=1}^{n}p_{i}^{T}p_{r}\langle v,x_{r}\rangle$ (D.1)

During training only the first $n_{1}$ positions are variable. Thus, the total
loss function at every step is:

$\ell(\\{p_{1},\cdots,p_{n}\\})=\sum_{i=1}^{n_{1}}(\hat{y}_{i}-y_{i})^{2}$
(D.2)

Before diving into the proofs, it is important to note that we write the
equations on a ring, where index "-1" points at the last element in the
sequence; i.e. index $n$. This makes it simpler to handle the boundary
conditions.

To make the convergence problem easier, we set $v=\theta$ at initialization
such that it does not need to explore any other direction. This alignment
assumption only holds for our proofs of negation for absolute positional
encoding (APE) and APE with augmentation, and the extra assumption makes our
proof of negation even stronger. We will prove with all of these
simplifications that APE will not still learn the structure for padded
positions, hence failing to generalize on the unseen.

######  Proposition (Formal statement of Proposition 5.1– First part).

The solution found by the gradient descent with infinitesimal weight-decay in
the population regime on the loss function (D.2) satisfies $p_{k}=0$ for all
$k\not\in[n_{1}]$. Therefore, the expected loss at each position, i.e.,
$\operatorname{\mathbb{E}}\\{\ell_{i}\\}=\operatorname{\mathbb{E}}\\{(\hat{y}_{i}-y_{i})^{2}\\}$
for $i>n_{1}$ converges to $\alpha^{2}+2\beta^{2}$.

###### Proof.

$\displaystyle\ell(P):$ $\displaystyle=\ell(\\{p_{1},\cdots,p_{n}\\})$
$\displaystyle=\sum_{i=1}^{n_{1}}(\hat{y}-y)^{2}$
$\displaystyle=\sum_{i=1}^{n_{1}}\big{(}\sum_{r=1}^{n_{1}}p_{i}^{T}p_{r}\langle\theta,x_{r}\rangle+\sum_{r=n_{1}+1}^{n}p_{i}^{T}p_{r}\langle\theta,\bm{0}\rangle-\alpha\langle\theta,x_{i}\rangle-\beta\langle\theta,x_{i-1}\rangle-\beta\langle\theta,x_{i+1}\rangle\big{)}^{2}$
$\displaystyle=\sum_{i=1}^{n_{1}}\big{(}\sum_{r=1}^{n_{1}}p_{i}^{T}p_{r}\langle\theta,x_{r}\rangle-\alpha\langle\theta,x_{i}\rangle-\beta\langle\theta,x_{i-1}\rangle-\beta\langle\theta,x_{i+1}\rangle\big{)}^{2}$

In the second equality we have separated the contribution of padded elements
from the rest. Because $x_{r}\sim\mathcal{N}(0,I_{d})$ and $\|\theta\|=1$
their inner product is a normal random variable that we call it $s_{r}$ for
$r\leq n_{1}$. Taking the gradient of the loss function with respect to an
arbitrary vector $p_{k}$, we will have:

$\displaystyle\nabla_{p_{k}}=$
$\displaystyle\sum_{i=1}^{n_{1}}2\big{(}\sum_{r=1}^{n_{1}}p_{i}^{T}p_{r}s_{r}-\alpha
s_{i}-\beta s_{i-1}\mathbbm{1}(i\neq 1)-\beta s_{i+1}\mathbbm{1}(i\neq
n_{1})\big{)}$
$\displaystyle\quad\times\Big{(}\mathbbm{1}(k\in[n1])\big{(}\mathbbm{1}(k=i)(2p_{k}s_{k}+\sum_{r\neq
k}^{n_{1}}p_{r}s_{r})+\mathbbm{1}(k\neq
i)(p_{i}s_{k})\big{)}+\mathbbm{1}(k\notin[n_{1}])\big{(}0)\Big{)}$

In which $\mathbbm{1}$ function gives 1 if inside condition holds, 0
otherwise. Therefore the gradient for $k\notin[n_{1}]$ is zero. As a result,
the updating rule under SGD with weight-decay rate equal to $\epsilon$ for
vectors outside the window is:

$p_{k}^{(t+1)}=p_{k}^{(t)}-\epsilon p_{k}^{(t)}$

And for $\epsilon$ being sufficiently small, we will have:

$p_{k}\to_{t}\bm{0}\quad\forall k\notin[n_{1}]$

Meanwhile generalization over elements outside of $[n_{1}]$ requires
$p_{k}^{T}p_{k}\to\alpha$ simply by matching the form of $\hat{y}$ to $y$.
Thus, APE cannot learn the structure outside of the seen window. ∎

###  D.2 Proof of Proposition 5.1: APE + Augmentation

Consider the same task (5.1) and the model (5.2). However, we consider
training with data augmentation in which we apply translational shifts to the
training data. Before, the training data had Gaussian-generated values in the
first $n_{1}$, and the other positions are zero-padded. Here, for each
$i\in[n]$ we shift the training data by $i$ positions. The shifted sequences
will have Gaussian values at positions $i+1,i+2,\cdots,i+n_{1}$ and are zero-
padded elsewhere (note that we’re considering circular shifts). The shifted
data will then be augmented to the training set. The augmented data is thus
generated according to the following distribution: For each $i\in[n]$:

$\begin{cases}x_{r}\stackrel{{\scriptstyle\rm{i.i.d.}}}{{\sim}}\mathcal{N}(0,\mathbb{I}_{d})\quad&\text{for}\quad
i+1\leq r\leq i+n_{1},\\\
x_{r}=\bm{0}_{d}&r\notin\\{i+1,\cdots,i+n_{1}\\}.\end{cases}$

The overall loss function for the case where data is shifted by $i$ positions
is:

$\ell_{i}(\\{p_{1},\cdots,p_{n}\\})=\sum_{t=1}^{n_{1}}(\hat{y}_{i+t}-y_{i+t})^{2}$
(D.3)

By summing over all the shifts, the total augmented loss becomes:

$\ell(P):=\ell(\\{p_{1},\cdots,p_{n}\\})=\sum_{i=1}^{n}\sum_{t=1}^{n_{1}}(\hat{y}_{i+t}-y_{i+t})^{2}$
(D.4)

######  Proposition (Formal statement of Proposition 5.1– Second part).

The solution found by the gradient flow with infinitesimal weight-decay in the
population regime on the loss function (D.4) satisfies
$\operatorname{\mathbb{E}}\\{\ell_{i}\\}=\operatorname{\mathbb{E}}\\{(\hat{y}_{i}-y_{i})^{2}\\}=\Omega(n/d)$
when
$x_{i}\stackrel{{\scriptstyle\rm{i.i.d.}}}{{\sim}}\mathcal{N}(0,\mathbb{I}_{d})\
\forall i\in[n]$.

###### Proof.

Inserting $y$ and $\hat{y}$ into this will give:

$\displaystyle\ell(P)$
$\displaystyle=\sum_{i=1}^{n}\sum_{t=-w}^{w}\big{(}\sum_{r\in
W_{i}}p_{i+t}^{T}p_{r}\langle\theta,x_{r}\rangle-\alpha\langle\theta,x_{i+t}\rangle-\beta\langle\theta,x_{i+t-1}\rangle\mathbbm{1}(t\neq-w)$
$\displaystyle\hskip
187.90244pt-\beta\langle\theta,x_{i+t+1}\rangle\mathbbm{1}(t\neq
w)\big{)}^{2}$

Borrowing our definition from last part $s_{i}=\langle\theta,x_{i}\rangle$ for
those positions that are not padded, we compute the expected gradients:

$\displaystyle\operatorname{\mathbb{E}}\\{\nabla_{p_{k}}\\}=$
$\displaystyle\operatorname{\mathbb{E}}\Big{\\{}\sum_{i=1}^{n}\sum_{t=-w}^{w}2\Big{(}\sum_{r\in
W_{i}}p_{i+t}^{T}p_{r}s_{r}-\alpha s_{i+t}-\beta
s_{i+t-1}\mathbbm{1}(t\neq-w)-\beta s_{i+t+1}\mathbbm{1}(t\neq w)\Big{)}$
$\displaystyle\qquad\qquad\quad\ \times\Big{(}\mathbbm{1}(k\in
W_{i})\big{(}\mathbbm{1}(k=i+t)(2p_{k}s_{k}+\sum_{r\in W_{i},r\neq
k}p_{r}s_{r})$ $\displaystyle\hskip 68.65627pt+\mathbbm{1}(k\neq
i+t)(p_{i+t}s_{k})\big{)}\Big{)}\Big{\\}}$ $\displaystyle=$ $\displaystyle
2\sum_{i=1}^{n}\sum_{t=-w}^{w}\Big{\\{}\mathbbm{1}(k\in
W_{i})\big{\\{}\mathbbm{1}(k=i+t)\big{(}2(p_{k}^{T}p_{k})p_{k}-2\alpha
p_{k}+\sum_{r\in W_{i},r\neq k}(p_{k}^{T}p_{r})p_{r}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad-\beta
p_{k+1}\mathbbm{1}(t\neq w)-\beta p_{k-1}\mathbbm{1}(t\neq-w)\big{)}$
$\displaystyle\hskip 93.95122pt+\mathbbm{1}(k\neq
i+t)\big{(}(p_{i+t}^{T}p_{k})p_{i+t}-\beta\mathbbm{1}(k=i+t-1)\mathbbm{1}(t\neq
w)p_{k+1}$ $\displaystyle\hskip
216.08769pt-\beta\mathbbm{1}(k=i+t+1)\mathbbm{1}(t\neq-w)p_{k-1}\big{)}\big{\\}}\Big{\\}}$
$\displaystyle=$ $\displaystyle 2\sum_{i=1}^{n}\Big{\\{}\mathbbm{1}(k\in
W_{i})\big{\\{}2\sum_{r\in W_{i}}(p_{k}^{T}p_{r})p_{r}-2\alpha
p_{k}-\mathbbm{1}(k\neq i+w)2\beta p_{k+1}$ $\displaystyle\hskip
178.50697pt-\mathbbm{1}(k\neq i-w)2\beta p_{k-1}\big{\\}}\Big{\\}}$

In the last line, we have taken the some over $t$ inside, where for each $k\in
W_{i}$ there is only one $t$ that $k=i+t$, and at most one such that
$k=i+t-1$, or $k=i+t+1$. To apply the sum over $i$ on the first component of
the sum, we have to count the total number of times that $k$ and $r$ fall on
the window when varying $i$. The answer is $(2w+1-|i-r|)$ for $|i-r|\leq 2w+1$
and zero otherwise. For other terms, every k happens $(2w+1)$ times,
therefore:

$\operatorname{\mathbb{E}}\\{\nabla_{p_{k}}\\}=4\sum_{r=-2w}^{2w}(2w+1-|r|)(p_{k}^{T}p_{k+r})p_{k+r}-4(2w+1)\alpha
p_{k}-8w\beta p_{k+1}-8w\beta p_{k-1}$

We insert this gradient in the gradient flow equation with weight-decay:

$\displaystyle\frac{d}{dt}p_{k}=-\nabla_{p_{k}}-\epsilon p_{k}=$
$\displaystyle-4\sum_{r=-2w}^{2w}(2w+1-|r|)(p_{k}^{T}p_{k+r})p_{k+r}+(4(2w+1)\alpha-\epsilon)p_{k}$
$\displaystyle+8w\beta p_{k+1}+8w\beta p_{k-1}$ (D.5)

Since $\epsilon$ is small and it appears in a component with non-zero factor,
we can neglect it.

It is not where each positional vector converges to that controls the
generalization, but rather their inner products that show up in (D.1). Let’s
define $A_{k,l}:=p_{k}^{T}p_{l}$, we compute how elements of $A$ evolve over
time:

$\displaystyle\frac{d}{dt}A_{k,l}=$
$\displaystyle(\frac{d}{dt}p_{k})^{T}p_{l}+p_{k}(\frac{d}{dt}p_{l})=-\\{(\nabla_{p_{k}})^{T}p_{l}+p_{k}^{T}\nabla_{p_{l}}\\}$
$\displaystyle=$
$\displaystyle-4\sum_{r=-2w}^{2w}(2w+1-|r|)\big{(}(p_{k}^{T}p_{k+r})(p_{l}^{T}p_{k+r})+(p_{l}^{T}p_{l+r})(p_{k}^{T}p_{l+r})\big{)}+8(2w+1)\alpha
p_{k}^{T}p_{l}$
$\displaystyle+8w\beta(p_{l}^{T}p_{k-1}+p_{l}^{T}p_{k+1}+p_{k}^{T}p_{l-1}+p_{k}^{T}p_{l+1})$
$\displaystyle=$ $\displaystyle
4\sum_{r=-2w}^{2w}(2w+1-|r|)\big{(}A_{k,k+r}A_{l,k+r}+A_{l,l+r}A_{k,l+r})+8(2w+1)\alpha
A_{k,l}$ $\displaystyle+8w\beta(A_{l,k-1}+A_{l,k+1}+A_{k,l-1}+A_{k,l+1})$
(D.6)

This means that elements of $A$ determine their own dynamics, and it suffices
to track them only. Next step is to put assumptions on the initialization of
positional vectors. Final model must acquire the same translational symmetries
as the underlying task for it to generalize at test time. We seed these
conditions at initialization to put it on the right track, and we will show
that even in this case off diagonal elements of A do not converge to desired
values. Applying the translational invariance to initial inner products of
$p_{k}$’s:

$p_{k}^{(0)T}p_{k+j}^{(0)}=p_{l}^{(0)T}p_{l+j}^{(0)}\quad\forall k,l,j$ (D.7)

######  Lemma D.1.

Condition (D.7) holds for all $t$ if positional vectors start under condition
(D.7) and evolve by (D.2).

###### Proof.

In other words the lemma says $A_{k,k+j}=A_{l,l+j}$. By induction, it only
suffices to show if condition holds at one time t, their derivatives are
equal. Using equation (D.2):

$\displaystyle\frac{d}{dt}A_{k,k+j}=$
$\displaystyle-4\sum_{r=-2w}^{2w}(2w+1-|r|)\big{(}A_{k,k+r}A_{k+j,k+r}+A_{k+j,k+j+r}A_{k,k+j+r})$
$\displaystyle+8(2w+1)\alpha
A_{k,k+j}+8w\beta(A_{k+j,k-1}+A_{k+j,k+1}+A_{k,k+j-1}+A_{k,k+j+1})$

In every term of above expression $k$ appears on both indices of A, hence it
can be replaced by l everywhere under condition (D.7). Thus:

$\frac{d}{dt}A_{k,k+j}=\frac{d}{dt}A_{l,l+j}$

∎

As a result, we can address the elements of A by the difference of its two
indices:

$A_{j}\leftarrow A_{l,l+j}$

$A_{0}$ stands for all diagonal elements of $A$, and $A_{j}$ for $j\geq 1$
addresses off diagonal elements with respect to their distance from the main
diagonal. Note that due to the circular constraint, the number of independent
variables is equal to $\lfloor\frac{n}{2}\rfloor$, and
$A_{\lfloor\frac{n}{2}\rfloor+1}=A_{\lfloor\frac{n-1}{2}\rfloor}$. This brings
us to a system of coupled differential equations for elements of A:

$\displaystyle\frac{d}{dt}A_{j}=$
$\displaystyle-4\sum_{r=-2w}^{2w}(2w+1-|r|)A_{r}\big{(}A_{r-j}+A_{j+r}\big{)}+8(2w+1)\alpha
A_{j}$ $\displaystyle+8w\beta(A_{-j-1}+A_{1-j}+A_{j+1}+A_{j-1})$

Using the fact that $A_{j}=A_{-j}$:

$\displaystyle\frac{d}{dt}A_{j}=8(2w+1)(\alpha-
A_{0})A_{j}-8\sum_{r=1}^{2w}(2w+1-|r|)A_{r}\big{(}A_{j-r}+A_{j+r}\big{)}+16w\beta(A_{j+1}+A_{j-1})$
$\displaystyle\quad j=0,1,\dots,\lfloor\frac{n}{2}\rfloor$ (D.8)

We will go through this system case by case:

Case 1. $w=0$, $\beta=0$
This means for each sample only one element is supervised and the rest of
elements are zero in that sample. The most basic task we can consider in this
setup is one where the output of each position is solely dependent on its own
input. So (D.2) simplifies to:

$\frac{d}{dt}A_{j}=8(\alpha-A_{0})A_{j}$ (D.9)

Writing it for $j=0$:

$\frac{d}{dt}A_{0}=8(\alpha-A_{0})A_{0}$

Which is a non-linear ODE with the following solution:

$A_{0}^{(t)}=\frac{\alpha}{1+(\frac{\alpha}{A_{0}^{(0)}}-1)\exp{(-8\alpha
t)}}$

$A_{0}$ converges monotonically to $\alpha$, either from below or above
depending on the values of $\alpha$. If $\alpha>A_{0}^{(0)}$, which holds with
high probability if we choose $\alpha>1$ and initialize $p_{k}$ with an
isotropic Gaussian vector [Ver18], then $a-A_{0}^{(t)}$ is always greater than
0, and according to equation (D.9):

$\begin{cases}\frac{d}{dt}A_{j}>0\quad&A_{j}>0,\\\
\frac{d}{dt}A_{j}<0\quad&A_{j}<0.\\\ \end{cases}$

Absolute value of all off diagonal elements increases over iterations. So if
at initialization $A_{j}=\theta(\frac{1}{\sqrt{d}})$, the final test loss for
position $i$ is:

$\displaystyle\operatorname{\mathbb{E}}\\{\ell_{i}\\}=$
$\displaystyle\operatorname{\mathbb{E}}\\{\big{(}\sum_{j=1}^{n}(p_{i}^{T}p_{j})s_{j}-\alpha
s_{j}\big{)}^{2}\\}=\operatorname{\mathbb{E}}\\{\big{(}\sum_{j=1}^{n}(A_{i-j})s_{j}-\alpha
s_{j}\big{)}^{2}\\}=\operatorname{\mathbb{E}}\\{\big{(}\sum_{j=1,j\neq
i}^{n}(A_{i-j})s_{j}\big{)}^{2}\\}$ $\displaystyle=$
$\displaystyle\sum_{j=1}^{n-1}|A_{r}|^{2}=\Omega(\frac{n}{d})>0$ (D.10)

∎

To sum up, the use of APE in a model does not enable generalization, even when
all samples in the training domain are augmented.

###  D.3 Proof of Proposition 5.1: RPE

We consider the same task as in (5.1), and investigate a factored attention
model similar to (5.2), with the difference that positional vectors are
relatively structured according to RPE. Note that the positional vectors are
applied only to the key (and this is consistent with the models considered in
(2.1)). The query does not include positional information and must be treated
as a constant vector in this scenario. We denote this constant vector by $u$.
The model thus becomes 222Note that if we had used $p_{i-r}$ instead of $u$,
it would have made no difference in our as the effective parameter is the
same. This becomes clear later on.:

$\hat{y}_{i}=h^{T}W_{V}\sum_{r=1}^{n}x_{r}(p_{i-r}^{T}W_{K}^{T}W_{Q}u)$ (D.11)

The formula presented here involves indices of $p$ from $-(n-1)$ to $(n-1)$.
Only a subset of these indices are used at each position. For example, for
position 1, only $p_{0}$ to $p_{-n_{1}+1}$ are considered. Absorbing the key
and query matrices into p and u, and letting $v=W_{V}^{T}h$, we obtain the
following equation:

$\hat{y}_{i}=\sum_{r=1}^{n}(p_{i-r}^{T}u)\langle v,x_{r}\rangle$ (D.12)

Consequently, we introduce a set of scalars that represent the model’s
parameters effectively: $a_{i-r}=p_{i-r}^{T}u$. The total loss function is:

$\displaystyle\ell(a,v)=$
$\displaystyle\sum_{i=1}^{n_{1}}(\hat{y}_{i}-y_{i})^{2}=\sum_{i=1}^{n_{1}}\big{(}\sum_{r=1}^{n_{1}}a_{i-r}\langle
v,x_{r}\rangle$
$\displaystyle-\alpha\langle\theta,x_{i}\rangle-\beta\langle\theta,x_{i-1}\rangle-\beta\langle\theta,x_{i+1}\rangle\big{)}^{2}$
(D.13)

At the first step, we prove that $v$ gets aligned with $\theta$ in spite of
convergence for positional vectors:

######  Lemma D.2.

For any initialization of v and any value for $p$ in equation D.12, running
SGD with small enough step-size and infinitesimal $\epsilon$ on optimization
problem (D.3) will align $v$ to $\theta$.

###### Proof.

Taking the gradient of loss function in (D.3) with respect to $v$:

$\displaystyle\nabla_{v}=$
$\displaystyle\sum_{i=1}^{n_{1}}2\big{(}\sum_{r=1}^{n_{1}}a_{i-r}v^{T}x_{r}-\alpha\theta^{T}x_{i}-\beta\theta^{T}x_{i-1}\mathbbm{1}(i\neq
1)-\beta\theta^{T}x_{i+1}\mathbbm{1}(i\neq n_{1})\big{)}$
$\displaystyle\quad\times\big{(}\sum_{r=1}^{n_{1}}a_{i-r}x_{r}\big{)}$

And in expectation:

$\displaystyle\operatorname{\mathbb{E}}\\{\nabla_{v}\\}=$
$\displaystyle\sum_{i=1}^{n_{1}}2\big{(}\sum_{r=1}^{n_{1}}a_{i-r}^{2}v-\alpha
a_{0}\theta-\beta a_{1}\theta\mathbbm{1}(i\neq 1)-\beta
a_{-1}\theta\mathbbm{1}(i\neq n_{1})\big{)}$ $\displaystyle=$
$\displaystyle(2\sum_{i=1}^{n_{1}}\sum_{r=1}^{n_{1}}a_{i-r}^{2})v-(2\alpha
n_{1}a_{0}+2\beta(n_{1}-1)(a_{1}+a_{-1}))\theta$

Decomposing $v$ to one aligned vector with $\theta$, and another orthogonal
element:

$v=v_{\parallel}+v_{\perp}$

If the step-size is $\eta$ and weight-decay equal to $\epsilon$, then the
updating rule for $v_{\perp}$ will be:

$v_{\perp}^{(t+1)}=v_{\perp}^{(t)}-\eta(2\sum_{i=1}^{n_{1}}\sum_{r=1}^{n_{1}}a_{i-r}^{2})v_{\perp}^{(t)}-\epsilon
v_{\perp}^{(t)}$

$\epsilon$ is sufficiently small, if we take $\eta$ small enough such that
$\eta(2\sum_{i=1}^{n_{1}}\sum_{r=1}^{n_{1}}a_{i-r}^{2})$ is smaller than 1 at
all times, the orthogonal term vanishes over time. ∎

After demonstrating that v becomes aligned, we will move its norm into
$a_{i}$’s, resulting in $\ v\ =1$ and $v=\theta$. It is time to concentrate on
positional encodings:

######  Proposition (Formal statement of Proposition 5.1– Third part).

The solution found by the gradient descent with infinitesimal weight-decay in
the population regime on the loss function (D.3) will align $v$ with $\theta$
and satisfies $\langle v,\theta\rangle a_{0}=\alpha,\langle v,\theta\rangle
a_{1}=\langle v,\theta\rangle a_{-1}=\beta$, and $a_{j}=0$ for all
$j\not\in\\{-1,0,1\\}$, and thus
$\operatorname{\mathbb{E}}\\{\ell_{i}\\}=\operatorname{\mathbb{E}}\\{(\hat{y}_{i}-y_{i})^{2}\\}=0$
when
$x_{i}\stackrel{{\scriptstyle\rm{i.i.d.}}}{{\sim}}\mathcal{N}(0,\mathbb{I}_{d})\
\forall i\in[n]$.

###### Proof.

Let’s compute the gradient of (D.3) with respect to a’s. Letting
$s_{r}=\theta^{T}x_{r}$, which is a normal random variable:

$\displaystyle\nabla_{a_{k}}=$
$\displaystyle\sum_{i=1}^{n_{1}}2\big{(}\sum_{r=1}^{n_{1}}a_{i-r}s_{r}-\alpha
s_{i}-\beta s_{i-1}\mathbbm{1}(i\neq 1)-\beta s_{i+1}\mathbbm{1}(i\neq
n_{1})\big{)}$
$\displaystyle\quad\times\big{(}\mathbbm{1}(k\in[i-1,i-n_{1}]s_{i-k})$

So that in expectation:

$\displaystyle\operatorname{\mathbb{E}}\\{\nabla_{a_{k}}\\}=$
$\displaystyle\sum_{i=1}^{n_{1}}2\mathbbm{1}(k\in[i-1,i-n_{1}])\big{(}a_{k}-\alpha\mathbbm{1}(k=0)-\beta\mathbbm{1}(i\neq
1)\mathbbm{1}(k=1)$ $\displaystyle\hskip 122.85876pt-\beta\mathbbm{1}(i\neq
n_{1})\mathbbm{1}(k=-1)\big{)}$ $\displaystyle=$ $\displaystyle
2(n_{1}-|k|)a_{k}-2n_{1}\alpha\mathbbm{1}(k=0)-2(n_{1}-1)\beta\mathbbm{1}(k=1)-2(n_{1}-1)\beta\mathbbm{1}(k=-1)$

Where in the last line to calculate the sum, we use the fact that $a_{k}$ is
observed $n_{1}-|k|$ times when traversing from $1$ to $n_{1}$. For values of
$k$ outside the range $[-n_{1},n_{1}]$, the gradient is zero. Hence with
weight-decay they will diminish over time. For values of $k$ between $-n_{1}$
and $n_{1}$, excluding -1, 0, or 1, the update equation is given by:

$a_{k}^{(t+1)}=a_{k}^{(t)}-2\eta(n_{1}-|k|)a_{k}^{(t)}-\epsilon a_{k}^{(t)}$

An again, for small values for the step-size and weight-decay they converge to
zero:

$a_{k}\to_{t}\bm{0}\quad\forall k\notin[-1,1]$

The other three variable go on a different route:

$\begin{cases}a_{0}^{(t+1)}=a_{0}^{(t)}-2\eta
n_{1}(a_{0}^{(t)}-\alpha)-\epsilon a_{0}^{(t)},\\\
a_{1}^{(t+1)}=a_{1}^{(t)}-2\eta(n_{1}-1)(a_{1}^{(t)}-\beta)-\epsilon
a_{1}^{(t)},\\\
a_{-1}^{(t+1)}=a_{-1}^{(t)}-2\eta(n_{1}-1)(a_{-1}^{(t)}-\beta)-\epsilon
a_{-1}^{(t)}.\end{cases}$

For infinitesimal values of weigh-decay and small step-size, they accordingly
converge to $\alpha$, $\beta$ and $\beta$. So at test time:

$\displaystyle\operatorname{\mathbb{E}}\\{\ell_{i}\\}=$
$\displaystyle\operatorname{\mathbb{E}}\\{\big{(}\sum_{j=1}^{n}(a_{i-j})s_{j}-\alpha
s_{j}-\beta s_{i-1}-\alpha\beta s_{i+1}\big{)}^{2}\\}=0$ (D.14)

Using RPE the model will be able to achieve zero test loss. ∎

###  D.4 Experiments for the linear attention model

#### Experiments for a more general setting.

We conducted another experiment for a symmetric task more generalized than
Equation 5.1 (here the dependency length is $m$ instead of 1), establishing
that out claims hold beyond the scope of our assumptions:

$y_{i}=\alpha\langle\theta,x_{i}\rangle+\beta_{1}\langle\theta,x_{i-1}\rangle+\beta_{1}\langle\theta,x_{i+1}\rangle+\dots+\beta_{m}\langle\theta,x_{i-m}\rangle+\beta_{m}\langle\theta,x_{i+m}\rangle$

For the result of Figure 15(a), we set $n=51$, $n_{1}=20$, $m=5$, and $d=200$.
As previously seen in Figure 9(b), although APE and APE with augmentation
achieves zero training loss, they fails to generalize when all positions are
non-padded. The causes of their failure are explained in Section D.1 and
Section D.2. We have plotted $[p_{i}^{T}W^{T}Wp_{j}]_{i\leq n,j\leq n}$ for
APE and APE with augmentation in Figures 16(a) and 16(b) accordingly. For
$i,j>n_{1}$, APE does not learn the structure outside the focus window as
described in Section D.1. But augmentation fails for a different reason, and
Figure 16(b) shows that near diagonal elements of $[p_{i}^{T}W^{T}Wp_{j}]$ are
learned properly while the rest move in random directions and ruin the
generalization.

Figure 15(b) illustrates that we can address this issue with the help of RPE,
here unwanted elements that keep us from generalizing in APE have been
eliminated automatically.

(a)

(b)

Figure 15: (a) Similarly to Figure 9(b), all three scenarios reach zero in the
training loss. At test time, all position get non-padded values. (b) The
generalization loss at each position in the sequence. Again, the model with
RPE greatly outperforms the other two.

(a)

(b)

(c)

Figure 16: (a) The set of effective parameters are
$[p_{i}^{T}W^{T}Wp_{j}]_{i\leq n,j\leq n}$. The model with APE does not learn
the structure for $i,j>n_{1}$. (b) Even though augmentation helps to learn the
main structure of the task every where, the two off-diagonal sides remain
nonzero and damage the generalization. (c) Shows $p_{i-j}^{T}W_{K}^{T}W_{Q}u$
where the index on this graph represents $i-j+(n-1)/2$. Note that
$\frac{1}{2}$ factor appears due to the ring condition, e.g. $i=1,j=n$ will
have $p_{-1}$ associated to it.

##  Appendix E Fixed map for RPE

If the maximum complexity across the training dataset is $2*d$ (Note that
$d-1$ stands for the maximum number of levels illustrated in Figure 1 among
all training samples), the trained model will ignore those previous positions
of both integers that do not contribute the current calculation. i.e., with
our notation defined in Section 3.1 for the tracked list of positions:

$\sigma_{\rm
training}(i)=\big{(}i-d,\cdots,i-1,i,i+l-d+1,\cdots,i+l,i+l+1\big{)}$

Thus, for a model utilized with RPE, the previous list will be memorized in
terms of relative positions:

$\sigma_{\rm training}^{\rm
rel}(i)=\big{(}-d,\cdots,-1,0,l-d+1,\cdots,l,l+1\big{)}$

As a result, $\sigma_{\rm training}^{\rm rel}(i)$ is the same list for every
$i$. Hence, if the model equipped with RPE learns to place its attention on
correct relative positions for in-distribution samples, it will naturally be
able to generalize to longer integers as long as the complexity of the samples
is limited.