Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure
Authors: Hanseul Cho
Jaeyoung Cha
Pranjal Awasthi
Srinadh Bhojanapalli
Anupam Gupta
Chulhee Yun
Venue: NeurIPS 2024 poster (2024)
Keywords: Length Generalization
Transformers
Position Coupling
Positional Encoding
Out-of-distribution Generalization
Arithmetic Tasks
Algorithmic Tasks
Primary area: natural_language_processing
TL;DR: To tackle the length generalization problem of decoder-only Transformer for solving arithmetic/algorithmic tasks, we inject the structure of the task into the Transformer by using the same position IDs for relevant tokens.

Abstract: Even for simple arithmetic tasks like integer addition, it is challenging for Transformers to generalize to longer sequences than those encountered during training. To tackle this problem, we propose *position coupling*, a simple yet effective method that directly embeds the structure of the tasks into the positional encoding of a (decoder-only) Transformer. Taking a departure from the vanilla absolute position mechanism assigning unique position IDs to each of the tokens, we assign the same position IDs to two or more "relevant" tokens; for integer addition tasks, we regard digits of the same significance as in the same position. On the empirical side, we show that with the proposed position coupling, our models trained on 1 to 30-digit additions can generalize up to *200-digit* additions (6.67x of the trained length). On the theoretical side, we prove that a 1-layer Transformer with coupled positions can solve the addition task involving exponentially many digits, whereas any 1-layer Transformer without positional information cannot entirely solve it. We also demonstrate that position coupling can be applied to other algorithmic tasks such as Nx2 multiplication and a two-dimensional task. Our codebase is available at [github.com/HanseulJo/position-coupling](https://github.com/HanseulJo/position-coupling).

Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure
Hanseul Cho"
Jaeyoung Cha"
Pranjal Awasthi Google Research
Graduate School of AI, KAIST [jhs4015, chaj jaajjayyyonanaaaaaaa Skaist.aiata........ ac.kr
Srinadh Bhojanapalli
Anupam Gupta NYU & Google Research anupam.geny edu
Chulhee Yun
Google Research berinadh@google..oo com
Graduate School of AI, KAIST
chulhee yunkkisst. ac. ac kr
Abstract
Even for simple arithmetic tasks like integer addition, is challenging for Trans- formers to generalize longer sequences than those encountered during training. To tackle this problem, we propose position coupling, a simple yet effective method that directly embeds the structure of the tasks into the positional encoding of (decoder-only) Transformer. Taking departure from the vanilla absolute position mechanism assigning unique position IDs to each of the tokens, we assign the same position IDs to two or more "relevant" tokens; for integer addition tasks, we regard digits of the same significance in the same position. On the empirical side, we show that with the proposed position coupling, a small (1-layer) Transformer trained on 30-digit additions can generalize up to 200-digit additions (6.67 of the trained length). On the theoretical side, we prove that 1-layer Transformer with coupled positions can solve the addition task involving exponentially many digits, whereas any 1-layer Transformer without positional information cannot entirely solve it. We also demonstrate that position coupling can be applied other algorithmic tasks such as N 2 multiplication and two-dimensional task. Our codebase available at github. com
Decimal Integer Addition (Reversed Response Zero-Padding)
100
Exaut- Motch Accuracy (%6)
Train lengths (1:30)
NoPE (6-laerr B-head)
60
Random-
(G-layer head) (98.26%
Index Hinting wi NoPE (6-layer S-head) (95, 19% at 45)
Index Hinting Random-Start
(90.85%
[Ours] Position Coupling (1-layer 4-head) 65%
200)
(urrs] Position Coupling (3-ayer B-headi (98.00% 200)
10
20
30
40
50
70
80
90
100
110
120
130
140
150
160
170
180
190
200
Operand Longth
Figure 1: Methods for Length Generalization in the Integer Addition Task. We report exact- match (EM) accuracies (markers: medians over experiments; light area: 95% confidence intervals). We employ the reversed format and zero-paddings (Lee et al., 2024) into the input sequence. With our proposed position coupling, we achieve more than 95% exact-match accuracy for up to 200-digit additions with decoder-only Transformers trained on up to 30-digit additions. For index hinting (Zhou et al., 2024a), we separately test absolute positional embedding (APE) with a random starting position ID (mimicking the original implementation by Zhou et al. (2024a)) and without positional encoding (NoPE) (Kazemnejad et al., 2023) (as tested by Zhou et al. (2024b)).
Authhors contributed equally to this paper.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Introduction
Since the appearance of a sequencc----sequueee deep neural architecture called Transformer (Vaswani et al., 2017), it has brought tremendous success in various fields including natural language process (NLP) (Chowdhery et al., 2023; Gemini et al., 2023; OpenAI, 2023; Thoppilan et al., 2022) and many applications such as mathematical reasoning and theorem proving (Lewkowycz et al., 2022; Trinh et al., 2024; Wu et al., 2022). Despite its triumph, it has recently been illuminated that Transformers often lack the ability of length generalization (Anil et al., 2022; Deletang et al., 2023; Press et al., 2022: Zhang et al., 2023). It refers to a special kind of out-of-distributie generalization capability to extrapolate the model's performance to longer sequences than those encountered during training. Understanding length generalization is of great importance because a lack of it provides evidence that language models do not genuinely understand the structure of given task. Improving Transformer's length generalization has received much attention, particularly because the time/memory complexities for training Transformers grow up to quadratically in the sequence length.
Even for simple arithmetic tasks such as integer addition, length generalization is still difficult for Transformers (Kazemnejad et al., 2023; Kim et al., 2021; Lee al., 2024; Nogueira et al., 2021; Zhou et al., 2024a,b). Humans can length-generalize in integer addition because they understand the essential principle of the task. Nevertheless, it is observed that Transformers typically learn to solve addition only up to the training sequence length (Lee et al., 2024), which is different from the true arithmetic algorithm that humans "implement". This raises an important question: can we make a Transformer truly understand the structure of a task so that it can generalize to the longer sequences without training on them? In other words, can we inject the known structure of a task into a Transformer so that it can automatically length- gength-generalize?
In this paper, we propose position coupling, simple yet effective method for length generalization that directly embeds the structure of the tasks into a Transformer. In contrast to the vanilla absolute position mechanism assigning unique and consecutive position IDs to each token, we assign the same position IDs to certain input tokens that are semantically relevant. Coupling such tokens together helps the model learn solve the task regardless of the length of the given input sequence. For example, in the addition task, it is important to consider the significance of digits, so we couple the positions at the same significance (unique in each operand and the answer).
1.1 Summary of Contributions
We propose position coupling to tackle the length generalization problem of decoder-only Transformers. Our approach injects the structure of the task into the absolute position encoding by assigning the same position IDs to relevant tokens (Section 3).
With position coupling, we achieve a robust and near-perfect generalization up to 200-digit additions by training Transformers on up to 30-digit additions, which is a 6.67 extrapolation of the operand lengths (Figure 1, Section 4). It is promising since it was unclear whether the length generalization on the addition task can be solved reliably with Transformers (Zhou et al., 2024b).
We theoretically prove by concrete construction that small (1-layer, 2-head) Transformer equipped with coupled position IDs can add two decimal integers whose lengths are exponential in the embedding dimension (Theorem 5.1). Interestingly, we observe striking similarity between the attention patterns from our theoretical construction and those extracted from a Transformer trained with standard optimizer (Section 1.1. As a complementary result, we also prove that any l-layer Transformer without positional information cannot fully solve any setptuuttttooossnnstti tasks such as addition (Section 5.2).
We empirically demonstrate that position coupling can effectively address various tasks beyond addition, including multiplication between N --digit and 2-digit integers (Section 1, in which we also provide theoretical construction of 2-layer Transformer that solves this task for exponentially large N). We also verify that position coupling can aid Transformers in learning tasks with multi- mualti-dimensionll structures (Section 6.2). Moreover, we evaluate position coupling on some other tasks (addition with multiple operands, copy/reverse) in Appendix B.
2 Preliminaries
We focus on decoder-only Transformers that solve the tasks using next-token prediction (See Ap- pendix A for brief background on it). Since we study deterministic tasks with a unique answer, we consider greedy decoding throughout the paper.
2

---

2.1 Data Formats
Each task in this work is represented as a collection of sequences of the form (query)=(response) given a query, our task is to infer the response correctly. Thus, we only care about the result of the next-token prediction for the '=' token and the tokens in the response (except its last token). That is, we only compute the losses and accuracies for those output tokens.
Previous works commonly observe that data formats play an important role in solving downstream tasks with Transformers because proper data format enables the model to learn a simple function to solve a task. Here we overview some well-known methods we apply, focusing on the addition task.
Reversed Format. Lee et al. (2024) observe that reversing the response leads to improvement in both performance and sample efficiency. For example, '653 + 49 = 702' becomes '653 + 49 = 207' in reversed format. This enables decoder-only Transformer to infer the response from the least significant digit to the most significant digit, similar to how humans add two numbers.
Zero-padding. Zero-paddings ensure that the length of both operands in query is the same and the length of a response is fixed when the length of the operand is given. That is, by padding the query and the response of an M M-digit + N-digit addition with O's, the input sequence becomes max(MM, NN -diigtt addition with (max M, N + 1)-dit response. For example, '653 + 49 = 702' becomes '653 + 049 0702'.
Wrapping with BOS/EOS token(s). It is conventional in NLP to put BOS/EOS (beginning//end-of sequence) tokens at the beginning/end of the sequence. Lee et al. (2024) use the same token 'S' for BOS and EOS tokens and observe that it beneficial to wrap each sequence with the $ token when solving the addition task. We do not observe any significant difference in the performance between sequences with the same and different BOS and EOS tokens.
2.2 Positional Embeedding//Encodingg (PE)
Vaswani al. (2017) introduce the absolute positional embedding (APE) to Transformers to inject the positional information into the model. The usual APE works as follows: given an input sequence of tokens, we assign a sequence of consecutive position IDs (integers). Each position ID is mapped to a unique PE vector, and the vector is either added or concatenated to the corresponding token embedding vector. We focus on the learned APE initially proposed by Gehring et al. (2017).
Length Generalization and PE. is actively studied whether PE is a crucial factor in solving the length generalization problem of Transformers. Kazemnejad et al. (2023) argue that decoder-only Transformers with no positional encoding (NoPE) can achieve length generalization of downstream tasks since Transformer decoder implicitly capture the generalizable positional information due to its causal nature. However, there is a line of works proposing new PE methods to improve the length generalization performance of Transformers (Li et al., 2024; Ruoss al., 2023).
3 Position Coupling: A Method for Length Generalization
We propose position coupling, which assigns position IDs that directly encode the structure of given tasks. Here, we explain the general position ID assignment rule of position coupling in two steps.
First, we partition the tokens of the input sequence. The detailed principles for grouping the tokens differ by task, but the common desiderata are the following: there are two or more groups of consecutive tokens, and each token in group must have unique semantic meaning so that one-to-one correspondence between tokens in different groups can be made.
Next, for each group of tokens, we assign sequence of consecutive numbers (usually, positive integers) as position IDs, starting from random number (at training time) or a fixed predetermined number (at evaluation time). We use random position IDs at training time for inducing length general- ization by enabling all position embedding vectors to be trained, up to pre-defined hyperparameter of maximum position ID (max_pos). Very importantly, we assign the same position IDs to the tokens in all groups that are relevant to each other for solving the given task: we refer to this as "coupling the positions". Lastly, we set O as the position IDs of special tokens like BOS/EOS tokens and the PAD token (padding for minibatch training and evaluation).
'The hy hyperparameter max_poc determines the maximum testable sequence length that a Transformer can handle. Note that the idea of random assignment of position IDs is similar to randomized PE (Ruoss et al., 2023) although it is different since assigns sequence of increasing integers which are generally not consecutive.
3

---

3.1
Position Coupling for Decimal Integer Addition Task
zero-padiingg
reversed sum
Input Tokens
6
5
3
0
4
9
0
7
[t00]
(MA
Position IDs
0
7
8
6
7
8
9
7
6
5
0
0
Position Couplings
Figure Position coupling for decimal integer addition task, displaying 653 + 49 702 with appropriate imput formats. The starting position ID "6" is an arbitrarily chosen number.
We illustrate position coupling for the decimal integer addition task (or addition task for short). To study the length generalization of the addition task, we regard each digit (0-9) single token. We will use an explicit example of the addition '653 + 49- = 702' for illustration.
Before applying the position coupling, we adopt an input format similar to Lee et al. (2024) so that we reverse the response, but we use zero-padding and wrapping with BOS/EOS token 'S' at the same time. For example, '653 + 49 = 702' becomes '$653 + 049 = 20708'.
We partition the tokens in the sequence into three groups: (1) first operand & '+', (2) second operand, and (3) '=' & response (which we call 'sum'). Then each number token is "unique" in the corresponding group in terms of significance, which naturally induces a one-to-one correspondence between (most of) the tokens across different groups. We group '=' and the sum together because these tokens are where we perform next-token prediction.
Now we assign the coupled position IDs to the tokens. Most importantly, we assign the same position ID to the digits of the same significance. Let say that the random starting number is 6. In our example, we assign 6, 7, and to the tokens in the operands, and assign 5, 6, 7, and to the tokens in the sum in reversed order: see Figure 2. We remark that, first, we assign and "=' tokens because they are adjacent to the number token with position ID 8, even if there are no "significances" for those tokens. Second, we assign 5 as position ID of the most significant digit of the sum (which may be '0' due to the zero sero-padding) just because it is next to the number token with position ID 6, even though there no other corresponding tokens in the other groups (operands). We also note that the '+' token is not grouped with the second operand and is not given the ID 5; this prevent unnecessary coupling between '+' and the most significant digit of the sum.
position IDs of '+'
Remark. A concurrent work by McLeish et al. (2024) proposes an analogous approach for solving arithmetic tasks, while they employ different input format. We provide detailed comparison with our work in Appendix A.
Comparison with Index Hinting. Even though the idea of implanting the structure of a task into the positional encoding is novel, there an existing approach named index hinting (Zhou et al., 2024a) that applies similar idea but to the input sequence. Index hinting is an input augmentation technique that places position markers front of the tokens to couple the semantically relevant tokens. For example, Zhou et al. (2024a) transform '653 + 49 702' into 'a0b6c5d3 a0b0c4d9 a0b7c0d2' with some zero- zero-padding paddings, where a, b, c, and are consecutive index hints. Here, the starting hint character is randomly selected during training, similar to method of choosing the starting position ID. The reversed format and BOSJEOS tokens can be applied as well.
One way in which index hinting differs from position coupling is that it doubles the input sequence length. This is because the position information and the token information do not merge: the index hints and the normal tokens are mapped to separate token embedding vectors which are alternately placed in the input embedding matrix. As a result, a Transformer must figure out the correspondence between each adjacent pair of index hint and normal token. Moreover, the doubled input length requires up to 4 the training time and memory consumption. In contrast, position coupling explicitly combines token and position information: every token embedding and corresponding position embedding are mixed into a single vector. Hence, a Transformer can effortlessly utilize the positional structure of the task, without hurting the training time. We highlight that, as will be mentioned in Section 4. position coupling exhibits better length generalization than index hinting.
Another difference is that the index hints should be inferred by Transformers in addition to the normal tokens in the response, which might be an additional burden. Our position coupling circumvents this difficulty, eliminating the need to estimate anything other than the tokens in the original response.

---

Experiments on the Addition Task
In this section, we empirically demonstrate that position coupling allows extensive length generaliza- tion of Transformers the addition task. We delve into the impact of training length and architecture on the length generalization performance and provide comparisons with NoPE, APE with a random starting position ID (we call random-start APE), and index hinting (Zhou al., 2024a).
Data Sampling. We opt for the balanced sampling terms of the number of digits (Nogueira et al., 2021). Given the maximum number of digits Dmxx. we do balanced sampling for each operand in two steps. First, we sample the number digits D E [1, Dmax] uniformly at random. Next, we sample operand from [10P-1, 10 1] uniformly at random, except for D = 1 where sample from [0, 9]. This procedure addresses the imbalance problem in the number of digits of operands.
Model and Training. We train decoder-only Transformer models from scratch. We properly choose pas so that the maximum testable length summands is 200. We not use packing or shifting for simplicity of implementation. Since we manually put coupled position IDs with a random starting index during training, we can train all the positions without packing and shifting. We run each experiment 8 times with 2 different random seeds for data generation and 4 different random seeds for model initialization & stochastic optimization unless mentioned otherwise. We summarize all hyperparameters in Appendix C.
4.1
Results
Decimal Integer Addition (1-layer, 4-head, 202)
ExectMatthh Accuracy (N) (Median over runi)
60
Traindd on 1-10 (5 48% 70)
40
Tained on 1:20 (EM- 25.55% at
1-30 (6M-16.33330
Trained on 1-40 JEMM 96. 200)
95% line
50
60
70
80
90
100
310
120
130
140
150
180
170
180
190
200
Operand Length
Figure 3: Ablation on the trained operand lengths (I-layer 4-head models).
Longer Trained Sequences Lead to Longer Generalizable Lengths. We train 1-layer 4-head models with D Dmax E {10, 20, 30, 40} and evaluate them on up to 200-digit additions. For each run of training, we choose and evaluate the best model in terms of the validation loss for 200-digit additions. The result showcased in Figure 3. We decide that a model successfully generalizes to a certain length of operands (referred as "generalizable length") if the median EM accuracy exceeds 95%.
We observe that the generalizable length becomes longer as we train on longer training sequences. The generalizable length is 70 for the models trained on additions involving 1-10 digits, 135 for models trained on 1-20, and 200 for 1-30 and 1-40. We believe that we could achieve even longer 1-40 if we use a larger max_pos. We note that
generalizable length for the models trained could scale up the generalizable length to 500 by training with lengths 1-160: refer to Appendix B.I. Although each test sample contains the operands of the same length, we also provide an extended evaluation on test samples with operands of different lengths: see Appendix B.2.
Addition Task
Dract-Match Accuracy (S))
(Madiann overs (Uns))
Train lengths
1-layer 8-head: 95.3% 200
2-ayer 8-head 05.9%
40
J-ayee 8-head: 98 18.0% 200
4-ayar 8-head: 07. 77.6% 700
20
S-aayr 8-head: 95. 6% 200
6-aayer B-head:
0
10
20
30
40
50
60
70
80
90
1000
110
120
130
140
150
160
170
100
190
200
Query Length
Figure 4: Ablation on the number of layers (trained with position coupling).
A Ablation on Number of Layers. Since Zhou et al. (2024a) and Zhou et al. (2024b) choose 6-layer 8-head models as their base models, we also test our method to deeper models. The results evaluated with models trained on 1-30 digits are displayed in Figure 4, whose experimental details are listed in Table 2. Overall, the performances are well extrapolated to test lengths (longer than trained lengths)
5

---

and are similar for the models with 1-5 layers. For the 6-layer model, however, the performance slightly deteriorates. We hypothesize that the performance degradation is due the bad implicit bias of deeper models (learning shortcuts only to achieve in-distribution generalization) when learning simple algorithm to solve the task. Since the theoretical construction of I I-layer addition Transformer (that will appear in Section 5.1) naturally extends to larger architectures, deeper models have at least as much generalization capability as shallower models. We believe exploring a theoretical explanation for the bad implicit bias of large models on low-complexity tasks is promising research direction. We also highlight that we present median accuracies over multiple runs, while Zhou et al. (2024b) report maximum accuracies. To better the comparison, also report the maximum accuracies (for the experiments in Figures and 4) in Appendix B.3, showing that our 6-layer models can achieve near-pertee near-perfect generalization for up 200-digit additions well.
Addition Task, 1-layer 4-head model
100
Dract-Match Accuracy (S5)
(Modian Saaiinnvveerr over 8 runs)
80
Train lengths
Zero-padding
(97.7% 200)
60
Zero-padding O Reverse Query&Arewer (96, 7% et 200)
Zero-padding
(5.6% 200)
Zero- Zero-padding X. Reverse Answer (0.00 (0.0% 2001
Zero Zero-padding Reverse Query&Answer (0.0% 200)
Zero pacolppadinnngggo X. Reverse (0.0% at 200
0
10
20
30
40
50
60
70
80
90
100
110
120
130
140
150
160
170
180
190
200
Query Length
Figure 5: Ablation on the data formats (1-layer 4-head models trained with position coupling).
Addition Task, 6-layer 8-head model
100
Clast-Merch Accuracy (%g)
(Median over 8 runs)
80
Train longths
Zero-padding
Answer (68.3% 200)
Zero pard-ngdding D, Reverse Query&Answer (92.2% at 200)
Zero-padding Reverse
388.2% 200)
Zero-padding X, Reverse Answer (94.5% at 200)
pero-ad X Reverse Query&AAnswer (13.8% 200)
Zero-pedding X, Reverse X (57.6% at 200)
0
10
20
30
40
50
60
70
80
90
100
110
120
130
140
150
160
170
180
190
200
Query Length
Figure 6: Ablation on the data formats (6-layer 8-head models trained with position coupling).
Ablation on Data Formats. Our input format is primarily selected to simplify the algorithm of solving the addition task, not through extensive ablation studies. Thus, we are not arguing that our choice of input format is empirically optimal for training Transformers. However, since the data format is one of the crucial factors in general machine learning, we test various formats for 1 1-layer 4-head (Figure 5) and 6-layer 8-head models (Figure 6). The results clearly show that the performance varies with input formats, as they affect the complexity of the algorithm that the model should learn.
Small models (1-layer 4-head) achieve near-perfect generalization when the numbers are zero- pero-padded and the response or all the numbers are reversed. We believe this is because the combination of zero- padding and reversing enabled small Transformer to learn simple Kero-padding seems crucial since it aids length generalization to some extent even without reversing. Without reversing any numbers, however, even the in-distribution performance slightly decays.
algorithm.
Larger models (6-layer 8-head) perform better than small models when the numbers are no longer zero-padded reversed. We believe this is because the task-solving algorithm without reversing or zero-padding that the model should learn is more sophisticated, which larger models can learn more casily. Contrarily, we observe a slight degradation in performance when we add zero-padding and reversing in the larger model, which suggests that the model may have learned "shortcut" due to its (overly) strong expressive power relative to the problem's 's complexity.
Comparison with NoPE and APE (with random starting position ID). Our experiments demon- strate that simple PE methods like NoPE and random-start APE cannot length-generaliz well on the addition task. In particular, we implement random-start APE to mimic the training process with the usual APE combined with packing and shifting. The results showcased in Figure I imply that naively

---

training all position embeddings does not necessarily help produce a strietly better model in terms of length generalization than that does not use position embeddings at all. We also remark that even training itself is difficult for shallower models (e.g., 1-layer) with NoPE and random-start APE.
Comparison with Index Hinting. We test index hinting by running the code implemented ourselves since the original code is unavailable. From Figure 1, we observe that index hinting indeed helps the model length-generalize more than the baselines (NoPE & random-start APE). However, the generalizable lengths of the models trained with index hinting do not extend further than 50; the models completely fail starting from the length 70. We also observe that Transformers with index hinting require enough depth to achieve high enough training and in in-disribbttion validation accuracies. Particularly, the training accuracies of 1 1-layer models do not deviate from near zero Thus, we only present the results for the 6-layer 8-head model as done by Zhou et al. (2024a).
Comparison & Combination with RoPE. We also examine the possibility of combining position coupling and RoPE (Su et al., 2024): See Appendix B.4 for the experimental results and details.
5 Theoretical Analyses on 1-layer Transformers
In the previous section, we provided empirical results exhibiting the outstanding performance of position coupling. One might ask why and how position coupling works so effectively. In Section 5.1, we provide theoretical explanation by carefully constructing a I-layer Transformer model that is capable of solving the addition task involving exponentially long operands when the input is encoded with position coupling. We also present the necessity of proper positional information for 1 1-layer Transformer to solve the addition task in Section 5.2.
5.1 1-Jayer Transformer with Coupled Positions can Perform Long Additions
For the sake of simplicity of presentation, we consider Transformer without any normalization layers, as conventionally done in theoretical constructions by previous works (Awasthi and Gupta, 2023; Yun et al., 20. For the sake of completeness, readers can find mathematical formulation of the decoder-only Transformer architecture in Appendix D.
Theorem 5.1. With the input format described in Section 3.1, there exists depth-1 two-head decoder- only Transformer with coupled positions that solves the addition task with next-token prediction.
Here, the operand length is at most 2 2[(d-17)/2) - 2, where the embedding dimension is d 2 21.
We provide our proof in Appendix E. We highlight that our proof is constructive and does not rely on any universal approximation result of neural networks.
Theorem 5.1 shows that a 1-layer 2-head Transformer is sufficient for implementing addition between two exponentially long integers. We emphasize that this result can be naturally extended to larger architectures with more layers/heads, with the help of residual connections.
5.1.1 Probing the Attention Patterns in Trained Transformers with Position Coupling
We discover a striking similarity between the attention patterns in our theoretical construction (Theorem 5.1) and those extracted from a Transformer trained with position coupling and a standard optimizer. In particular, the manually constructed attention patterns described in Tables 1 and 17 in Appendix closely resemble the actual attention patterns in Figure 7.2 Drawn from this discovery, we claim that a Transformer trained with position coupling spontaneously learns two separate components of the addition task: (I) adding two numbers without carries, and (2) predicting the carries.
Let us revisit the example in Figure 2 and consider predicting '7' (position ID 6) as the next token of '0' (position ID 7). Note that the token '7' is the result of combining the digit-wise sum 6+0=6 and a propagated carry 1. To find out the sum without carry, it is enough for the model to attend to the two previous positions with ID 6: tokens '6' and '0'. On the other hand, to predict the carry, the model may attend to the three positions with ID 7: tokens '5', '4', and '0', The reason why we should care about '0' is that considering the sum 5+4 (=9) of the two digits in the operands is not sufficient to determine the existence of the carry. By looking at the token '0' in the response (with position ID 7), we can detect that the actual sum in this position is 10 (=5+4+1, where 1 is another carry propagated from the previous position) and hence we need to propagate carry 1 to the next position (with ID 6).
2 Note that they match up matrix transpose, which is due the difference in the formulations.
7

---

Attention
Causal
HEAD 1
Mask
Direction of Next taken Prediction
Diraction Moaningfull Attention
Attention
Causal
HEAD
Mask
Figure 7: Probing attention matrices of I-layer 2-head Transformer with position coupling, trained on up to 5-digit additions. (Left) There are two heatmaps (clipped to zero below 0.01) corresponding to the (transposed) attention matrices observed from the attention heads. Averaged over 10K se- quences of 6-digit additions. (Right) We magnify parts of the attention matrices that are involved in inferring the response (sum). The arrows explain the process of inferring the next token '0' from '3'.
Now we inspect the aforementioned claim by examining the attention matrices of an actual trained Transformer. In the model, discover two different pattems of attention matrices," 3 playing distinct roles. The first attention pattern (top of the figure) seems to correspond to the addition without carries: each token in the response (including '=') attends to two positions needed find out the sum without carry. Conversely, the second attention pattern (bottom of the figure) seems to correspond to the carry prediction: again, each token in the response attends to three positions required to find out the carry.
Remark. Similarly to our analysis, Quirke and Barez (2024) study the attention patterns of a 1-layer 3-head decoder-only Transformer model trained solely 5-digit addition. They also observe that each head handles different subtasks addition, such as digit- wise summation and carry detection.
5.2 1-layer Transformers Require Positional Information
In Section 4.1, observed that I-layer Transformers fail to perform the addition task without position coupling. Here, we provide a partial result that theoretically explains why this happens inevitably, particularly in the case of NoPE. We start with general proposition: I-layer Transformer without positional encoding cannot distinguish queries that are identical up to permutation when inferring the first token of the response using greedy next-token prediction.
Proposition 5.2. Consider any depth-I finite-head decoder-only Transformer model T without positional encoding (NoPE), Given an input sequence I and its arbitrary permulation I', if the last tokens of and I' are identical, then the next tokens predicted by T will also identical for both sequences when applying greedy decoding scheme.
The proof is deferred to Appendix F. According to the proposition above, the 1-layer Transformer without positional encoding will always output the same values starting from the "=" token, provided that the combination of query tokens identical, even their order varies. However, the addition task is permutatioon-sensiiiy meaning that the permuted queries may result in different responses. Therefore, the I-layer Transformer cannot completely solve the task without positional encoding. is important to note that this result remains unchanged regardless of the input format: neither reversed format nor index hinting provides any benefit. We also highlight that this impossibility result can be extended to any other permutation sensitive tasks, such arithmetic tasks and copy/reverse tasks.
Based on this, we write code to directly calculate the maximum EM accuracy on the m-digit addition task that a 1-layer decoder-only Transformer can achieve (see Appendix F for the code). The accuracies rapidly decrease to zero: 6.2% for 3-digit addition, 1% for 4-digit integers, and 0. 0.13% for 5-digit integers. We leave it for future work to investigate the necessary conditions of the architecture for implementing addition when other positional encoding schemes are employed.
'The attention matrices depicted in Figure? 7 are square, lower-triangular (due to causal attention pattern), and tow-stochastic (all entries nonnegative and the sum of each row equals 1).
8

---

6 Applying Position Coupling Beyond Addition Task
To demonstrate the versatility of position coupling, we consider two other tasks in this section: N x multiplication and two-dimensional (2D) task. Other example tasks (e.g., addition with multiple summands, copy/reverse allowing duplicates) can be found in Appendix B.
6.1 Position Coupling for N x Multiplication Tasks
Here, we study length generalization on the N-digit x 2-digit multiplication task in terms of the length N of the first operand, while fixing the length of the second operand by 2. Similar tasks have been studied before (Duan and Shi, 2023; Jelassi et al., 2023); we discuss further in Appendix A.
We reverse and zero-pad the response, setting the length of it N 2. We couple the position starting from the least significant digits of both operands and response, decrementing the ID as we move to their most significant digits: showcased in Figure verify the efficacy of position coupling compared to NoPE and random-start APE. We observe that a l-layer model fails even with position coupling. even for training. However, as the depth increases to 2 or more, it immediately becomes capable of length generalization.
Figure 17 in Appendix B.5. The experimental results
(N-digit)x (2-digit) Multiplication
Match Accuracy (Median)
100
60
trained (1
Position Coupling, 1-aaer 8-head (EM-0.01%
Position Coupling. 2-layer 8-head (FM-05 54% 70)
Position Coupling. 3-ayer 8-head (EM-96,
80)
Position Coupling 4-layer 8ihead (EM=96. 81% 90)
NoPE, 3 3-ayer Bihead (EM 99 41.41% at 42)
Random-Stare
8-head (EM-08.04%
line
20
40
60
80
100
120
140
160
180
200
Operand Length
Figure 8: N x multiplication task, trained on sequences of length 1-40.
Unlike addition, position coupling for N x 2 multiplication is less intuitive, predicting the token in the middle of the response requires multiple digits from both operands while each token in the response is linked with at most 2 tokens in the query. Perhaps surprisingly, can still construct Transformer that provably solves this task for exponentially long sequences.
Theorem 6.1. Given an appropriate format of the input sequence, there exists a depth-2 decoder-only Transformer model with coupled positions that can perform the N x 2 multiplication task with next-token prediction. Here, the number of the total heads is 10 and the length of the first operand is at most 2 ( (d334)/6] 3, where we denote the token embedding dimension by d 46.
We defer the proof Appendix G. This result suggests that the proposed position coupling scheme for the N x 2 multiplication task sufficiently captures the inherent structure of the task, and thus provides the potential for the trained model to generalize across unseen lengths, Also, we believe that Theorem 6.1 is optimal in terms of the number of attention layers, as the depth- model exhibits total failure even for in-distribution samples our experiment.
6.2 Two-dimensional Position Coupling for Minesweeper Generator Task
Now, we investigate the extension position coupling for handling a 2D task, where the query and the response originally 2D objects. In particular, we define and investigate a task we call minesweeper generator. Given a rectangular board where each cell is filled with either 'M' (mine) or (an empty cell). the task is generate new board of the same size, having each cell filled with:
'M', if the corresponding cell in the original board contains 'M';
The count of mines adjacent cells, if the corresponding cell in the original board contains '*'
Data Format & Position Coupling. We introduce two position coupling modules: one for the row direction and another for the column direction. Following this, we flatten the board to feed it into Transformer: see Figure 9. Within the model, an embedding vector for each token (cell) is generated by adding the token embedding vector and corresponding two PE vectors.

---

Ireut Tokens
M
Roos Couplng
3
3 4
Flaroon
Powtten Ds
1
M
2
Oolumn CamnnDousrry
Figure 9: Position coupling for the two-dimensional 'minesweeper generator' task. (Left) The idea of assigning coupled position IDs. (Right) The model receives a flattened sequence input tokens and two-dimensional position IDs.
Experiments. To assess the efficacy of position coupling, we contrast its performance with NoPE. The training samples are designed with the width and height of the board between 5 and 9 inclusively. We allow the width and height be different for training samples. We evaluate the test performance on square board with a width between 5 and 14 inclusively. We also employ 4-layer 8-head model for position coupling and a 6-layer 8-head model for NoPE. In particular, for position coupling, we use the same embedding layer for both position coupling modules, as this approach empirically performs better than using distinet embedding layers for each module (see Appendix B.8)
The experimental results are described in Figure 10. Position coupling maintains over 98% accuracy until a width of 12 and near 90% accuracy even a width of 14. In contrast, NoPE fails even for in-distribution samples. One might be concerned that the generalizable length of 12 seems only slightly higher than the trained length of 9. However, we stress that our query 2D board, therefore the actual length generalization is from 81 to 144.
Dasc Daac: March Acurscy (14) (Mtainn))
Minesweeper Generator
2.0
Tran ar widheehiggt (2-8)
Noft (5 syer 3- b-head) (D4-99.50% at
Penliio Coup ing (4 layer S-haad) 12W-28..0% a2))2)
20
Widheeggtt
Figure 10: Minesweeper generator task, trained on sequences of length (5-9)x(5-9),
Conclusion
Achieving length generalization of Transformers even in the simple case of the addition task has been a challenge that received lot of attention. We propose position coupling, variant of learned APE, which enables capturing task structure to improve the length generalization performance of Transformers for addition. We show that Transformer trained on 1-30 digit addition generalize up to 200-digit addition. We also provide the construction of a 1-layer Transformer model capable of adding two exponentially long integers when position coupling is applied. Furthermore, we verify the efficacy of position coupling for length generalization in other arithmetic and algorithmic tasks.
I Limitations & Future Directions. We intentionally limited ourselves to the tasks with explicit structure between the tokens in each sequence. This is because we are proposing a method to instill the known structure of the task into Transformer by training coupling of positions for tasks whose structure is implicit or black-box (c.g., for general NLP tasks) remains fascinating next step: we leave the methodology for uncovering hidden structures and autonomously creating appropriate couplings (without manually designing them) for future work.
short sequences. Designing the
We also leave two challenging arithmetic tasks to addition with a varying number of summands, i.e., determining if the model can generalize to summing multiple integers when trained on samples with fewer summands. The second task is multiplication, where the lengths of both operands can vary. Note that our method is further extended to solve these two challenging length generalization problems in recent work (Cho et al., 2024).
for future work. One is the
10

---

Acknowledgments and Disclosure of Funding
This work was partly supported by a National Research Foundation of Korea (NRF) grant (No. RS-2024-00421203) funded by the Korean government (MSIT), and an Institute for Information & communications Technology Planning & Evaluation (ITTP) grant N0o.RS-2019-I119000 Intelligence Graduate School Program (KAIST)) funded by the Korean government (MSIT). HC, JC, and CY acknowledge support from a Google Gift on the research related Long Context Transformers. The experiments contained in this work were supported in part through a Google Cloud Platform Credit Award.
Artificial
References
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization the unseen, logic reasoning and degree curriculum. In International Conference on Machine Learning, pages 31-60.
PMLR, 2023. A.2
Kartik Ahuja and Amin Mansouri. On provable length and compositional generalization. arXiv preprint arXiv:2402.04877 2024. A.2
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. Advances in Neural Information Processing Systems, 01388566-35550 2022. 1
Pranjal Awasthi and Anupam Gupta. Improving length- augth-gennrallzzatio in transformers via task hinting. arXiv preprint arXiv:2310.00720 2023. 5.1
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1071v:1607.064400, 2016.
Hanseul Cho, Jaeyoung Cha, Srinadh Bhojanapalli, and Chulhee Yun. Arithmetic transformers can length-generalize in both operand length and count. arXiv preprint vrXiv:2410.15787 2024. 7
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference machine learning, pages 933-941 PMLR, 2017. D
Gregoire Deletang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, and Pedro A Ortega. Neural networks and the chomsky hierarchy. In The Eleventh International Conference on Learning Representations, 2023. URL /tttpp:://ppeeeee
Shaoxiong Duan and Yining Shi. From interpolation extrapolation: Complete length generalization for arithmetic transformers. arXiv preprint arXiv:2310.1 1984, 2023. 6.1, A.2
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. Advances in Neural Information Processing Systems, 36, 2023. 10
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann Dauphin. Convolutional sequence to sequence learning. In International conference on machine learning, pages 1243-1252. PMLR, 2017. 2.2
Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023. 1
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 1wX1v.1606.00415, 2016. D

---

Kevin Jarrett, Koray Kavukcuoglu, Marc' Aurelio Ranzato, and Yann LeCun. What is the best multi-stage architecture for object recognition? In 2009 IEEE 12th international conference on computer vision, pages 2146-2153. IEEE, 2009. D
Samy Jelassi,
Stéphane d' Ascoli, Carles Domingo-Enrinn Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers, arXiv preprint arXiv:230.115000 2023. 6.1,A.2
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy, The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2023. 1, 1, 2.2, C
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 Proceedings of the 2021 Conference Empirical Methods in Natural Language Processing, pages 7031-7037, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. 1
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024. URL /sttps://openrevee www.:e/oooromdddddeeoaaaaaao A.2
1, 1, 2.1,
Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843-3857, 2022.
Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang, Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. In The Twelfth International Conference on Learning Representations, 2024. URL
2.2
David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Mikulik. Tracr: Compiled transformers laboratory for interpretability. Advances in Neural Information Processing Systems, 36, 2023. 10
Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, and Tom Goldstein. Transformers can arithmetic with the right embeddings. arXiv preprint arXiv:2405. 17399, 2024. 3.1, A.2
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference on machine learning (ICML-10), pages 807-814, 2010. D
Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with
simple arithmetic tasks. arXiv preprint arXiv:2102. 3019, 2021. 1,4 4
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. C
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. URL https: /ttp:://///preeeee.eeoo wttp:e/oooomm.ii-d-doPPPPoo.... 1
12

---

Philip Quirke and Fazl Barez. Understanding addition in transformers. In The Twelfth International Conference on Learning Representations, 2024. URL httpp:::/ooparereeeeeeeeeeee/e/s wwwp:rrwwww....eeoooniniiiiia rIx1YXWWZb. 5.1.1
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. Journal of machine learning research, 21(140):1-67, 2020. C
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. arXiv preprini 1.Xiv:2305.168433 2023. 2.2,
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 2020. 1, D,D
Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Positional description matters for transformers arithmetic. arXiv preprint arXiv:2311.14737 2023. A.2
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiigg 568:127063 2024. 4.1, B.4
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXivv2201.0022
2022. 1
Trieu H. Trinh, Yuhuai Wu, Quoc V. Le. He He, and Thang Luong. Solving olympiad geometry with- out human demonstrations. Nature, 620 E8, 2024. doi: 10.10366411566-024007113. URL https: //t.i//doiiooa org/10.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 1, 2.2, C, D
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 1080011090, PMLR, 2021. 10
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et Huggingface's transformers: State-of-the-ar natural language processing. arXiv preprint 190.v:1110 2019. C
Yuhuai Wu, Albert Qiaochu Jiang, Wenda Li, Markus Rabe, Charles Staats, Mateja Jamnik, and
Christian Szegedy. Autoformalization with large language models. Advances in Neural Information Processing Systems, 35:32353-32366 2022. 1
Changnan Xiao and Bing Liu. Conditions for length generalization in learning reasoning skills. arXiv preprint urXiv:2311.1611 16173, 2023. A.2
Changnan Xiao and Bing Liu. A theory for length generalization in learning to reason. arXiv preprint arXiv:2404.00566 2024. A.2
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020. C
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeenc functions? In International Con- ference on Learning Representations, 2020a. URL https: httep:://openreviee net/forum?id- ByxRMONtvr. 5.1, D.D
Chulhee Yun, Yin- Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O(n) connections are expressive enough: Universal approximability of sparse transformers. Advances in Neural Information Processing Systems, 33:13783- 13794, 2020b. 5.1
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019. C, F,
13

---

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115. 2021. G.8, G.8
Yi Zhang, Arturs Backurs, Sebastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wag- ner. Unveiling transformers with LEGO: A synthetic reasoning task, 2023. URL https: /ttpen/eviww.nee//// wwt./neeetddddieaaeeeo (atprrummP--jjD--RRREE- RfQfrb. 1
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M. Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. In The Twelfth Iniernational Conference on Learning Representations, 2024a. URL htttps://openrevier net http://ormmimiimemee IuHnmHX. 1, 3.1, 4.1, 4.1, A.2, B.7, 10
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou, Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024b. 1, 1, 1.1, 4.1, A.2, B.3, C
14

---

Contents
Introduction
1.1
Summary of Contributions
2 Preliminaries
2.1 Data Formats
2.2
Positional Embeddings/Encoding (PE)
3
Position Coupling: A Method for Length Generalization
3.1
Position Coupling for Decimal Integer Addition Task
4
Experiments on the Addition Task
4.1
Results
5
Theoretical Analyses 1-layer Transformers
5.1
1-layer Transformer with Coupled Positions can Perform Long Additions
5.1.1 Probing the Attention Patterns in Trained Transformers with Position Coupling
5.2
1-layer Transformers Require Positional Information
Applying Position Coupling Beyond Addition Task
6.1
Position Coupling for N 2 Multiplication Tasks
6.2
Two-dimensional Position Coupling for Minesweeper Generator Task
7
Conclusion
10
Omitted Backgrounds
17
A.1
Next-token Prediction with Decoder-only Transformers
17
A.2
Related Works
17
More Applications & Experiments of Position Couping
19
B.1
Decimal Integer Addition Task: Scale-up Length of 500
19
B.2 Decimal Integer Addition Task: Operands of Different Lengths
19
B.3 Decimal Integer Addition Task: Maximum Exact-Match Accuracies
20
B.4
Decimal Integer Addition Task: Comparison & Combination with RoPE
20
B.5 Position Coupling for N 2 Multiplication Tasks
B.6 Addition Task with Multiple Summands
B.7 Position Coupling for Copy/Reverse Tasks
22
B.8 Position Coupling for Minesweeper Generator Tasks
22
C Experiment Details and Hyperparameters
24
Decoder-only Transformer Architecture
29
Formal Construction Addition Transformer with Position Coupling
30
E.1 Notation
30
E.2 Input Sequence
30
15

---

E.3
Encoding Function
32
E.3.1
Token Embedding
32
E.3.2
Coupled Position IDs and Position Embedding
32
E,4
Transformer Block - Causal Attention Layer
33
E.4.1
Attention Head I: Digit-wise Addition without Carries
33
E.4.2
Attention Head 2: Carry & EOS Detection
36
E.4.3
Residual Connection
40
E.5
Transformer Block Token-wise Feed-forwand Layer
41
E.5.1
Subnetwork 1: Construction for SUM (dimension 7-16).
42
E.5.2
Subnetwork 2: Construction for IS EOS (dimension 17).
43
E.5.3
Residual Connection
44
E.6
Decoding Function
44
Impossibility Addition with No Positional Encoding
46
(Formal) Construction of N x Multiplication Transformer with Position Coupling
48
G.I Notation
48
G.2 Input Sequence
48
G.3 Encoding Function
49
G.3.1 Token Embedding
49
G.3.2
Coupled Position IDs and Position Embedding
50
G.4
Construction Idea
51
G.5
Transformer Block Causal Attention Layer
52
G.5.1
Attention Head 1:
Detecting the Ones Digit of the Second Operand
53
G.5.2
Attention Head 2: Detecting the Tens Digit of the Second Operand
54
G.5.3
Attention Head 3: Position Masking
56
G.5.4
Residual Connection
58
G.6
Transformer Block -
Baken-wisse Feed forward Layer
58
G.6.1
Residual Connection
59
G.7
Transformer Block 2 - Causal Attention Layer
59
G.7.1
Attention Head 1: Copying the Ones Digit of the Second Operand
60
G.7.2
Attention Head 2: Copying the Tens Digit of the Second Operand
61
G.7.3
Attention Head 3: Copying the Appropriate Digit from the First Operand
62
G.7.4
Attention Head 4: Copying the Appropriate Digit from the First Operand II
64
G.7.5
Attention Head 5: Copying the Appropriate Digit from the First Operand III
66
G.7.6
Attention Head 6: Copying the Appropriate Digit from the First Operand IV
68
G.7.7
Attention Head 7: Copying the Appropriate Digit from the First Operand V
69
G.7.8 Residual Connection
71
G.8
Transformer Block 2 - Token-wise Feed- forward Layer
72
G.8.1 Residual Connection
74
G.9
Decoding Function
74
16

---

A Omitted Backgrounds
A.1 Next-token Prediction with Decoder-only Transformers
Calculate Loss
7
2
Transformer Decoder
6
5
3 query
4
9
7
2
[R04] (AAX
response
Figure 11: Schematic of solving an integer addition task instance using next-token prediction with a decoder-only Transformers. BOS/EOS mean siwtuiiigg-eenoo---uuuunnee tokens, respectively, PAD means a padding token, used for matching the sequence lengths in single minibatch of sequences. Here we assume a basic input format (plain, no pero-padding), which different from that we used in our experiment.
decoder-only Transformer returns an output sequence of the same length as the input sequence. One difference from Transformer encoder that the attention mechanism in Transformer decoder occurs only in a single forward direction due to the causal attention mask. Due to this causal nature, the Transformer decoder is mostly used for inferring the next token of each token, just based on the information of the current and the previous tokens.
A.2 Related Works
Length Generalization in the Addition Tasks. Lee al. (2024) observe that reversing the output in the addition task enables the model to learn simple function. Shen et al. (2023) propose "Random Spacing" and "Recursive Scratchpad", achieving near-perfect generalization from 10-digits 12- digits addition. Zhou et al. (2024a) introduce "index hints", position markers placed in front of each token, both the input and output of addition tasks. Most recently, Zhou et al. (2024b) demonstrate possibility of extrapolation to the length 100 with training length 1-40 in the addition task by combining appropriate input format and advanced PE, yet they also observe that the performances are not robust and highly depend on the random seeds.
Length Generalization the N M Multiplication Task (M is fixed). Jelassi al. (2023) investigate N x 3 using encoder-only model and Duan and Shi (2023) study N with an encoder-decoder Transformer architecture, Besides the architectural difference, Jelassi et al. (2023) fail to observe length generalization with RPE and only achieve it by supplementing small number of long samples to the training set. Furthermore, although Duan and Shi (2023) provide perfect length generalization results even for test samples 10x longer than those observed during training, their approach requires a retraining step with
bias correction on attention score matrices.
Analyzing Length Generalization in Theoretical Perspectives. An emerging line of research seeks to theoretically address why length generalization is difficult and under what conditions it can be achieved. Abbe al. (2023), the authors demonstrate that various neural network models have an implicit bias towards min-degree interpolators, which may not be ideal for various reasoning tasks. Xiao and Liu (2023, 2024) investigate problems whose reasoning processes can be formulated as directed acyclic graph (DAG) structures, introducing the concept of maximal input element distance to identify a sufficient condition for length generalization. Recently, Ahuja and Mansouri (2024) formulate the conditions of function classes required to guarantee the length generalization of the empirical risk minimizer function.
Comparison with McLeish et al. (2024) A very recent concurrent work by McLeish et al. (2024) proposes a new position embedding method called "Abacus". From methodological perspective, Abacus is almost identical to our position coupling except for two main differences: Abacus reverses both the query and the response and does not use padding. From now on, we outline the differences between their work and ours beyond the methodology.
17

---

In terms of the model architecture, they use a depth- 16 decoder-only Transformer model. They combine their method with looped Transformers and input injection and report an improved perfor- mance. In contrast, our main results are obtained with shallower models (up to 6 layers) with standard Transformer architecture of stacked decoder layers.
Besides the addition task, they study multiplication, sorting, and Bitwise OR. On the other hand, we study multiplication, triple addition, copy/reverse, and 2D task. Specifically, for the multiplication task, their study mainly considers the case where the length of both operands could vary up to 15. In contrast, we focus solely on the N task, fixing the length of the second operand by 2. While we achieve length generalization up to 90 90-digit multiplication by training the model on up to 40-digit multiplication, they report near-pen serr-perfect n-diistribution performance but poor length generalization.
Finally and notably, we provide novel theoretical analyses, including (1) the constructive proof that depth-1 Transformer equipped with position coupling can completely solve the addition task for exponentially long digits and (2) the impossibility the same model being capable of the addition task. We also present theoretical results for the N 2 multiplication task.
18

---

B More Applications & Experiments of Position Couping
B.1 Decimal Integer Addition Task: Scale-up to Length of 500
Here, demonstrate the scalability of our proposed position coupling approach for large lengths of up to 500. Specifically, we again train a depth-1 decoder-only model for the addition task and evaluate the performance for instances with up to 500 digits. The results are shown in Figure 12. We notice that at train length of 160, we achieve excellent length generalization for 500-digit addition. On the other hand, training on sequences of length up 40 or 80 is insufficient for extreme length generalization. The results demonstrate that position coupling, as an approach, is highly scalable.
Decimal Integer Addition
100
Exicct Motch Accurscy (N) (Medlan over nuna)
60
[EM-96.30K 200)
Train on 1-160 (EMw05, 15% 500)
40
80
100
160
200
300
400
500
Operand Length
Figure The exact-match accuracies obtained by training depth- transformer on the addition task. We see that while training with sequences of length up to 40 and 80 is insufficient for generalization to large lengths, training length 160 we achieve strong performance for lengths up to 500. The experimental details can be found in Table 3.
B.2 Decimal Integer Addition Task: Operands of Different Lengths
In the main text, mainly focus the test examples of additions where the lengths of both operands are the same. For the sake of completeness, we also report the evaluation results for the cases where the operand lengths be different (although the zero- paro-pgdding is applied ensure the consistency of the data format). See Figure 13.
1-tayer 4-haad,
runs)
1-layer 4-haad,
runs)
100
100
60
60
19
1-layer 4-head, trained on 1-30 (8 runs)
1-layer 4-head, trained on 1-40 (8 runs)
100
100
40
40
Figure 13: Exact-matee accuracies (%) additions with operands of different lengths. Different heatmap corresponds to different trained length of operands (1-10, 1-20, 1-30, and 1-40, expressed with red box for each). For each heatmap, the z-axii and the y-axis are for the length of the first and the second operand, respectively.
19

---

B.3 Decimal Integer Addition Task: Maximum Exact-Match Accuracies
prior work by Zhou et al. (2024b) provides a similar analysis on the addition tasks as ours. Combining appropriate input format and advanced PE, they achieve 298% EM accuracy for 100-digit additions with a 6-layer 8-head model trained on 1-40. Moreover, they achieve generalizable length of 45 for model trained 1-30, 25 for 1-20, and 10 for 1-10 (no length generalization). One big difference between their analysis and ours they report the maximum accuracy for each testing length trials, while report the medians. Thus, choose a bit lower threshold (95%) for generalizability than theirs. For better comparison with Zhou et al. (2024b), we report the maximum exact-match (EM) accuracies. See Figures 14 and 15.
Decimal Integer Addition (max p05-202)
100
Exect Exact Match Accurscy 1%)
(Max svar runss
Triind on 1-10 [EM-96 14% 65
Triined on 1-20 (EM-96. 62% 1400
JEM-
Taained on 1-40 (9M-99
39% 200)
98% line
60
70
80
90
100
330
320
130
140
150
160
170
180
190
200
Length
Figure 14: Ablation on the trained lengths (1-layer 4-head model trained with position coupling). Here, report maximum EM accuracies over runs for each tested length.
Addition Task
100
Cxact Mertc--tcchhhcc Accuracy (%)
(Mox over runs)
Train lengths
1-layar 8-head
2-layer B-head: 98.8% 200
3-aeer 8-head:
4-byyr B-head. 98.8% 200
5-layer
6-layer 8-head: 96 2% at 200
10
20
30
50
150
160
170
180
190
200
Query Length
Figure 15: Ablation on the number of layers (trained with position coupling). Here, we report maximum EM accuracies over runs for each tested length.
B.4 Decimal Integer Addition Task: Comparison & Combination with RoPE
We examine further the possibility of combining our position coupling and RoPE (Su et al., 2024).
RoPE incorporates the positional information into the key and the query vectors by rotating them. Suppose position ID is assigned to a key vector k. for every two consecutive entries of k (i.c., (k2j-1, k2))), we rotate by predefined angle 0; multiplied by m:
A2i-1 Kzi
12 cos m0; k2 sin m0, 12y-1 sin md; + lzz cos m6,
We also apply a similar rotation to the query vectors (say, q with position ID n).. As result, the attention score for this key-query pair becomes function of k. q. and the relative distance
Unlike the original implementation of RoPE, we apply rotations based on the assigned position IDs according to our position coupling method. We incorporate this RoPE variant into every attention head of every layer. One more difference implementation that, during training, we randomly sampled an integer scaler & E 2,, 4, 5, and multiplied it by the rotation angle. By such random re- scalilin of rotation angles, we expect the model could handle unseen large rotation angles at test time.
20

---

The result on 12-layer models is showcased in Figure 16 (the orange line). Unlike vanilla RoPE (the blue line), which fails immediately outside the trained lengths (1-40), our combination RoPE and position coupling achieves a much better generalization up to operand lengths 100.
Extct-Match Accuracy (%)
Addition Task, 12-layer model
(Median over 3 runs
100
75
50
Train lengths
25
RoPE
RoPE Position Coupling
10
20
30
40
50
60
70
90
100
Query Length
Figure 16: RoPE-based position coupling, 12-layer model.
B.5 Position Coupling for N Multiplication Tasks
We present example of the position coupling method for N 2 (N- (igit by 2-digit) multiplication, which is omitted from the main text. See Section 6.1 for the experimental results.
product
[B08]
7
X
=
5
0
6
[508]
0
4
5
8
6
7
6
5
2
0
Figure 17: Illustration of position coupling for N 2 multiplication task.
B.6 Addition Task with Multiple Summands
The position coupling scheme for the vanilla addition task (with two operands) can naturally extend to the addition task with multiple summands: assign position IDs in ascending order from most significant digits to least significant digits for every operand and the response. Here, we focus on the addition of three summands.
Exacbbraton Accurncy 154 (Mediin orer (art)
Addition
lingtha
1-byer hhewd (DM
Zlayer f-hred (DM 175 w:)
3tayer 6-heed (DM-26 9:900
wayer B-nead (EV--3..22 at
40
Random-Sttrf APE, byyer Bhheed (EMio 500
20-
Figure 18: Exact-match accuracy (median over 4 runs) for triple addition task, trained on sequences of length 1-40 with position coupling, NoPE, and random-see start APE. For further experiment details, refer to Table 6.
Experiments. We train sequences with operands of 1-40 digits. Our choice of max_poo is 102, so we test the operands of up to 100 digits. We investigate the performance of 3 different architectures, each with different depth. The experimental results described in Figure 18. 1-layer models keep their generalization capability until 100 digits, whereas the 3-layer models exhibit great stability across random seeds and achieve the highest generalizable length of 90.
Lastly, we note that the result Theorem 5.1 can extended to addition tasks with multiple summands with slight adjustments to the feed-forward layer in the construction.
21

---

B.7 Position Coupling for Copy/Reverse Tasks
Data Format Position Coupling. Each token of the query sequence is digit (10 distinet characters). We couple the positions in the query and the response by their correspondence. Note that the position ID assigned the equal token different for the two tasks because as per our design principle (Sec 3.1), the equal token is grouped to the response tokens and position IDs have to be consecutive numbers within each group.
[806)
b
a
C
=
a
b
a
(809)
(R06)
b
a
=
a
b
(000)
0
4
6
2
4
6
0
5
7
5
4
0
Figure 19: Illustration of position coupling for copy/reverse tasks.
Experiments. We compare position coupling with NoPE and random-start APE. We train a model on lengths 1-40 and evaluate its performance on lengths from 300, intervals of While 1-layer 4-head model is used for the position coupling, we observe that the same architecture fails to memorize training samples for both NoPE and random-start APE. Therefore, we use 6-layer 8-head model for the latter cases as it is commonly used in the literature (Zhou et al., 2024a).
March Accuracy (Medlan)
Copy
1000
Train lengt (1-40)
NoPE (6-layer B-head] (0M-96228%- 75)
Random Start 8-head) (EM-95.50%
40
Position Coupling (1-layer 4-hwad] EEM-100.00% 300)
10
20
30
40
50
60
80
90
100
1001100110
120
130
140
150 160 170 180 190 200 210 220 230 250 270 300 Query Length
Reverse
100
80
Train lengths (1 (1-40)
Accurscy
NoPE (6-layer B-head] (EM-36.86%) at
Random-Start (6-layer (EM-98
40
Position Coupling (1-layer 4-heacl (EM-100.000 at 300)
10
20
30
40
50
60
70
80
90
100
100110220
130 140 150 160
180 190 200 210 230 250
270 280 280 300
Query
Figure 20: Exact-match accuracy (median over 4 runs) for (a) copying task and (b) reversing task, trained sequences of length 1-40 with position coupling, NoPE, and random-start APE, For further experiment details, refer to the Table 7.
The experimental results are described Figure 20. For both copy and reverse tasks, position coupling exhibits near-perfect accuracy across the entire test length (7.5x for the trained length). In contrast, NoPE and random-start APE immediately fail to length-gente genereeeeerzeee
B.8 Position Coupling for Minesweeper Generator Tasks
Here, we present the extra experimental results for training the minesweeper generator task with position coupling. Specifically, compare the performance of two configurations: where the model shares the same positional embedding layer for both position coupling modules, and another where the model uses separate positional embedding layers for each position coupling module.
The results are described Figure 21. When sharing the same positional embedding layer, position coupling achieves over 98% accuracy 12x 12 board, and maintains near 90% accuracy
22

---

14x 14 board. However, with distinct positional embedding layers, position coupling only successfully generalizes 10x 10 board. We currently do not have clear explanation for why the former method exhibits significantly better performance than the latter one. We leave the investigation and explanation of this phenomenon for future work. Minesweeper Generator
DocchMatch Accuracy (10) (Median)
80
60
widthssheighhs (5-9)
20
Position Coupling with Shared Embedding (4 layer 8 hrad) (EM-98.07% at 12)
Position Coupling Distinct Embeddlig (4-layer 8-head) (EM-39.
WidttHHelght
Figure 21: Exact-match accuracy (median over runs) for minesweeper generator task, trained on sequences of length (5-9)x(5-9)) 9) with position coupling. For further experiment details, see Table 8.
23

---

C Experiment Details and Hyperparameters
Position coupling can be easily implemented on top usual librarics of training transformer models like HuggingFace (Wolf et al., 2019) and Flaxformer' since these libraries support arbitrary array of position IDs (in the case of using APE). All we need is to build up short routine implementing the assigning rule of position IDs when establishing the dataset and data loaders. To compare with NoPE, we use the code base provided by Kazemnejad et al. (2023) for most of the experiments. It contains custom implementation of decoder-only T5 (Raffel et al., 2020) established top of PyTorch (Paszke et al., 2019) and Huggingface, including several PE methods. We additionally implement custom RMSNorm module (Zhang and Sennrich, 2019) and various positioning schemes of normalization layers (e.g., PreNorm (Xiong et al., 2020), PostNorm (Vaswani et al., 2017), and their combination), follow the implementation details of Zhou al. (2024b).
Table I: Hyperparameter summary for decimal integer addition task: comparison between trained lengths (Figures 3 and 14).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
Number of Attention Heads
4
Embedding Dimension
512
Dimension per Head
128
Hidden Width of Fred-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU (Shazeer, 2020)
Normalization Layer
RMSNorm (Zhang and Sennrich, 2019)
Normalization Layer Position
PreNorm and PostNorm
Training Steps
50,000
Batch Size
1,000
Optimizer
Adam (Kingma and Ba, 2015)
Learning Rate (LR)
0.0001
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0.1LR)
LR Cool-down
Maximum Position ID (max_pos)
202
Training Dataset Size
1,000,000
Evaluation Dataset Size
100,000
Device
NVIDIA RTX A6000 48GB
Training Time
10 hours
"giihhbb.ccm/poogg/////afoeene

---

Table 2: Hyperparameter summary for decimal integer addition task: comparison between the number of layers (Figures 4 and 15).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
1-6
Number of Attention Heads
8
Embedding Dimension
1024
Dimension Head
128
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU
Normalization Layer
RMSNorm
Normalization Layer Position
PreNorm and PostNorm
Trained Lengths of Operands
1-30
Training Steps
50,000
Batch Size
1000
Optimizer
Adam
Learning Rate (LR)
0.00003
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0.ILR)
LR Cool-down
Maximum Position ID (max_pos)
202
Training Dataset Size
1,000,000
Evaluation Dataset Size
100,000
Device
NVIDIA RTX A6000 48GB 10 hours
Training Time
Table 3: Hyperparameter summary for decimal integer addition task: generalization up to length 500
(Figure 12).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
Number of Attention Heads
2
Embedding Dimension
512
Dimension per Head
256
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU
Normalization Layer
LayerNorm (Ba et al., 2016)
Normalization Layer Position
PostNorm
Training Steps
1,000,000
Batch Size
128
Optimizer
Adam
Learning Rate (LR)
0.0001
LR Warm-up
Linear (From 0 to LR), 500 steps Cosine Decay (From LR to 0.0)
LR Cool-down
Maximum Position ID (max_pos)
1003
Training Dataset Size
1,000,000
Evaluation Dataset Size
100,000
Device
TPU V4 Chips
Training Time
4 hours
25

---

Table Hyperparameter summary for decimal integer addition task: extracting attention patterns
(Figure 7).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
1
Number of Attention Heads
2
Embedding Dimension
512
Dimension Head
256
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU
Normalization Layer
RMSNorm
Normalization Layer Position
PreNorm and PostNorm
Trained Lengths of Operands
1-5
Training Steps
50,000
Batch Size
100
Optimizer
Adam
Learning Rate (LR)
0.00005
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0.ILR)
LR Cool-down
Maximum Position ID (max_pos)
17
Training Dataset Size
100,000
Device
NVIDIA RTX A6000 48GB
Training Time
S 6 hours
Table 5: Hyperparameter summary for N multiplication task (Figure 8).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
1-4 (Ours), 3 (NoPE & Random-start APE)
Number Attention Heads
8
Embedding Dimension
512
Dimension per Head
64
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU
Normalization Layer
RMSNorm
Normalization Layer Position
PreNorm and PostNorm
Trained Lengths of Operands
1-40
Training Steps
50,000
Batch Size
200 (Ours), 800 (NoPE & Random-start APE)
Optimizer
Adam
Learning Rate (LR)
0.0001
LR Warm-up
Linear (From 0 to LR), 1% total steps
LR Cool-down
Cosine Decay (From LR to 0.1LR) 203 (Ours), 1023 (Random-star APE)
Maximum Position ID (max_pos)
Training Dataset Size
50,000 (Ours), 500,000 (Others) 100,000
Evaluation Dataset Size
Device
NVIDIA RTX A6000 48GB
Training Time
hours
26

---

RMSNorm Table 6: Hyperparameter summary for addition task with three summands (Figure 18).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
1-3 (Ours), 6 (NoPE & Random-stte APE)
Number of Attention Heads
4
Embedding Dimension
512
Dimension per Head
128
Hidden Width of Feed-forward Layer
2048
Activation Function of feed-forward Layer
GEGLU
Normalization Layer
Normalization Layer Position
PreNorm and PostNorm
Trained Lengths of Operands
1-40
Training Steps
50,000
Batch Size
1000 (Ours), 800 (Others)
Optimizer
Adam
Learning Rate (LR)
0.0001
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0, O.I) 102 (Ours), 1023 (Random-star APE)
LR Cool-down
Maximum Position ID (max_pos)
Training Dataset Size
1,000,000 100,000
Evaluation Dataset Size
Device
NVIDIA RTX A6000 48GB < 12 hours
Training Time
Table 7: Hyperparameter summary for copy/reverse task (Figure 20).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
(Ours), 6 (NoPE & Random-start APE) (Ours), 8 (NoPE & Random-start APE)
Number of Attention Heads
Embedding Dimension
512
Dimension per Head
128
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU
Normalization Layer
RMSNorm
Normalization Layer Position
PreNorm and PostNorm
Trained Lengths of Query
1-40
Training Steps
50,000
Batch Size
1000 (Ours), 500 (Others)
Optimizer
Adam
Learning Rate (LR)
0.0001
LR Warm-up
Linear (From to LR), 1 1% of total steps
LR Cool-down
Cosine Decay (From LR to 0.1LR) 301 (Ours), 601 (Random-starr APE)
Maximum Position ID (max_pos)
Training Dataset Size
1,000,000 100.000
Evaluation Dataset Size
Device
NVIDIA RTX A6000 48GB 8 hours
Training Time
27

---

Table 8: Hyperparameter summary for minesweeper generator task (Figures 10 and 21).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
4 (Ours), (NoPE)
Number of Attention Heads
8
Embedding Dimension
512
Dimension per Head
64
Hidden Width of Feed-forward Layer
2048
Activation Function of fordaro Layer
GEGLU
Normalization Layer
RMSNorm
Normalization Layer Position
PreNorm and PostNorm
Trained Lengths of Query
(5-9) x (5-9)
Training Steps
100,000
Batch Size
200
Optimizer
Adam
Learning Rate (LR)
0.0001 (Ours), 0.0002 (NoPE)
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0.ILR)
LR Cool-down
Maximum Position ID (max_pos)
15
Training Dataset Size
100,000
Evaluation Dataset Size
100,000
Device
NVIDIA RTX A6000 48GB
Training Time
S 30 hours
28

---

D Decoder-only Transformer Architecture
Here detail the architecture of a depth- I,, H--had decoder-only Transformer (Vaswani al., 2017). For a simple presentation, we ignore the normalization layers, as in Yun et al. (2020a).
Let V be the (ordered) vocabulary, set of all tokens. Given an input sequence I E VV and its length N, the encoding function Enc : V + _dxN maps it to
(0) Enc(Z).
It is a sum of the token embedding and the position embedding.
Next, there are L Transformer blocks that sequentially transform this input. We denote by Tf. RdxN _ RRx the operation of the 1-th block (1 E [T]),, so that
X 1 Tf, (x(
(2
The block Tf, consists of (causal) attention layer Att; Rd R and a (token-wise) feed-forward layer FF, Rd N
+ Rdx each of which contains residual connection:
Tf, = (id + FF,) o (id - Att;),
where we denote by id: Rd + RdxN identity map.
Each attention layer Att; consists of attention heads. Its h-th head (h E [H] has matrices @ k E R V E R xd and (/ E mdxdia as its parameters. 6 With these
matrices, borrowing the notation from Yun et al. (2020a), the attention layer with an input X E IRd can be written
H
At) = Σ
U
V
x softmax (KK)) X) Q.! x)
(4)
Here the softmax operator takes a square matrix M NxN and outputs an N N upper-triangular column-stochastic matrix
My
[softmax(M))
Eigre, Mip 1 (15)))
(5)
where 1 (8) is indicator function for predicate &: it equals if & is true and 0 otherwise. Note that the upper triangularity captures the auto-regressive behavior of the causal attention. For the sake of convenience, we denote by ( == (1-1) + Att (x-1)) E gaxN which a consequence of residual connection right after the attention layer.
Each feed-forward layer FF; is two-layer perceptron having w dRd x 600
E RR w
dd b E Rd as its parameters. It applies the following map to each column y of an input Y:
+ w
(1) ((W(")
(1) y + (19) + 5
(6)
where @ is a component-wise activation function. That is, the ford-forward layer is defined as
(Y) == o o(w(") Y + 1 6b 1
where 1. is the d-dimensional vectors filled with I's. Here we mainly the ReLU operation <(.) = max {-, 0} (Jarrett et al., 2009; Nair and Hinton. 2010), but there are many other popular choices such as GeLU (Hendrycks and Gimpel, 2016), GLU (Dauphin al., 2017), ReGLU, and GEGLU (Shazeer, 2020).
The final component of the Transformer model is the decoding function Dec R VN which is composed of a linear readout and a (token-wise) arg-max operation. Here, the linear readout is simply a linear layer having Wout E [gld as its parameter. The decoding function produces the output sequence
Dec(e (1)) E V.
'One can dø maxt, h max di Q.. a d constructions a bit messier with redundant entries O.
h as inner dimension each head. This makes our formal
'Every entry is non-negative and the sum of entries in each column is 1.
29

---

E Formal Construction of Addition Transformer with Position Coupling
Here we show how implement the addition by employing single-layer two-head decoder-only Transformer equipped with position coupling. We restate the theorem for the sake of readability.
Theorem 5.1. With the input format described in Section 3.1, there exists a depth-1 two-head decoder- only Transformer with coupled positions that solves the addition task with next-token prediction. Here, the operand length is at most 2l(d-17)/21 - 2, where the embedding dimension is d 2 21.
Organization of the Proof. A whole section dedicated to prove Theorem 5.1.
We start with the notation (Appendix E.1).
We review and formalize the format of the input sequence (zero-padding, reversed format, and wrapping with BOS/EOS) (Appendix E.2).
We define the encoding function Enc with a table of concrete example (Appendix E.3), where Enc maps an input sequence of length N to d N encoding matrix (()
We devote lot of pages to the detailed construction of the parameters of a causal attention layer Att; to generate desired attention patterns (Appendix E.4). The attention layer has two attention heads playing distinct roles: (1) preparing for sum without considering carries; and (2) preparing for the earry prediction EOS detection.
We provide construction of token-wise feed-forward neural network FF, which is two-layer ReLU network (Appendix E.5). It consists of two subnetworks playing different roles; (I) producing one-hot vectors, each of which indicates digit of the sum (response); and (2) binary values indicating whether the position is the end of the sequence.
We conclude the proof by defining the decoding function Dec which performs the linear readout and the arg-max operation to generate the output sequence (Appendix E.6).
We illustrate the roadmap of the proof in Figure 22.
E.1
Notation
For the architecture of the decoder-only Transformer, follow the notation introduced in Ap- pendix D.
Let e denote the i-th standard basis vector of Rd. For example, ei = the m m identity matrix. Let Op and 1p denote the p-dimensional vectors filled with O's and I's, respectively. Similarly, let 0mxn denote the m x n zero matrix. For a positive integer n, we frequently use the set [n] := {1...., n.. For any matrix A, denote the i-th row and j-th column of A by A. and A.j. respectively. Given two non-negative integers a and b, let E(a, b) be the length of longer one between and For example, E(12, 3456) =
Oj Let Im be
Consider an ordered vocabulary V (0, 1, 4, 6, 'S' that plays the role of both the sagiining--o-sequenee token.* 8 We denote V4 as k-th element of V. For instance, V4- 3 and V1s = $. Lastly, since we employ only one Transformer block, we omit the superscripts (I) the parameter matrices/vectors
7,8,
+, +, We include a special token (BOS) token and the end-of-sequence (EOS)
and the size of dimensions (1)
QK, and d
E.2 Input Sequence
We seek to perform an addition a + b = using next-token prediction. To this end, we want to transform it into an input sequence = $4 + B = C of an appropriate format. Note that the FOS token is the last token that needs to be predicted, so exclude EOS in the input sequence. Let é ;== ((a, b).
We first zero-pad the shorter one between a and 6 to match the length of the part A and part B é, Sometimes, the sum might be longer than a or b due to a carry. To make the length of the part C
BOS and EOS tokens do not need to be identical. We regard them as the same token just for the simplicity of the presentation.
30

---

Addition:
"653+49=702"
Input Sequence: (E.2)
Enc
Taken Embeddings
(E.4.2)
Attention Layer (E.4)
Input Encoding:
(
0
(E.4.1)
Head 2: Preparing
Position Embeddinge
Head Preparing Sum w/o Carriss
Carry Prediction &
Att
Besed on Position Coupling
Residual Connection
EOS Detection
PRE SUM
Tokan Embeddings
PRE SUM
FRE
Position Embeddings
Sub-NN 1: One-hot Vectors Indicating Sum
Sub-NN 2: Binary:
FF
Residual Connection
EOS or not
Token Embeddiings
'BUM"
PRE CARRY
X
PRE EDS
Position Embeddings
Feed-Forward Layer (E.5)
Dec
Output Sequence: (E.6)
20705
Figure 22: Roadmap to the formal construction of addition Transformer with position coupling.
consistent, also put zero-pad in front of to set its length as l + 1. Also, the addition with next-token prediction, we reverse the sum e to make the part C. For example, we have a sum 3812 + 98 = 3910, we use $3812 + 0098 = 01930 as an input sequence; if a sum 98 + 9907 = 10005 is given, we use $0098 + 9907 = 50001 as an input sequence. The red digits zero-paddings, and the blue digits are the reversed sum.
calculating
To recap, the input sequence I =
E V of length N 38 + consists six parts:
1. the BOS token 0i
"8'
2. the first operand A = 02 ... 0e+i where 0; E (o....,.)
3. the addition symbol 0e42 = 't';
4. the second operand B = T6+3 02442 where o; E (O....,9):
5. the equality symbol 02143 '=';
6. the (reversed) sum C' 02644 03844 where d; E (O.....9).
Note that the part C might be incomplete (i.e., N 38 + 4) at the inference time; infer the digits of the part C one by one using next-token prediction. Throughout this section on a formal only consider the train time setup in which we infer all the digits of the part C at once using simultaneous next-token prediction in a single forward pass. Precisely, we want to use an input sequence I = o.... ON to produce an output sequence O = o...oo' where
construction, however,
2 ... NN-1 C' 02t+1 ... N. and d'y '8' (EOS).
31

---

E.3
Encoding Function
We plan to produce an imput encoding, given an input sequence I designed as above. The encoding matrix ( is of size x N: each column represents embedding vector for token, while each row represents a particular named dimension. What we mean by named dimension is that we give name each dimension for clear description of our formal construction.
We construct an input encoding by concatenating the token embedding and the position embedding, which can be viewed
of two different embedding matrices of the same size.
Table Example initial encoding. Here we consider the input sequence $653 + 049 2070 and the starting position ID is chosen 2. The vectors vfj are defined in Equation (11). The gray rows will be filled later.
$
5
+
4
0
0
1: NUM
3
0
4
0
0
0
2: IS_BOS
1
0
0
0
0
0
0
0
0
0
0
0
0
3: FULL_ON ONES
1
1
1
1
1
1
4: PRE_ SUM
0
0
0
0
0
0
0
0
0
5: PRE CARRY
0
0
0
0
0
0
0
6: PRE EOS
0
0
0
0
0
0
0
0
7-16: SUM
10
10
10
10
10
10
10
0
10
10
0
10
10
0
10
10 0
17: IS EOS
0
0
0
0
0
0
0
0
18 18-(P- + 17): POS
0p
0
01
0
P
01
01
1
12
01
(P 18) + 17): POS_2
Op
04
0/
06
27
0
2/5
E.3.1 Token Embedding
The token embedding consists of 17 dimensions: we call them
1=NUM, 2=15 _BO__BBB_________ 3=FULL_OO ONES.
4=PRE_SUM, 5=PRF_ 6=PRF__OO0
(7....,16) -S.,,,,) and 17=1S_FOS. FOS.
Initially, we let the last 14 dimensions be empty (i.e., all zeros). Thus, we explain the first three dimensions, NUM, IS BOS, and FULL_ONES
.....,9) Dimension (NUM). For number token (O...., 9), we put itself in the dimension NUM. For the other tokens (+. =, S), we put O.
Dimension (IS_BOS). For special token " put in the dimension I$_BOS. Otherwise, put O.
Dimension 3 (FULI_ We put everywhere in this dimension.
E.3.2 Coupled Position IDs and Position Embedding
Before constructing position embedding, we specify the coupled position IDs for the addition task. Let namma__ppp be hyperparameter of the maximum position IDs, where position IDs are non-negative integers. Basically,
match the significance of the digits: e.g., least significant digit is always coupled to the other least significant digits. To this end, we first randomly choose a starting position ID E [max_ _oo_ppooss-
1]. (For that, max pos L 2 must hold.) Then we allocate the position IDs of token in the input sequence I 01
as
0,
i = 1, i = 2, .,. +
+ i 1,
p())
8 8 (E + 22),
i = l 3,
+ +3, i = 28 + 4, 38 1 4.
8 (38 4)
32

---

Recall that N = 38- + 4. Also, observe that for i 12+
p(i) p(i + p
+
(10)
which couples the position of (E - i 2)-th significant digit in the first operand (A), the second operand (B), and the sum (C). Also, the position of tokens '+' and '=' are coupled. Lastly, the only
token that has the position ID 0 is the special token '8'.
Before moving on to the positional embedding, we define UK (k E (2 as
of = (11))
( D
i=1 E R
(11)
where (,.bk)) is defined as the i-th (from left) digit of D-digit binary representation of l: For example, if D 2,
v =
-3 1-1 1I =1 -1]" of 1-1 -11"
(12)
We remark that the points o are the vertices of D D-dimensional hypercube with side length 2, centered at the origin." 9 Note that for k + L, | = D, (up uP) - 2.
(13)
Now we explain the position embedding. consists of 2P dimensions, which eventually become from 18-th to (2P + 17) dimension after concatenation. If p(i) 0, we let O2p as position embedding vector. For the positive position IDs p(1) 2 1, we let a concatenation
P(0
(14)
+()++1
position embedding vector of token de. (In case of p(1) =2p, we use of instead of oo - We call the former P dimensions for the position embedding as POS_ I and the latter P dimensions as P.
Concatenating the token embedding and the position embedding, get the input embedding X (0), See Table 9 for an example. As a result, the total embedding dimension is d = 2P + 17. Note the maximum possible position ID that can be represented with of maa__po = 22" 22)2-17)/21
Therefore, the length of an operand must be S max max_pos - 2 = (d 17)
E.4
Transformer Block Causal Attention Layer
The goal of the causal attention layer is to fill in the zero- sero-blanks' 10 of the encoding matrix at dimensions PRE_SUM, PRE_ TRE_CARRY, and PRE_ EOS. We divide the roles into two different heads.
E.4.1 Attention Head 1: Digit-wise Addition without Carries
The goal of the first head is to perform a digit wise addition and to fill in the blanks of the en- coding matrix at dimension DRE_SUM. Later, using this dimension, combined with the dimension PRE_CARRY, we will able to perform the next-token prediction for addition. For now, we do not care about the carries, which will be dealt with in later section. Formally, we aim to perform o; + 0++++1 for each E (2, ... l + }} and put its result at the (36 + - i)-th position (column) of the dimension PRE SUM (row). To this end, utilize position embedding.
Recall that d 2P + 17 and let dok,
1. Let M be number determined later. Let Opxp)
Q ( (VMP(cH7 (e )OEES
Opx17
VMIr
E Iddk.ixd
(15)
O
01XP)
K1 =
OPX17 17
OPxP
VMIP)
E RRdoK.id
(16)
MMP(e
Oxxr
01xP
"hhe choice of the vectors v2 not strict. They only need to have the same length and be distinguishable (for least constant order) in terms of inner products. That is, there should be noticeable difference between |elll and (of, oP)) for of L. Succh an idea of filling in the blacks of the encoding matrix is borrowed from the literature of RASP language(s) (Friedman et al., 2023: Lindner et al., 2023; Weiss et al., 2021; Zhou et al., 2024a). This can be done with the help of residual connections.
33

---

The linear transformations with matrices Q1 and K₁ do two different jobs at once. (1) Q1 (K1, resp.) takes the dimensions POS (POS 2, resp.) from the input encoding matrix and scale them up by VM; ; (2) Q1 (K,, resp.) takes the dimension FULL ONES (IS BS__,BO, resp.) and scale up by VMP. For concrete examples, please refer to Tables 12 and 13. By these, the attention score matrix C1 := (K, ( (T Qix becomes in Table 10. The blanks Table 10 are the numbers smaller than (P - 2); the asterisks (***) are the entries (or lower triangular submatrices) ignored by the causal sof tmax operator; the dots represents the hidden MP'..
Table 10: Exact attention score matrix C1 (with explicit row/column indices) of Head 1.
row col j ==1 2 3 t+2 e 8 2 28 + 3
32 + 2 32 + 3 36 +
MP MP MP MP MP MP MP
MP
MP
MP
*
MP
MP
MP
$
MP
MP
8+2 2
2+3
*
*
MP
MP
*
W
28 +
MP
#
#
*
#
*
38 +
38 +
11
38 + 4
Now consider the attention matrix A, softmax(Cii) NXN, Its exact form is bit messy due to the softmax operation of finite numbers. However, one can observe that, if the number M is large enough, gets close the column-stochastic matrix Ti NxN described Table 11. The blanks in Table 11 are zeros; the dots represent the omitted nonzero entries.
Table 11: Limiting attention matrix Ti (with explicit row/column indices) of Head 1, as M gets large.
row col j
2+2
E+3 l f+
+
3
2 38
+ 38 + 4
172
172
1
173
173
1/3
0
1/2
0
0
1/3
0
1/3
0
{+1
0
0
0
1/2
0
0
1/3
0
0
E+2
0
0
0
0
0
0
8+3
0
0
0
0
0
1/3
0
1/3
0
+ 2
0
0
0
0
1/3
0
28 3
0
0
0
0
0
0
0
0
0
38 +
0
0
Let R A - Ti E N be the error matrix, which is upper triangular. Its exact form is messy
as well, but we can obtain the bounds their entries. Consider a pair of indices (i, [N] such
that < Let 23 1/TT (1, If [Til +,2 Rily
and
MP
I2 MP + +ij - x M(P 22 2M + +(j (i=1
- [Rly S
(17)
On the other hand, if [ills 0. [Rilis and
e M(P-2)
Rikk S MP +(j x. 2
(18)
2.e2M + (j xzz)
34

---

Now let dy.1 = 1 and
Vi (( T Rdddiidd
(19)
U, = eRRE_SU Rdxdv.
(20)
The linear transformation with matrix U1 takes the dimension NUM from the input encoding matrix, scales it up by 3, and puts it to the dimension PRE_SUM. A concrete example is provided in Table 14.
Obtaining V (0) A1, its every entry is zero except at the dimension PRE_SUM. Observe that (V,V, X (0) (PRE_SUM)] = 0, because in the input encoding matrix, the dimension NUM starts with Also, note that enough to focus on the columns/ j E {22 + 3,..... 34 + 4] since we only care about the next-token prediction of the tokens after 02243 ='='. Specifying the dimension (i.e., the
particular row) for these columns, we have
x
(30+443)+4- + X (0)
(0)
36 + 2}, if E (26+ ,,....
[U] V T; (PRR_SUMM)
(NEM)
0
if E (3l 3,34 + 4),
(21)
if {22 .,, 36 2),
(22)
0
if E {3t + 3,36 4].
Refer to Table 15 for a concrete example of computing U V (0) T. Also, for the softmax errors,
[U (0) Rí (AR(FF__UMM)
Σ
3X (0)
(NUM) RRlii-
(23)
25151
Specifically, if E (26 3,...,,
3l + 23 (thus 2i 3),
[U V (0) R (PRR_SUM)) =
Σ
3x (0)
Ri
Σ
(x
R R1ly.
(NUM)s
(NUM)
(3(( 55
25i5j
(31+444) is (48+5)))
negalive
positive
(24)
where
2 2.9(j 3)
Σ
3X (0)
Rilis
(25)
(NUM)
3 + (j 3)
holds by Equation (17), and
05
Σ 2013
3x (0) (NUM)E
R1ig v
27(j - 3)
(26)
3e2M + (j 3)
(38-4)- o (48+5)-3
holds by Equation (18). On the other hand, if E {36 3,38 4].
27(j 1)
0 < UU V RII(PRF =
Σ
3x
(0)
(NUM); R Ri S
(27)
22M + (j- 1)
2015)
One can easily prove these inequalities by using the bounds of RRiiy' 's and the fact that the entries in x [NUM] lie in the interval [0,9].
If we let M > logg( N 1) + 3, we ensure that [U, V.X X (0) Ra (PRE_SUM)a) smaller than 0.1 for
each j E {2 + ... 36 + 4}. The proof is simple: it is enough to check
27(N - 3)
27(N - 1) 2M + (N - 1)
1
(28)
M + (N 3)
10
10
35

---

Table 12: Example of Q₁ X continuing from Table 9.
$
6
I-P:
OP
MM VM vM
Vileg vMug VMu Vilg
MMv VM vM VMP
MM VMP
V VMP
P+1:vM /MP
VMP VMP VMP VMP VMP VMP VMP
VMP
Table 13: Example of K X (0), continuing from Table 9.
I
$
5
3
+
0
4
2
0
7
0
I-P:
OP
VMel VM VMag
VM MM
VM
/ /M
M
v /Mug
v VMv
P+ 1:
V
MM
0
Table 14: Example of U₁ V1 X (0), continuing from Table 9.
I
$
6
5
3
+
0
4
2
0
7
I: NUM
0
0
0
0
0
0
0
0
0
0
0
0
2: 1$_
0
0
0
0
0
0
0
3: FULL_ ONES
0
0
0
0
0
0
0
0
0
0
4:
0
18
15
0
0
12
27
0
6
0
21
5: PRE
PRE__ARR CARRY
0
0
0
0
0
0
0
0
0
0
6: PRE EOS
0
0
0
0
0
0
7-16: SUM
010
0
0
10
10
10
10
0
10
0
10
10
10
010
10
010
17: IS S_EOS
0
0
0
0
0
0
18-end: POS_1,POS.
02P
02P
02P
02P
02P
02P
02P
02P
02P
02p
02p
02p
02P
Table 15: Example of U, continuing from Table 14. See Table 11 for the definition of T
I
$
6
5
3
+
0
4
2
0
7
1: NUM
0
0
0
0
0
0
0
0
0
0
0
2: IS_ BOS
0
0
0
0
0
0
0
0
3: FULL_ ONES
0
0
0
0
0
0
0
0
0
0
0
4: PRE_SUM
0
0
9
7.5
4.5
0
6
12
9
6
0
5: PRE_CARRY
0
0
0
0
0
0
0
0
0
0
0
6: PRE_EOS
0
0
0
0
0
0
0
0
0
0
7-16: SUM
110
0
10
0
10
10
0
10
10
0
10
0
10
10
10
10
10
010
17: IS EOS
0
0
0
0
0
0
18-end: POS_I.POS_2
02P
02P
02P
02P
02P
O2P
02P
02P
02P
O2P
O2p
02p
02P
E.4.2
Attention Head 2: Carry & EOS Detection
The goal of the second head is to fill in the blanks of the encoding matrix at dimensions PRE_ _EE_R and PRE _es_ At dimension PRE EOS, we will put (approximately) EOS token ('S"), otherwise, we will put strictly smaller numbers like (approximately) 2/3 and 1/2.
the next token would be the
What we will put at dimension PRE_CARE is the evidence of the presence of an additional carry, which is not quite straightforward to understand. Let us take a look at some examples. Consider an addition + 9 12. Since is greater than or equal 10, the least significant digits in the operands generate carry 1. But in some cases, a pair of digits with sum less than 10 can make carry. Next, consider an addition 53 + 49- = 102. In the second least significant digits, An addition of and occurs. However, carry is already produced in the least significant digits (3
12), so the total sum including the earry is 10, not Thus, it also produces a carry. But how can we know the presence of carry while only looking at the second least significant digits? The answer to observe
36

---

the second least significant digit in the sum, O of 102. Somehow, the consequence of adding 5 and 4 (or 10, implicitly) so it makes carry.
To generalize this explanation, let a and be digits of the operands in the same significance, and e be digit of the sum in the same significance addition of a and b generates a carry is that
and b. We find that the rule of recognizing that the
then a carry is generated, a b- {9, 10), Otherwise,
(29)
then the carry is not generated.
Thus, crucial to store the information + b or any related one somewhere. In fact, we can store + b + at dimension a+ c and used later in the forward layer. Formally, we aim to perform o +0+++++111111 + 03c45-8 for each i {2, .... L + 1} and put its result at the (3l + 5 i'-th position (column) of the dimension PRE_ EIERY (row). ) To this end, we again utilize our position embedding.
of the encoding matrix, and it can be transformed into
Recall that d = 2P + 17 and let dok, = P Let
Q2
( ( (EPLLL
Opx17
VMIP 0Pxp
E RdQxaxd
(30)
OIxP
01XP
K₂
OPX17 MMP(e BOS)
VMIP
0pxP
Rddokaxd
(31)
O1xP
01xP
The linear transformations with matrices Q2 and K₂ do two different jobs at once. (1) they take the dimensions POS from the input encoding matrix and scale them up by VM; (2) Q2 (K2, resp.) takes the dimension FULL_ ONES (IS BOS, resp.) and scale up by VMP. For concrete examples, refer to Tables 18 and 19. By these, the attention score matrix C2 == (x2 ) Q2
becomes
as in Table 16. The blanks in Table 16 are the numbers less than equal to M(P 2); the asterisks (**)) are the entries (or lower triangular submatrices) ignored by the causal sof tmax operator; the dots represent the hidden MP's.
Table 16: Exact attention score matrix C2 (with explicit row/column indices) of Head 2.
row col i=1
3 + 3 38 4 t+1 e+2 f+3 2l +2 2 28 26 +4
MP
MP MP
MP MP MP MP MP
MP
MP
2
MP
MP
MP
E 1
MP
MP
MP
E 2
#
MP
MP
2+3
MP
MP
28 +2 2
W
MP
MP
28 3
MP
28 4
#
#
&
MP
#
38 3
MP
38
W
MP
Now consider the attention matrix A == softmax(C NxN Similarly to the previous head, the number M is large enough, it gets close to the column-stochastic matrix T2 E x described Table 17. The blanks in Table are zeros; the dots represent the omitted nonzero entries.
Let R2 = A T2 E N be the error matrix, which is upper triangular as well. Its exact form is messy well, but can obtain the bounds of their entries. Consider pair of indices (i.j) [N] such that < j. Let vt = T222 E (1,2, 4). If (Talsp = - [Rlli < 0 and
MP
I x. 2 (j xz)))
- <
(32)
Zj MP (j M(P
37

---

Table 17: Limiting attention matrix T2 (with explicit row/column indices) Head 2, as M gets large.
row col
j =
2
t+1
2+
2
8+3
28 +
2
2
+
26
+ 4
3 + 3
3 + 4
1/2
1/2
1/2
1/3
1/3
1/3
1/4
1/4
1/2
2
1/2
0
0
1/3
0
0
0
1/4
0
8+1
1/2
0
1/3
1/4
0
8 2
1/2
O
0
1/3
0
0
0
8+3
*
*
1/3
1/4
0
#
*
26 + 2
1/3
0
1/4
0
0
28 +3 3
#
#
*
*
1/3
0
26+4 4
1/4
0
38 +3 3
W
W
*
1/4
0
36 + 4
1/2
On the other hand, if [Talis- = 0, Ralis > and
(((22)
R2lú S 2zeNP + (j 2. (1-
(33)
xye2M + (j xz))
Now let dy.2 2 and
V23 = 4(e (2 BOS)
E IR. d
(34)
U² (e (CRRE_CARRY CARRY
cd CiRE_EOS
E Rdxdv.2.
(35)
The linear combination with matrix U2V2 does two jobs at once. First, takes the dimension NUM from the encoding matrix, scales it up by 4, and puts it to the dimension TRE_CARRY Second, it takes the dimension IS_BOS BOS from the encoding matrix, scales it up by 2, and puts to the dimension PRE EOS. TRE_EOS. A concrete example is provided in Table 20.
Obtaining U2 V2 (0) Ag, its every entry is zero except at the dimensions PRE_ CAERY and PRE_EOS. Observe that [Ug V2 (2)x(()) CARRY) 0, because in the input encoding matrix, the dimension NUM starts with O. Also, note again that enough to focus on the columns j {2( + 3.... 3 + 4}, since we only care about the next-token prediction of the tokens after 02 = '='. Specifying the
dimensions (i.e., the particular rows) for these columns, we have
(NUM)(F++2) + -x(
(xIM))
(2/ + 3) 28 + 3,
UU2 V. (1) T3. (PRE_CAMRY))
(0) (xUM)(3r45-g) (0)
(NNM)(44+6-9) + (0)
if (2f 1,... ... + 3). if = 32 +4,
[NUNy
0
(36)
0
if (2f 3,33 4}.
(37)
(a6+4))- + (4840))- to
if E (2f 44. ,,...., 38 3),
2/3 if = 28 +
U V2 (0) (oS))
1/2 if (2f + 4,....,3.. +3,
(38)
if 3 4.
Refer to Table 21 for concrete example of computing V,V (T),,, Also, for the softmax errors,
UU2VXX (0)
R (PRF_
4x (0) (NUM)
R1lg-
(39)
25i5)
U (0) R
R2]((PE__OS))
Σ
2x (0)
R1ly.
(40)
[18 BB_y
38

---

Let us first obtain bound of the softmax error term at dimension PRE PRE_CARRY Jf = 26 3, since
X (0) x
X (0)
(26+3) O.
(NUM)
12 (0) (PRE (PRE__CAAARRY) CARRY)(()R))
Σ
4X (0)
(NUM) [Rils
(41)
25 7
and
36(2()
0
Σ
4X
(NUM)) RRilis -
(42)
3. + 2
251 1964+2
j E (28 .,....,3
38+ 33,
[U2 (22 R2l (PMEE___C))RY)
Σ
4x (0)
Rilij
Σ
4x (0)
[R]]is
(NUM)
(NUMs)s
2<i<j-1
(3(3++5)
nogativee
(#(++6)
posslvve
(43)
where
4x Rilis S 3-9(j - 4)
Σ
0
(NUM))
4c2M + (j- - 4)
(44)
and
36(j - 44)
05
Σ
4X (0)
(NUM)) Rilis
(45)
4 + (j -
2
i7 (36+55)
i (48+0))-
And if = 38 + = N,
1U2V; X (0) R2 (PRE__RRRREE CARRY)N
4x (0)
RRINNN
4 (0)
R J,N,
(46)
[NUM)N
(NUM)s
2515N-1
negative
paxitive
where
18(N -2)
0 < -4x (0)
(NUM) RIINN S 202M + N --2
(47)
and
36(N 2)
05
Σ
4X (0) (NUM)
[R]
IN 2º2M N 2
(48)
2-1SN-1
Next, we obtain a bound of the softmax error term at dimension PRE EOS. Since
Σ
2x (0)
(1S (IS_ Rilig
2x (0)
(is_nos) R
(49)
13i5
the error term can be bounded
2(j 3)
if 28+ +3
3(322M + - 3)
2(j 4)
0 < - UU (0) R2 (PRO_COS)) S
if (20 + 4,...,33.
34 14..,,3++33 3},
(50)
4(4c2M +j- - 4)
(j -2
if = +4+
2c2M + 2
39

---

We then can ensure that both |U2V2X (0) R2 (PRF, (PRF_ suM)a] and |U2V2X (0) R2l (PRF, [os]_ smaller than 0.1 for each j E {2e ,, + 4). by letting M 2 1/2 log(N) + 3. The proof is similar to the one that is presented for head 1.
Table 18: Example of Q2 (0), continuing from Table 9.
I
$
5
+
4
0
0
1-P:
OP MM VMug vMog vMog V VMof
VMog Vilf VMof VM
V
VMuf
P+1: VMP VMP VMP VMP VMP
VMP VMP VMP
VMP
VMP
VMP
VMP
VMP
Table 19: Example of K2 X (0), continuing from Table 9.
$
5
4
0
0
1-P:
Op
VMug v Vau
V
v /M
v VM
VM
v VMug
v VMeg vf
VMug
v VMug
P+1: VMP
0
0
0
0
0
0
Table 20: Example of U2 2 (0) continuing from Table 9.
5
3
+
0
4
7
1: NUM
0
0
0
0
0
0
0
0
0
IS BOS
0
0
0
0
0
0
0
3: FULL FULL_ONES ONES
0
0
0
0
0
0
0
0
0
0
0
PRE_SUM SUM
0
0
0
0
0
0
5: PRE_CARRY
0
24
20
12
0
0
16
36
8
0
28
0
PRE_ROS ROS
0
0
0
0
7-16: SUM
010
010
0
010
110
10
010
010
010
010
010
010
010
17: IS_FOS
0
0
0
18-end: POS_I
02p 02P
02P
02P
02P
02P
02p
02P
02P
02P
02P
O2P
02p
Table 21: Example of U2 V2 X (0)
T2 continuing from Table 20. See Table 17 for definition of T.
5
3
+
0
0
7
0
1: NUM
0
0
0
0
0
0
0
0
0
0
0
IS_ _SOS
0
0
0
0
0
0
3: FULL_ONES ONES
0
0
0
0
0
0
0
0
0
0
0
PRE_SUM
0
0
0
0
5: PRE_CARRS
0
12
10
6
0
8
12
16
0
14
9
13
0
PRE_ROS
2/3
2/3
2/3
2/3
1/2
1/2
7-16: SUM
110
0
0 10
0
010
010
010
10
0 010
0 10
0
010
010
17: IS EOS
0
0
0
0
18 18-end: POS
02P
02p
02P
02P
02P
02p
O2P
02P
02P
02P
02P
02P
02p
E.4.3 Residual Connection
So far we have computed the output of Att; operation. Passing through the residual connection, the output of the attention layer the sum of the original input encoding matrix and the output of Att operation:
Σ
U U. XO) + Σ U.V.
X R
(51)
hE(1.2
bE(1.2)
softmax ciTor term
40

---

Since the term E1e(1 U.V. has nonzero entries only at dimensions PRE_ [EE__ARRY and PRE_EOS, the residual connection plays role of "filling in some blanks" in the input encoding matrix. A concrete example of the output of residual connection presented in Table 22, ignoring the softmax error term, whose entries have an absolute value smaller than 0.1.
Table 22:
Example output of residual connection, continuing from Tables 9, 15 and 21. Here we ignore the softmax error terms the orange rows. The gray rows will filled later.
I
$
6
5
4
9
2
0
7
0
NUM
0
5
0
0
4
9
2
0
7
2: IS_ BOS
1
0
0
0
0
0
0
0
0
0
0
0
0
FULI__ONE,
4: PRE_SUM
0
0
9
7.5
4.5
0
6
9
12
9
6
0
0
FRE_CARRY
12
0
8
12
16
14
9
13
6: PRE_EOS
2
1
1
1
2/3
2/3
2/3
2/3
1/2
1/2
1/2
7-16: SUM
010
10
0 10
10
0 10
0
0
0
0
0
00
0
00
17: IS_
0
0
0
0
0
0
0
0
0
0
0
0
0
18-(P + 17): POS 1
0p
4
44
02
(P + 18)-(2P + 17): POS 2
0p
04
066
07
01
06
07
066
0
03
E.5
Transformer Block - TToen--wie Feed-forward Layer
The goal of the feed-forward layer to fill the blanks of the encoding matrix at dimensions SUM and IS EOS. Be careful that the ferd-forward layer can only implement token-wise mappings: token- wise mapping takes inputs only from the entries in the same column of the encoding matrix. Besides, the architecture of our feed-forward layer (except for the residual connection) is one-hidden-layer ReLU network.
For token for i {2e + 3,..... 36 + 3} (from '=' token to the second ), we will put standard unit vector _ dimensions SUM if the next token is & {0. 99).
Recall from the discussion in Appendix E,4.2 that we can judge whether a carry generated at a certain position by exploiting only the digits (of the operands and the sum) in the same significance. Bringing the notation, let and be digits of the operands in the same significance, and e be a digit of the sum in the same significance as a and b. Then the rule of recognizing that the addition of and b generates a carry is that
If Otherwise: if + b (-1,0), then the carry not generated.
C (9, 10}.
then a carry is generated,
(52)
A simple case analysis shows that the value of + b c must be one of 11 O, 9, and 10. Let us briefly check this claim our example:
6 - 7-
no carry from 6+0
(53)
5+4 - - 99; 9;
there is carry from 5+4
(54)
3+9 - = 10.
there is a carry from 3+9
(55)
Recall that noisy version of a + b + is already stored dimension PRE_ CARRY of (1), and is exactly at dimension NUM. Thus, we can (approximately) implement a + c for token by
y!!) ((PFF_CARM) (URE_CARRRY)
2Y (NUM)))
(56)
This is kind of token-wise linear transform, so we do not need to consume any hidden layer (with ReLU activation Φ) implement it.
Combining with y (PRE (PRF_sUM))" a noisy version addition without carry, we can indeed implement the addition. Note that a digit-wise addition should be done as
digit-wise addition = (addition without carry + 1 (carry propagates)) mod 10.
(57)
We first describe the formal construction of feed-forward network FF1 for dimensions SUM and SS_S and then explain the intuition behind the construction. For the example result of applying the fred-forward network is presented in Table 23.
41

---

E.5,1 Subnetwork 1: Construction for SUM (dimension 7-16).
Given vector y [yll== E
Rd define linear function g IR
+ R8
UPRE_CARRY CARRY
2YNUM
y/s + 1/4 2y1
(
YPRE_SUM
+ +.21
0.21
(58)
10
10
and consider a one-hidden-layer ReLU network fo R + R (k = 0, 1, .............................. 9) defined as
(z) = 2 20( - (k 0.5)) o(x k) o(x (I 0.5)) + 6(x - (1: + 1)) +xx (A: 9.5)) - Φ(x - (k: + 10)) o(x (1: 10.5)) o(x (I 11))
(59)
Then we construct subnetwork of our fced-forwart network for token σy by
FF1 (y( (SUM))
[[o (g (x.!"))
So (1 (y (2)
(60)
fx(x)
2
5
8
10
11
12
13
14
(x(x)
2
11
12
14
f2(x)
2
11
12
14
(x)
Figure 23: Example plots of (xx defined in Equation (59). (k: = 0, 1,2,3 3)
Explanation. The purpose of the first subnetwork is to generate a 10-dimensional one-hot vector whose position of indicates the next digit: e for the answer of next-token prediction 'k' There are two cases where we need to predict the next token as 'k"
Case 1: (Addition without carry) mod 10 and no carry propagates.
Case 2: (Addition without carry) = mod 10 and there is propagating carry 1.
In the first case, due the softmax error (with magnitude most 0.1),
(0)
0.1, k + 0.1] n [A + 9.9, k 10.1),
(61)
(PRE_SUM))
(0)
2Y (0)
(NUM) E [-1.11 0.9 n [-0.1.0.1] c [-1.11 0.1].
(62)
[PRE_CARRY]
In the second case, again due to the softmax error (with magnitude most 0.1).
Y(0) (0)
(PRE_SUM)) + 1 [A: 0.1. + + [K 9.9, A 10.1].
(63)
Y
(1) 10 [-1.1, - 0 c (-1.1.0.]]
(64)
(PRE_CARRY)
(NUM),
42

---

In both cases,
Y (0)
2Y- (0)
(0) (PRE_SUM)
(PRE CARRY)j
(NUM))
+ 0.21 E [k, + n 1 + 10, k + 10.32]
(65)
10
C + 0 [ + 10, A 10.5).
(66)
We can map the column +f This job is done by the function 9. Note that the resulting sets [k, 0.5) n 10, k + 10.5] are
the set [k, k 0.5) +0, k 10.5] if the next token is J- 'k'
disjoint for different k's.
Recall that our objective is output I the dimension k + 6 (among the dimensions 7, .... 16 in SUM) and to output O to the other dimensions in SUM if we need to predict 'k" as the next and to map the
token. To this end, is enough to map the set k. + 0.5] n [ 10, k + 10.5] other sets (for different k's) to O. This can be done by a ReLU network fa(2)) is a ReLU network having two bumps at intervals [k 0.5, 1] and [k 9.5, k + 11]. In particular, fx( E [k: k + 0.5] U [k 10, k 10.5]: see Figure 23 for an illustration.
Lastly, we have desired one-hot vector output for each by taking composition between g and Vo((),...,J())) as written in Equation (60).
E.5.2 Subnetwork 2: Construction for Is EOS (dimension 17).
We move the dimension IS EOS. For token o, for j E {2f ....., 32 + 4), & is the next token, we will put (k=8) to dimension IS $__oo: if A is the special token 's' and 0 otherwise. To this end, we define ReLU network h R - IR as
h(x)= = 10g (æ 0.8) - 10$ (x - 0.9).
(67)
Then, we can construct subnetwork of our feed-forward network for token oy by
(os_) = (r (iRE__00))) FFi (r
(68)
Explanation. Note that for columns E {26 + ,, + 4), if we consider the presence of softmax errors with magnitude at most 0.1, the values that x (PRE_EOS); can have lie in the set [0.4, 0.6] n [2/3 0.1, + 0.1] n [0.9, 1.1] C (-ox,08) n oo). We want to output if (
BOS)_ 2 0.9 and 0 otherwise: this can be done with the ReLU network h with two neurons.
Remarks:
In total, we consume 8 x 2 82 ReL.U neurons in our feed-forward network FF1. However, it is possible to construct the addition Transformer with a smaller number of neurons, with slight modification in the linear readout of the decoding function (Appendix E.6).
Unlike in the attention layer, now we do not have to worry about softmax errors in the output since the forward ReLU network plays the role of denoiser.
Table 23: Example output after applying the fied-forrwrd network.
+
2
1: NUM
0
0
0
0
0
0
0
0
2: S__
0
0
0
0
0
0
0
0
0
3: FULL ONES
0
0
0
0
0
0
0
0
0
0
4: PRE SUM
0
0
0
0
0
5: PRE CARRY
0
0
0
0
0
0
0
0
0
0
6: PRE_FOS
0
0
0
7-16: SUM
e
100
eg
e5
10
10
10
e3
10
10
e
10
e1
10
e
17: IS EOS
0
18 18-(P + 17): POS 1
0p
(P+ (2P 17): POS 2
0p
43

---

E.5.3
Residual Connection
The last task of the feed-forward layer is pass FFi (Y through the residual connection. As result, we have
x = Y FFi (r(v).
(69)
concrete example of the output the second residual connection is showcased in Table 24.
Table 24: Example output of residual connection, continuing from Table 23, Here we ignore the softmax error terms in the orange rows.
9
=
2
0
7
I: NUM
0
0
0
4
9
0
2
0
7
0
2: IS_BOS
0
0
0
0
0
0
0
3: FULL ONES
1
1
1
1
1
1
1
1
4: PRE SUM
4.5
9
12
9
5: PRE CARRY
0
12
10
6
0
8
12
16
0
14
9
13
0
6:
2/3
2/3
2/3
2/3
1/2
7-16: SUM
e
e
e10
eg
c5
E2
10
67
10
e
10
E3
10
eg
e
17: __OOSS
0
0
0
0
0
0
18- + 17): POS
0p
0/
04
P
P
0g
18) 17): POS 2
0p
06
04
87
U/
04
0/
E.6 Decoding Function
As mentioned in Appendix D, the decoding function performs a linear readout (with a weight matrix Wout IR Vixd) and a (token-wise) arg-max max operation. That is,
(x()) = (Vk,)i= E V.
Dec
(70)
where V48 the k-th element of V and
arg max { W x((x = -[o mrll"
(71)
The objective of the decoding function is to perform a proper next-token prediction for addition. especially utilizing the dimensions SUM and _S_FEO of (1)
We now construct the weight matrix W Wout- For a token o;, if the value of dimension IS EOS of X is O, then the linear readout output the dimensions SUM as it is to return one of a number token (0-9). the other hand, the value of dimension IS IS_EOS 1, then the linear readout outputs a large number (like 100 for example) for the token 'S' to return EOS (8). This can be implemented by the weight matrix W described in Table 25. Also, an example of applying the linear transform showcased in Table 26.
44

---

Table 25: The transposed weight matrix W
of the linear readout in decoding function.
out
V
2
3
4
+
$
1-6: NUM-PRE EEO
0g
0g
0
06
0g
0g
0g
0g
0g
06
0g
0
0g
7: SUM1
1
0
0
0
0
0
0
0
0
0
0
0
8: SUM₂
0
0
0
0
0
9: SUM3
0
0
0
0
0
0
0
10: SUM4
0
0
0
0
0
0
0
0
11: SUM5
0
0
0
0
0
0
0
0
12: SUMG
0
0
0
0
0
0
0
0
13: SUM7
0
0
0
0
0
0
14: SUMs
0
0
0
0
0
0
0
0
0
15: SUMg
0
0
0
0
0
0
0
16: SUM10
0
0
0
0
0
0
0
17: IS EOS
100
18-end: POS 1, POS 02r
02F
2r
OZP
02r
02r
02r
02r
02r
02r
O2p
02r
02r
Table 26: Example output of linear readout (W ( continuing from Tables 24 and 25. The yellow cells represent the maximum value of each column, from the '=' token's column to the rightmost column (used for next-token prediction).
5
7
0
1
0
0
0
0
0
0
0
1
0
I
0
0
0
0
0
0
2
0
0
0
0
0
0
0
0
0
0
3
0
0
4
0
0
0
0
1
0
0
0
0
0
0
0
5
0
0
0
6
0
0
0
1
0
0
0
0
0
7
0
0
0
0
8
0
0
0
0
0
0
9
0
0
0
0
0
0
0
11
0
0
0
0
0
100
100
100
100
100
0
100
Table 27: Example output sequence O Dec (x continuing from Table 26. The yellow cells in
the bottom row exactly predict the next tokens.
6
3
9
2
8
$
6
9
0
45

---

F Impossibility of Addition with No Positional Encoding
For the sake of readability, restate the proposition below.
Proposition 5.2. Consider any depth-1 finite-head decoder-only Transformer model T without positional encoding (NoPE). Given an input sequence and its arbitrary permutation I', the last tokens of I and I' are identical, then the next tokens predicted by TT will also be identical for both sequences when applying greedy decoding scheme.
Remark. We assume the lllayer (L= = 1) H-head Transformer achitecture specified in Appendix D. Although omits normalization layers, we remark that Proposition 5.2 remains valid even for the architecture with standard layer normalization (Ba et al., 2016) or its variants (e.g., Zhang and Sennrich, 2019).
Proof. We keep following the notation about matrices introduced in Appendix E.I. Throughout the proof, we denote the value/vector/matrixx related I' by appending it.
Let encoding matrices generated from the input sequences I' E V as
X Enc(T) RaxN and X' Enc(I') RdxN
(72)
Since there is positional encoding, the encoding function Enc(-) maps the same tokens the same columns. In particular, Z; = T3 implies X.; X.j. Since we assume that T' is a permutation of such that IN = x, there exists bijection π [N] [N] such that I Iz(i) for each i E [N] and T = N. Then, it follows that X', (=(i)) for each and, specifically, X.N = X.N.
Recall that the single H-head attention layer Att RaxN + IR operates as Att(X) = Σ (X) where the attention head h is defined as
Head; (X) == Un V. softmax ((, TQ,X) E RdxN
where Qh, Kn E R.dK xd Vn E dR xd and Up E Rdxdv,
Claim: Headd, (X) )lNN = [Head, (ieaa, (X')), for all l E [M].
The claim suffices to prove the proposition because of the following: first, the claim implies that the last (N-th) columns of the attention layer outputs are the same, i.e., [Att(X)) = AAt(('') (X')),N Note that the operations after the attention layer-residus connections, FF, and Dec-all operate in token-wise column-by-column manner: the j-th column of the output of a token-wise operation is a function of j-th column of the input for the operation. Therefore, the last column of the attention layer output totally determines the next-token prediction at N-th input token. As a result, the predicted next-tokens are the same for I and I'.
The rest of the proof is devoted to proving the aforementioned claim. Fix any [H]. Let
[softmax ((K,)) ((X,))))X)))
[81
SN]
(73)
sofftaxx XIK, Qx))) = |s] s']'
(74)
which are both stochastic (sum to 1) column vectors. Considering that we are taking the last column of the softmax output, it follows that s' $=()) for each E [N]: this can be proved by applying the definition of the softmax operation and the fact
X') Qa X X K, Qu. T m0)) Q.X.N. [(,)) T Q
(T())))
(75)
Consequently, since
N
N
s(X',
Σ (t(1) X. ( = s.Xo
(76)
1=1
we have
X'. [softmax
Q (Xx,''))'''''']l__l) = X /softmax Qnx))...
(77)
Therefore, the claim holds. This concludes the proof.
46

---

Here, we provide the Python code that calculates the maximum possible exact-match accuracy that 1-layer Transformer with NoPE can achieve for the m-digit addition problem.
from itertools import product from collections import defaultdict
total Change here
counter_dict = iifaultdictddica
for b in product (product (range (10), repeat-n), product (range (10),
repeat-m)):
1f a [0] -- 0 or b[o] 0: continue
10
total += tuple (srrted
12
amm (n) join a)))
13
b_num in (ot..""..nnnpppp.p... (mtp b)))
14
ab_sum a_num b_num
if ab_sum in counter_dict(c]:
counter_dict[c]] [ab_sum] else:
17
18
counter dict [c] [ab_ sum]
19
count sum (max (d. values ()) for d in counter_dict items ()
22
print ("m =",
23
print ("Permutation Invariant Additions Count:", count)
24
print ("
Total m-digit Additions Count:", total)
print ("
Ratio:", count / total)
27 п..
28
[Example Outputs]
31
Permutation Invariant Additions Count:
32
Total m- m-digit Additions Count: 81
33
Ratio: 1.0
Permutation Invariant Additions
Count: 2668
36
Total m-digit Additions
Count: 8100
37
Ratio: 0.329382716004982774
38
-
3
Permutation Invariant Additions
Count: 50150
40
Total m-digit Additions
Count: 810000
41
Ratio: 0.06191358022691988
42
Permutation Invariant Additions
Count: 765139
Total m-digit Additions
Count: 81000000
45
Ratio: 0.009446160049882711
46 =
47
Permutation Invariant Additions
Count: 10033314
Total m m-digit Additions
Count: 8100000000
49
Ratio: 0.0012988807407407407
50
47

---

G (Formal) Construction of N x Multiplication Transformer with Position Coupling
Here we show how to implement the N x 2 multiplication using a depth-2 decoder-only Transformer equipped with position coupling. Our construction involves 3 heads in the first Transformer block and 7 heads in the second Transformer block, requiring a total of 10 heads.
Theorem 6.1. Given an appropriate format of the input sequence, there exists a depth-2 decoder-only Transformer model with coupled positions that can perform the N x 2 multiplication task with next-token prediction. Here, the number of the total heads is 10 and the length of the first operand is at most 2 (dd34))66) 3, where we denote the token embedding dimension by d 2 46.
We note that construction for the N 2 multiplication task permits the of multiple FFN layers at the second decoder block. However, we believe that there exists a potential improvement in our construction, wherein single FFN layer could suffice for each decoder block, leveraging the expressivity of the neural network. Additionally, we do not provide a detailed error analysis but assume that the softmax operation with sufficiently large attention weights reduce small attention scores to zero values, thereby clearly revealing the desired attention patterns.
G.1 Notation
Consider an ordered vocabulary V = (,,1,,,,,4,5,,,,,,,, 8)= We include a special token '8' that plays the role of both the beginning-of-sequeen token. We denote V2 as k-th element of V. For instance, V4 and V13 $. Unlike the addition task, our construction for the multiplication involves multiple layers and hence we do not omit the superscripts (!) in the parameter matrices/vectors and the size of dimensions.
(BOS) token and the end-of-sequence (EOS)
G.2 Input Sequence
Our objective is to next-token prediction for implementing x b bb=c To this end, want to transform into an input sequence I SA x B == of an appropriate format. Let la and & represent the lengths of a and b, respectively, and we denote their sum as L = la + l. While our immediate focus on the case where 4 2, it worth noting that our approach can be extended to the case where l, > 2. as the key insight for the construction does not rely on &,. Thus, we present the input sequence and encoding function in a more general form applicable to &, 2
Unlike the addition case, we do not zero-pad both a and b. Instead, we only zero-pad the response, as the length of e may either equal the of the lengths of and b, or be less than the sum of their lengths by 1. Hence, we zero-pad in front of for the latter case to fix the length of by f. We also reverse the response e to make the part C. For instance, if we have 312 x 24 7488, the input sequence transforms $312 x 24 88470. we have 589 62 36518, then the input sequence would be $589 x 62 = 81563. The red digit is zero-padding, and the blue digits are the reversed product.
To recap, the input sequence I 0102 0102...00 .N. V of length N 28 + 3 consists of six parts:
1. the BOS token 0 = '8"
2. the first operand A ... 0f, where o; {O,...,9): 9);
3. the multiplication symbol 02,t2 'x:;
4. the second operand B = 17,,33... 0242 (note that e = la + Eb) where o; E (0,...,)9;
5. the equality symbol c2+3 ===;
6. the (reversed) product C' = oe+1 02443 where E ((....,,)).
Note that the part C might be incomplete (i.e., N 28 + 3) at the inference time; we infer the digits of the part C' one by one using next-token prediction. Throughout this section on a formal construction, however, we only consider the train time setup in which we infer all the digits of the part C' at once using simultaneous next-token prediction in a single forward pass. Precisely, we want to use an input sequence I = ...... NN to produce an output sequence O = .o..ooy where 0'43 ... ON-I V.ya...O'y_ C' 0i+4 ... ON and d'y "8' (EOS).
48

---

G.3
Encoding Function
We now explain the imput embedding for given an input sequence I designed as above. The embedding matrix X is of size x N: each column represents embedding vector for token, while each row represents particular named dimension. We concatenate the token embedding and the position embedding, which can be viewed sum of two different embedding matrices of the same size.
Table 28: Example initial encoding. Here consider the input sequence $7595 x 79 500006 and the starting position ID is chosen as & = 1. The vectors f are defined in Equation (79). The gray rows will be filled in later.
I
S
7
5
9
5
X
7
9
5
0
0
0
6
I: NUM
0
7
5
9
5
0
7
9
0
5
0
0
0
6
2: FULL ONES
1
1
1
1
1
1
1
1
1
1
1
1
3: IS BOS
1
0
0
0
0
0
0
0
0
0
0
0
4: IS MUL
0
0
0
0
5: IS_ LEQUALL
0
0
0
1
0
0
6: IS_ IS_OP2_ONE
0
0
0
0
0
0
0
0
0
0
0
7: 1$_OP2_ 1___P___ENN TEN
0
0
0
0
0
0
0
0
0
0
8: OP2 ONE
0
9: OP2_TEN TEN
10: OP1 SHIFTO
0
11: OP1 ISI_TI
12: OP1 SHIFT2
13: SHIFT3
14: OP1_ OPI_SHIFT4
15: RESULTI
16: RESULT2
17: RESULI3
18: RESULI4
19:
20: PRE_CARRY
21: PRE_BOS
22: PRE_EOST
23-32: PROD
010
010
10
10
10
10
10
10
10
10
33: I$_EOS.
0
0
0
0
34: MASK
0
0
0
0
0
0
35-(P+34): POS 2 MASK
0p
0p
OP
0P
0P
(P435)-(PP+34); P
0p 0/
074
17
17
07
10
12
(2P+35)(3P+34) POS 2
Or vi
06
16
17
(3P435)(1P433); pos_3
Op ug
06
(4P435)(6P434) POS
Op 0g
wio
(6P435) P
Op
17
U/
Vio
Vi
vio
vn
vio
Ug
of
G.3.1 Token Embedding
The token embedding consists of (34 + P) dimensions, where P represents the dimension for the position embedding which will be described in the very next section. While the token embedding dimension for the addition task was independent of P. our construction strategy for the multiplication task involves copying the position embedding into the token embedding. This is why we have the P term in our token embedding dimension. For the first 34 dimensions, we label them as:
1=NUM, 2=FULL_ONESS 3=1S BOS, 4=1S MUL, 5=1S EEU_LL
6=IS_OP2_ONE, 1iiiS__P2__E____ 8=0P2_ ONE, 9=OP2_TEN,
11=OPI_SHHIFTH 10-:0P1_SIIP2,,__, 1:03=O11_SHIIITT 14=0P1_ SSIFF4,
15=RESULTI, 16=RESULT2, =7=RESULT3, 18=RESULT4.
19=PRE_PROD, 20= 2OR 21=172_ 22
(23....32))PRODD, 33=15 EOS, 34=MASK,
49

---

and for the last P dimensions (f35, ...., 34 + PJ), we named them as POS 2 MASK.
The initial token embedding fills only NUM, FUIL, ONES, _o_,,, iS MUI, and is_ S_FAU,,, leaving the other (29 + P) dimensions empty (i.e., all zeros). These (29 + P) dimensions will be filled by passing through the layers. Here we describe how fill the first dimensions.
Dimension (NUM). For number token (0....,9), we put itself into the dimension NUM. For other tokens (x,=88) 8), we put O.
Dimension (FULL_ONES), We put everywhere in this dimension.
Dimension 3 (IS_BOS)) For a special token 'S', we put 1 into the dimension I$_BOS. Otherwise,
put 0.
Dimension (IS_MUL.. For a special token 'x', we put into the dimension IS IS_MUL. Otherwise, we put 0.
Dimension (IS [I__EQUL). For special token '=', put into the dimension iS_ IS_EQUAL. Otherwise, we put O.
G.3.2 Coupled Position IDs and Position Embedding
We now specify the allocation of coupled position IDs for the N M multiplication task as the following: given an input sequence 01 ON.
0,
i i= 1,
+ i 1-i-2+60 lo.
i = ...... La + 2, i La +3, ... l 1+ +3,
p(1)
(78)
8 s 3,
8 --1+3+22 2 i= + 28 2
Compared to the addition case, the position allocating function p becomes more complicated since the length of two operands can be different, but the core remains simple: coupling the position IDs for the least significant digit in the first operand (A), the second operand (B), and the result (C), and then decreasing the IDs as the digit position increases for each A, B, and C.
Now explain the position embedding. We utilize the same of (k: E [22) defined for the addition task, specifically
2h = (-1))(00)) D
E R
(79)
where 6, (D.k) is defined as the i-th (from left) digit of D-digit binary representation of 1. Using
D we design the position embedding for each position ID p(i) by
u & P()
200+
v p(0+2
(80)
p(i)+
+4])+4
The first P dimensions of the position embedding are named as POS_ and subsequent sets of P. dimensions are named as POS 2, POS_ POS 4, and POS 5, respectively, Thus, the position embedding is 5/-dimensional vector. In case p(i)+ +
+j( (j E [4]) exceeding 2"", we use v p- instead of u P P(())5" If p()) 0, let Ogp position embedding vector.
By concatenating the token embedding and the position embedding, get the input embedding (0)), Specifically, the position embedding is placed under the token embedding ((P + 355 to (6P + 34)-th dimension). See Table for an example. As result, the total embedding dimension is d = 6P + 34. Note the maximum possible position ID that can be represented with of is max max_pos 2 (2d-344) Therefore, the length of the first operand must be la < max pos
i 1 2(d-31/0 - tb 1. For the case when &s 2, this inequality becomes &, < 2 2[14-31/6] 3.
50

---

G.4 Construction Idea
Here, we provide an example that demonstrates how we construct the N 2 multiplication. Consider the calculation 7595 x 79 = 600005. While typical method for computing such a multiplication is illustrated in Table 29, we consider an alternative approach, as shown in Table 30. In this method, we pair the digits from the first and second operands at each step where the sum of their digit positions the same, and then calculate the sum of the pairwise products. For example, the number 116 in Table 30 is generated by X 9 + 5 X 7, and the number 108 is generated by 5 x 9 9++ 7, where blue indicates numbers from the first operand and red indicates numbers from the second operand. The main reason for considering such method is to provide clearer intuition for determining which numbers from each operand we should attend to when predicting the next token.
Table 30: Multiplication II
Table 29: Multiplication
7
5
9
5
X
7
9
5
5
X
7
9
1
6
6
8
3
5
1
0
8
5
3
6
9
8
6
0
0
0
0
5
4
6
0
0
0 0 5
Suppose the current input sequence is $7595 79 5000. During this step, the model is tasked with predicting 0 (the 0 just before 6) for the next token. As illustrated in Table 30, this 0 is computed from the
of 9, 9, 1, and additional representing the carry from the previous step. Similar to the explanation in E.4.2, we highlight that the carry can be detected by computing (hundreds digit of 116) (current token): yielding a result of 9, indicating the occurrence of a carry 1.
8 (ones digit of 98) + 0 (tens digit of 108)
In summaryy the correct prediction of the next token 0 (the 0 just before 6) can be achieved by summing the main summation part and the carry part, where the main summation part is computed using 49, 98, 108, and the carry part is calculated using 98, 108, and 116. Additionally, it's noteworthy to detail the breakdowns:
49 0 x 9 + 7x
98 = 7 x 9 + 5+7,
108 5 x 9 9
116 9 x 9 +
Thus, for predicting the next token, we need 0, 7, 5, 9, 5, 7. Here, we highlight that this structure, requiring 5 consecutive tokens from the first operand and every token from the second operand for the next-token prediction, remains unchanged for any prediction time and any query length.
As will see in the later subsection, depth-2 decoder-only Transformer model can be constructed to TEN by 7, and SPP__SHIFTU to OPl SSSHIF44 by 0, 7, 5, 9, and 5, respectively. One may be concerned that 0 is not given in the first operand the input sequence. This requirement of 0 beyond the most significant digit arises in the later stage of the prediction, i.e., predicting the token that is near the most significant digit of the response. Although 0 is not explicitly given in the first operand, our construction can automatically manage as if the O were originally at the start of the first operand. A similar situation occurs in the early stage of the prediction that 0 is required before the least significant digit of the first operand, and our construction is also capable of handling this issue.
fill OP2 ONN by 9, OP2_TEN
Consequently, the embedding vector of the current token (the 0 preceding 60) will be structured as the most table in Table 31, with some irrelevant dimensions omitted for readability. We then utilize a feed-forward layer to
RESULTI with OP1, SHIFTO OP2_ONE ONE oPI_SHIFTI OP2 OEN,
51

---

RESULT2 with SPI_SHIFT] X OP2_ONE + SPI_SHIFT2 X OP2_TEN,
RESULT3 with OP1 2P__HHIFT2 OP2_ONE OPI_SHIFT3 OP2_TEN,
RESULT4 with OPI_SHIFT3 x OP2_ONE + OPI __SHS_FFT4 x OP2_TEN.
The result is illustrated in the center table of Table 31. Next, we employ an additional feed-forward layer to fill
PRE_ PROD with ones digit of RESULTI + tens digit of RESULT2 + hundreds digit of RESULT3,
PRE_CARRY with ones digit of RESULT2 + tens digit of RESULT3 + hundreds digit of RESULT4.
These computations yield the result illustrated the right-most table of Table 31. Once this process is done, we can finally predict the next token by the following two steps:
0, if CRE_Y NUM {-2, -1, 0},
CARRY
1, if PRE_CARRY NUM E {8, 9, 10),
2, PRE_CARRY - NUM E { {18, 19, 20),
= PRE_PROU + CARRY (mod 10).
Table 31: Illustration of the construction idea.
I
0
I: NUM
I: NUM
I: NUM
2: FULL_ONES
1
2: FULL_ONES
2: FULL_ONES
3: S_B
0
3: IS BOS
0
3: IS BOS
MUL.
0
4: MUL
4: MUL
IS EQUAL
5: IS. EQUAL
5: IS. EQUAL
8: OP2. OP2_ONE ONE
9
8:
8: OP2. ONE
9: OP2 TEN
7
9: OP2
7
9: OP2
7
10: SHIFO
0
10: OP1 SHIFTO
10: OP1 SHIFTO
OP1_SHIFTI
11: OPI_SHIFTI
11: OP1 OPI_SHIFTI
OP1_SHIFT2
12: 1P1_OPP1_TTTTTT
12: 101_OPP1SSTTTT
13: OPI_SHIFT3
9
13: OPI_SHIFT3
13: OPI_SHIFT3
14: OP1_SHIFT4
5
14: OP1__HIFTA
14: 1P1__OIHTTT
15: RESULTI
15: RESULTI
49
15: RESULTI
49
RESULT2
RESULT2
16: RESULT2
17: RESULT3
17: RESULT3
108
17: RESULT3
108
18: RESULT4
18: RESULT4
116
18: RESULT4
116
19: PRE_PROD
19: PRE_PROD
19: PRE_PROD
19
20: CAR_Y
20: PRE_CARRY
20: PRE_CARRY
(P+35)-2P434: 1
(P+35)-27+30) POS..
0.
(P+35-224400) POS.
03
(27435-37744) P
(2P435)(9P430) _os_2
(2P435-88P430 POS
0
(3P435-(1P430) P
(PP45)-(PP3)) rOS.3
vg
(3P439- POS.3
Ug
(4P435)-5PP300 P
(AP43)-65443) POS
Vg
(4P439)65P43) POS
Ug
P
(6 POS_5
(5P435)-6P44)
G.5 Transformer Block - Causal Attention Layer
To implement the concept introduced in Appendix G.4, it is essential to design a Transformer block capable of generating an embedding matrix depicted the left-most table of Table 31. The goal of the first Transformer block is to fill ii__s_____ $$_______________________ OON (6-th dimension) and IS S__oo2__r__ TEN (7-th dimension) by if the token corresponds the ones or tens digit of the second operand, respectively, and otherwise. These two dimensions enable the filling of OP2_ ONE (8-th dimension) and OP2_TEN second Transformer block. Furthermore, we will MASK (34-th dimension) in the first block, which will serve as base for filling OP1 1ip___SFTOO to OP1 _HHFFT4 in the second block. Thus, we currently have objectives(ts_op2 be addressed by individual head.
(9-th dimension)
IS S__ TEN, MASK), each of which will
52

---

G.5.1
Attention Head 1: Detecting the Ones Digit of the Second Operand
The goal of the first head is to make the dimension IS I$_OP2 0P200P____0 ONE as a one-hot row vector, where is placed only at the token corresponding the ones digit of the second operand.
Recall that d = 6P + 34 and let doK,1 = P + Let M > be a sufficiently large positive real number. Let
1
O (Px(P+34)
OpxP
VMIp
0pxp
OpxP
OpxP
BRdok.ixd
P431 VMP (eFULL_ONES)
O1XP
O1xP
0 01xP
0 01xP
01xP
(81)
k
(P+3) P431
VMIP
OpxP
Opxp
0pxp
Opxp
E RRdxxaxd
(82)
VMP
O1XP
0 OIXP
O1XP
O1XP
0 01xP
Unlike the construction for the addition task, we do not provide the table for the exact matrix and detailed error analysis due to their complex characterization. Instead, provide an illustrative example for each step. We will also simply regard M as sufficiently large real scalar and thus the attention values can be clearly separated after going through the softmax operation.
The matrix Q ( maps the embedding matrix X into query matrix x())(()) R(P+1)xN, where the first P rows are obtained by copying from the dimensions POS_2 2 and scaling by VM, while the last row is the copy of the dimension FULL ONES scaled by /MP. Similarly, the matrix KL maps the embedding matrix to key matrix K X E [R In this case, the first P obtained by copying from the dimensions POS and scaled by vM, with the last row being the dimension iS BOS, scaled by VMP. For concrete examples, refer to Tables 32 and 33.
x
Table 32: Example of Q ()), continuing from Table 28.
I
$
7
5 9
5
7
5 0
I-P:
V vMlgg VMug Vflf VM vM vWlf vNoo vMuy VMlg vMuf VMW VMly VMlv VMP VMP VMP VMP VMP VMP VMP VMP /MP JMP /MP JMP /MP VMP
P+
VMP
Table 33: Example of K X (0), continuing from Table 28.
I
7 5 9
5
7
5 0 O
6
I-P: vMuy Vfuff vMuf vuseg of varef Ja vnuf vuuu VMuuf vMuff Vadus /ad /M P+1:/MP
By these, the attention score matrix c(2) = (K (1) x (0)) Q."h) x (0) and the attention matrix A == softmax((())) E NxN can be obtained. We provide the example of A(" Table 34. Specifically, entry in ( (1) is non-zero if and only if the inner product between the query and key vectors equals MP.
Now let dy 1 and define
(11) 2[(eis_MuL el ((S_)QU))) T E IRdvuxxd
(83)
(1) ele eis_or2_oNe E ddddv.1
(84)
The matrix U.7 (,(1) (0) takes the dimension IS_ _S_MUL and IS EQUAL from the embedding matrix (0) subtracts one from the other, scales the result by 2, and puts it to the dimension is $S_PP_2___ _UM__UMM.
Consequently, the matrix U.!)) vi x ((1)) is matrix that matches the size of the input embedding matrix X and is filled with zeroes, except for unique located at the ones place of the second
operand the input sequence, in the dimension IS_ $__OP___O0 ONN (6-th). concrete example provided in Tables 35 and 36.
53

---

Table 34: Example of A11 (with explicit row/column indices and sufficiently large M), continuing
(1)
from Tables 32 and 33.
row col 2
4
6
8
9
10
11
12
13
14
15
1
1
1
1/2
1/2
1
1/3
1/4
1/4
1/3
1/3
1/2
2
0
0
0
0
1/3
3
0
0
0
0
0
0
0
0
0
0
0
0
1/3
0
0
0
0
0
0
1/4
0
0
5
0
0
0
0
0
1/2
0
0
1/4
0
0
0
0
1/2
0
1/3
0
0
0
0
0
0
0
O
1/4
0
0
0
00000
0
0
0
0
0
0
1/3
1/4
O
0
0
0
0
0
0
11
10
0
0
00
0
0
00000
0
0
000000000
0
0
000000000
1
1/4
1
1
1/4
0
0
0
0
12
0
000000000
0 0
00000
0
1/3
0
13
0
0
1/3
14
0
00
0
1/2
15
0
0
Table 35: Example of ( v-(1 (0), continuing from Table 28, (Irrelevant dimensions are omitted for readability)
7
5
X
7
5
0
6
1: NUM
0
0
0
0
0
0
0
0
FULL_ONES
0
0
0
0
0
0
0
0
0
3: IS IS_BOS
0
0
0
0
0
0
0
0
0
IS IS_MUL
0
0
0
0
0
0
0
0
0
5: IS EQUAL
0
0
0
0
0
0
0
0
0
0
0
IS $_0P2_ ONE
0
0
0
0
0
Table 36: Example of U v (1) are omitted for readability)
X (0)
( continuing from Tables 34 and 35. (Irrelevant dimensions
I
S
7
5
5
7
9
5
0
0
0
6
I: NUM
0
0
0
0
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
IS BOS
0
0
0
0
0
0
0
0
0
4: IS MUL
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
IS ES_ALLLL
0
0
0
6: IS_ [S_OP2_ONT)
0
0
0
0
0
0
0
0
0
0
G.5.2 Attention Head 2: Detecting the Tens Digit of the Second Operand
In the previous head, set the dimension IS 10l_0P2_ON___ ONE (6-th dimension) to one-hot vector, where is placed only in the token corresponding to the ones digit of the second operand. The objective of Attention head is to fill the dimension IS ___oP_____ TEN (7-th dimension) similarly to IS IS_OP2_ONE, but with placed only in the tens digit of the second operand.
The design of the query, key, and value weight is not significantly different from the previous head. Compared the construction of attention head I, only push / M Ip to the next block for designing
54

---

0 (1) Specifically, 2 and K2 are defined
2
0 (px(F+34)
OpxP
OPxP
MMIp
OpxP
0pxP
Rdox..xxd
VMP [FFULL_ ONESS
01xP
01xP
01xP
0 01xP
01xP
(85)
K2
(P4())34) VMP (eis_
VMIP
OpxP
OpxP
0pxP
0pxP
E RRoxx,1xd
(86)
01xP
01xP
01xP
01xP
01xP
where 12 is set to P 1. We refer to Tables 37 and 38 for specific examples.
Table 37: Example of Q X (0), continuing from Table 28.
I
7 5 9
X
5 0
Or VMug vMvg VMuf vMog /uuu /M Jurog vMog JMug /Mvf vMug /Mog //Mu /Mog P+1.AM VMP VMP VMP VMP VMP VMP VMP VMP VMP VMP VMP VMP VMP VMF
I-P:
Table 38: Example of K! X (0), continuing from Table 28.
I
5
5
I-P: 0p vMug vMof vMug vMug oMfff vMlef vi vMoff v vMee volff Moof /u VMlef
P P+1.vM LMM
0
0
Q²¹¹ X (0) and the attention matrix _ == softmax(Cx") NxN can be obtained, and the example of A is provided in Table 39.
By these, the attention score matrix c. (K x
Table 39: Example of A (with explicit row/column indices and sufficiently large M), continuing from Tables 37 and 38.
(I)
row cool 2
3
4
6
7
8
9
10
11
12
13
14
15
1
1
1
1
1/2
1
1
1/3
1/4
1/4
1/3
1/3
2
0
0
0
0
0
1/3
3
0
0
0
0
0
0
0
0
0
0
0
0
0
1/3
0
0
0
0
1/4
0
5
0
0
0
0
0
0
0
0
0
0
1/4
0
0
0
0
1/2
0
1/3
0
0
O
0
0
0
0
0
0
0
0
O
1
1/4
0
0
0
0
0
0
O
1/4
0
11
10
9
0 0
0000
0000
0
0
000002000
0
00000
0
0
0
0
00000
0
0
0000
0
0
0
-0000
1/3
20000
1/4
1/4
0
0
0
0
12
0
1/3
13
0
0
00003
0000
0
1/3
00000
14
0
0
15
0
Finally, we set (1( and U.! the same to that of the previous head. That is, with dy.12 = 1,
v 2 2(e
FOuA))) Iddv,1ad
(87)
U (1) et TEN Rxdxxy
(88)
and the example of U.)) V (1) (0) and U V x A is provided in Tables 40 and 41. Con- (1) (1)
sequently, the matrix U. V
x A (1) is matrix that matches the size of the input embedding
55

---

matrix and is filled with zeroes, except for unique located at the tens place of the second operand in the input sequence, with the dimension is 1S_ TEN (7-th dimension).
Table 40: Example of U2 for readability)
V2 X (0) continuing from Table 28. (Irrelevant dimensions are omitted
(1)
I
$
7
9
5
7
9
5
0
6
1: NUM
0
0
0
0
0
0
0
0
0
0
0
0
2: FULL_ ONEES
0
0
0
0
0
0
0
0
0
0
0
0
3: _S_BOS
0
0
0
0
0
0
0
4: IS 3S_MUL
0
0
0
0
0
0
0
0
0
0
0
0
5: B_BQUAL
0
0
0
0
0
0
0
7: $s__s1 OP2 TEN
0
0
0
0
0
0
-2
0
0
0
Table 41: Example of U V2 are omitted for readability)
())
11)
X
A
continuing from Tables 39 and 40. (Irrelevant dimensions
I
7
5
7
5
0
6
I: NUM
0
0
0
0
0
0
2: FULL ONES
0
0
0
0
0
0
0
0
0
0
0
0
3: IS __BOS
0
0
0
0
0
0
0
4: IS MUL
0
0
0
0
0
0
0
0
0
0
0
0
5: IS_ [S_EQUAL
0
0
0
0
0
0
0
0
7: IS_ 1S_OP___________ TEN
0
0
0
0
0
0
0
G.5.3
Attention Head 3: Position Masking
The goal of Attention head is to generate binary mask the dimension MASK (34-th dimension), with '0' placed before the multiplication symbol (x) and '1' placed starting from the multiplication symbol to the end.
To this end, we set d2K,13 and design query and key weights by
Q (1)
(e ooNes) )
(89)
K 3
(1) (ei)_Mu)) E Rdde.aad
(90)
The matrices Q (0) and ((()) X (0) take the dimension FULL_ONES and IS_MUL, respectively, from the input embedding matrix. For concrete examples, please refer to Tables 42 and 43.
Table 42: Example of Q!' (0), continuing from Table 28.
I
7
5
9
5
X
7
5
0
0
0
6
1
1
1
1
Table 43; Example of K x
(0) continuing from Table 28.
I S
9 5
1: 0
0
0
By these, the attention score matrix ((()) == (K x TQ!!! (0) and the attention matrix
A (1) := (oftmax(C((1) NxN can be obtained and the example of A is provided in Table 44.
56

---

Table 44: Example of Ag (with explicit row/column indices), continuing from Tables 42 and 43.
(1)
j
2
3
4
5
6
7
8
9
10
11
12
13
14
15
row cool
1
1
1
1
1
1
1
1/2
1/3
1/4
1/5
0
0
0
0
0
0
0
0
0
2
0
1/2
1/3
1/4
1/5
0
0
0
0
3
0
0
1/3
1/4
1/5
0
0
0
0
0
0
0
0
0
4
a
O
O
1/4
1/5
0
0
0
0
0
0
5
O
0
O
O
1/5
0
0
0
0
0
0
0
0
0
6
O
1
1
1
1
1
1
1
7
0
0
1
0
0
2
0
0
0
8
O
0
0
0
0
000002000
10
11
9
0
000
00000
0
0
0
0
0
00000
000002000
0
00000
0
0
00010
0
00000
0
0
0000
0
0
00000
00000
0
0
00000
-0000
12
0
00000
13
0
14
0
15
Finally, we set ( and U by dv,13 = and
(2) (eis (us. E Rdvaaxd
(91)
U MASK E 12dxdv113
(92)
The example of U/ V (1) (0) and U V. x (0) A is provided in Tables 45 and 46. Conse- (1)
quently, the matrix Ug V X (0) A is matrix that matches the size of the input embedding matrix and is filled with only at the dimension MASK (34-th dimension) starting from the token to the end of sequence, and 0 otherwise.
At this point, the objective of attention head 3 may somewhat unclear. We note that the output of Attention head 3 will be utilized to fill the dimensions POS POS_2_M 2_MASK the subsequent FFN layer, and this POS 2 MASK plays erucial role in designing the key matrices in the Attention heads 3 to at the second Transformer block.
Table 45: Example of ( V (1) (0), continuing from Table 28. (Irrelevant dimensions are omitted for readability)
I
$
5
9
X
=
0
1: NUM
0
0
0
0
0
0
0
0
0
0
0
0
2: FOESSOONES
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_BOS
0
0
0
0
0
0
0
0
0
0
0
4: IS_MUL
0
0
0
0
0
0
0
0
0
0
5: IS SS_EQUAL
0
0
0
0
0
0
0
0
0
0
0
0
34: MASK
0
0
0
0
0
0
0
Table 46: Example of US are omitted for readability)
V3
(1)
X
A
(1)
continuing from Tables 44 and 45. (Irrelevant dimensions
I
$
5
9
X
9
0
0
0
1: NUM
0
0
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_BOS
0
0
0
0
0
0
0
0
4: IS _ML
0
0
0
0
0
0
0
0
0
0
0
0
0
5: IS EQUAL
0
0
0
0
0
0
0
0
34: MASK
0
0
0
0
57

---

G.5,4 Residual Connection
So far we have computed the output of Att; operation. Passing through the residual connection, the output of the attention layer becomes the sum of the original input embedding matrix and the output of Att; operation:
y x
Σ
U V.
A
(93)
kE(1,2,3)
An example of the output of residual connection is presented in Table 47.
Table 47: Example output of residual connection, continuing from Tables 28, 36, 41 and 46. Uncolored rows represent the initial embedding. Yellow rows indicate the rows filled by the attention heads in the first Transformer block. A pink row indicates the row that will be filled by the subsequent FFN layer.
I
$
7
5
9
5
7
9
5
0
0
0
6
I: NUM
0
7
5
9
7
9
0
5
0
0
6
2: FULL_ONES
1
1
1
1
1
1
1
3: IS_BOS
0
0
0
0
0
0
0
0
0
0
4: IS S_MML.
0
0
0
0
0
0
0
0
0
0
0
0
5: IS EQUAL
0
0
0
0
0
1
0
0
6: IS 1S_OPP2__2 _P2 ONE
ONE
0
0
0
0
0
0
0
0
0
0
7: [S_OP2_TTEN
0
0
0
0
0
34: MASK
0
0
0
0
0
1
1
1
35-(P+34): POS
MASK
0p
Op
OP
P
OP
OP
P
OP
OP
0p
P
0p
0g
(P+35)- POS
0p U
2
2
26
(2P+35)-(3P-34)) POS_2
Op o
17
17
11
47
u
(3P+35) POS_3
Op v
t
07
0
11
0
0.
87
v
v
(4P+35) POS_4
0p 0
0/
0g
(5P435)(6P-34) POS_5
0p 07
0g
0g
vi0
Vil
0g
V10
Vi1
V10
0g
0g
0f
0g
G.6 Transformer Block - Token-wise Feed-forward Layer
The goal of the feed-forward layer involves filling the dimensions POS _ MASK. Specifically, for each token dg if the dimension MASK is (i.e., y(
(MASKj+ = 1). we want to fill the dimensions POS 2_ MASK by copying the the corresponding token's 's POS otherwise, want to fill with Op. Be careful that the feed-forward operation is restricted to token-wi wise mapping, meaning it only takes inputs from entries within the same column of the encoding matrix.
Construction for POs_2_MAST Given vector y (951ş=1 E Rdd, define functions gr, Ju Rd + IR for every E [P] as
(y) hy) := Yros 2.1 22/MAKK
Yros_2, 2yMMSK
(94)
(95)
where Yross E IR is the 1-th dimension yros_ E RR. (1 E 1,2..... P). Consider simple oaehhidden-layer ReLU networks fi Rdd + R. defined
fi(y) = d(g(t)) - (z)(
Using the fact that yros 2,1 is either -1 or 1, we can easily check that f(y)) Yros_ if YMaSK is 0, and fi(y) = 0 if YMASK is 1.
Now, we can construct the width-2P feed-forward network that outputs the desired value at the dimension POS_ 2_____MAR MASK by
FFi (y
[ (Y:!!)
fr (Y.!))]
ex1x
(96)
(Pos 2 MASKK)
58

---

and 0 for any other dimensions. The example output for this layer is presented in Table 48.
Table 48; Example output of FFN layer at the first Transformer block, continuing from Table 47.
I
$ 7
7 9
0 0
35-(P+34); POS_ rOs_2_MASK
Op 0i 0g
P
of 0p Op 0p 0p Op Op Op Op 0p 0p
G.6.1 Residual Connection
The last task of the feed- fordafr layer is to pass FF1 (Y(I) through the residual connection, As result, we have
x y + FFi (r()).
(97)
This is the end of the first Transformer block, and a concrete example of X is illustrated in Table 49.
Table 49: Example embedding matrix after the first Transformer block. The yellow rows represent the results introduced during the first block, while the gray rows will be filled in the second block.
I
$
7
5
9
5
7
9
5
0
0
0
6
I: NUM
0
7
5
9
5
0
7
9
0
5
0
0
0
0
6
2: FULL_ONES
HIL
1
1
1
1
1
1
1
1
1
3: IS. BOS
0
0
0
4: IS_MUL
0
0
0
0
0
0
0
5: IS ESSUALL
O
0
0
0
0
0
1
0
0
0
0
6: 1S_OP___P000___ ONE
O
0
0
0
0
1
0
0
0
0
7: IS_ 2__0P22___ TEN
0
0
0
0
8: OP2 OP2_ONE ONE
0
9: OP2_TEN
0
0
10: OP1 SHIFTO
0
11: SHIFTI
12: OP1 SHIFT2
13: OPI_SHIFT3
14: OP1 OP1_SHIFT4
15: RESULTI
16: RESULI2
17: RESULI3
18: RESULI4
19: PRE_PROD
20: PRE_CARRY
PRE_EOS
PRE ERS_
23-32: PROD
010
010
110
O10
O10
010
O10
O10
O10
O10
33: IS _S_OOS
0
0
0
O
0
34: MASK
0
0
0
0
35-(P+34): POS _ MASK
OP
O
OP
0 P
OP
0P
Op
OP
O
(P4:5)-(2P433); POS_
Or 0g
17
17
(2P435)(3PP434); POS
Op v."
03
0
0
(3P435)(1P+34) POS
Op ug
(4P+35) _
Op
V7
Vn
P
vn
Vip
(5P435) P
Op
vil
vio
vil
vio
of
G.7 Transformer Block - Causal Attention Layer
Consider scenario where the model is at the step of predicting the i-th least significant digit of the multiplication result. There are two goals for the causal attention layer at the second Transformer block. The first goal is to generate the embedding matrix as the left-most figure in Table 31, that is,
59

---

fill OP2 ONE, OP2 TEN, and OPI_SHIFTO to OPl SHIFT4 with the ones digit of the second operand, the tens digit of the second operand, and the i. (i 1). (i 2). (i-3), 3), 4)-th least significant digit of the first operand, respectively. Our construction assigns each head to each dimension. The second goal is to fill PRR_ROS] and PRE__OSS _os2 with appropriate values. These 2 dimensions will be utilized in the subsequent FFN layer to predict whether we should predict the next token as EOS or not. Also, we note that filling these dimensions can be implemented within the same head for iei_FI and OP! _1IF1 respectively, thus requiring total of seven heads.
G.7.1 Attention Head 1: Copying the Ones Digit of the Second Operand
The objective of Attention head is to fill the dimension OP2_ONE with the ones digit of the second operand. To do we design the weights by defining dQ and
Q = (eTULL
E Rdgkaax
(98)
K(2) = (eis_ (eis_or2_one) ONE)
E R.qkkax
(99)
We also define dy.21 1 and
V122) = (e E Bgdvaax
(100)
U = ed
ONF__ONE E R
(101)
concrete example of (22))0000))))) X(1, x(), AP. U v X(),, and U ( (1) { is pro- vided in Tables 50 to 54. One might be concerned that in Table 54, the dimension OP2 OOE is completely filled with '9', but only the latter part. However, we note that given our focus on
next-token prediction, it suffices to accurately fill values starting from the = token, and filling the preceding tokens with placeholder values does not cause any issues.
Table 50: Example of Q x (1), continuing from Table 49,
7
X
7
6
1:
1 1 1 1
1 1 1
1
1
Table 51: Example of K
(1), continuing from Table 49,
$
7
9
X
9
5
0
0
1:
0
0 0
0
0
0
1
0
0
0
0
Table 53: Example of U v2 (1), continuing from Table 49, (Irrelevant dimensions are omitted for readability)
$
5
X
9
0
I: NUM
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_BOS
0
0
0
0
0
4: MUL
0
0
0
0
0
0
0
0
0
0
0
0
5: IS EQUAL
0
0
0
0
0
8: OP2. ONE
0
5
9
0
9
0
0
60

---

Table 52: Example of A (with explicit row/column indices), continuing from Tables 50 and 51.
(2)
j
2
3
4
6
8
9
10
11
12
13
14
15
row col
1
1
1
1
1
1
1/2
1/3
1/4
1/5
1/6
1/7
0
0
0
0
0
0
2
0
1/2
1/3
1/4
1/5
1/6
1/7
0
0
0
3
0
0
1/3
1/4
1/5
1/6
1/7
0
0
0
0
0
0
0
4
0
O
1/4
1/5
1/6
1/7
0
0
0
0
0
5
O
0
O
O
1/5
1/6
1/7
O
0
0
0
0
0
0
6
O
DO
D
1/6
1/7
O
0
00
0
0
OOOD
7
0
O
1/7
29090
181111100
0
11111110
0
8
0
O
2
10
11
9
0
000
00000
0
00000
0
-0000
0
00000
0
0
00000
0
00000
0 1000
00000
00000
00000
0
00000
0000
00000
00000
12
0
13
1
000
0
14
0
15
Table 54: Example of U are omitted for readability)
V
2)
X
A
continuing from Tables 52 and 53. (Irrelevant dimensions
I
$
X
0
0
1: NUM
0
0
0
0
0
0
2: FULL ONES
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_BOS
0
0
0
0
0
0
0
0
0
4: IS_MUL.
0
0
0
0
0
0
0
0
0
0
0
0
5: IS [EQUULLL
0
0
0
0
0
0
0
0
8: OP2_ONE
0
4
21/4
33/7
9
9
9
9
9
G.7.2 Attention Head 2: Copying the Tens Digit of the Second Operand
The objective of Attention head 2 is to fill the dimension OP2_TEN with the tens digit of the second operand. We take a similar approach to Attention head 1, but the main difference is that we utilize the dimension IS 1__OP2_TEN instead of IS [S_OP2_ONE for generating the key weight. We design the weights by defining doK,2 = 1 and
Q ()) = (eFULL
E Rddx
(102)
K = -
(is_OP2_TEN)
T
E 2Rd
(103)
We also define dy.22 = and
(2) = (ex T E Rd
(104)
0 = d
1 TEN E pdxdv.22
(105)
(2) is pro- vided in Tables 55 to 59. Once again, the dimension OP2. TEN is not entirely filled with '7' in Table 59. As mentioned in the previous head, this does not cause any issues because the front part (before =) does not affect the final prediction unless additional attention blocks are introduced after the second Transformer block.
concrete example of 982² (1), x1), K. (2) x A2, U, V2 x and U, V X (1) A
Table 55: Example of Q (1), continuing from Table 49.
?
||$
I
$
7
5
9
5
X
7
5
0
0
0
0
6
1:
1
1
1
1
1
1
61

---

Table 56: Example of K x continuing from Table 49.
(2)
I
$
7
5
9
5
7 9 5
0
0
0
0
1:
0
0
0
0
0
0
Table 57: Example of A
(with explicit row/column indices), continuing from Tables and 56.
2
3
4
5
6
7
8
9
10
11
12
13
14
15
row coll
1
1
1
1
1
1
1
1
1/2
1/3
1
1/5
1/6
0
0
0
0
0
2
1/2
1/3
1/4
1/5
1/6
0
0
0
0
3
1/3
1/4
1/5
1/6
0
0
0
0
0
4
O
0
O
1/4
1/5
1/6
0
0
0
0
0
5
O
O
1/5
1/6
0
0
0
0
0
6
O
O
O
O
6
0
0
0
7
0
D
0
1
1000
8
0
0
2000
9
0
0
000000000
10
0
-0000
0
00000000000
11
0
12
0
14
0
Table 58: Example of U (2) for readability)
(1), continuing from Table 49. (Irrelevant dimensions are omitted
I
$
7
5
9
5
X
7
9
0
0
0
1: NUM
0
0
0
0
0
0
0
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
3: IS [S_BOS
0
0
0
0
0
0
0
0
0
0
0
0
4: IS MSLM
0
0
0
0
0
0
0
0
0
0
0
0
5: IS_ IS_EQUAL
0
0
0
0
0
0
0
0
0
9: OP2_TEE TEN
0
5
0
9
0
0
0
Table 59: Example of U. (2() (1)
A
continuing from Tables 57 and 58. (Irrelevant dimensions
are omitted for readability)
$
5
9
X
9
=
0
0
1: NUM
0
0
0
0
0
0
0
0
2: FULL_ ONES
0
0
0
0
0
0
0
0
3: IS [SO
0
0
0
0
0
0
0
0
0
0
0
4: IS_MUL
0
0
0
0
0
0
0
0
0
5: IS EQUAL
0
0
0
0
0
0
0
0
0
0
0
0
9: OP2. TEN
0
2/2
2114
26/5
13/5
7
7
G.7.3 Attention Head 3: Copying the Appropriate Digit from the First Operand
The objectives of the first and the second Attention heads were to extract the ones and tens digits of the second operand and display them in the dimensions OP2_ONE _ONEE and OP2_TER respectively. For Attention head we mainly focus on the first operand. Specifically, in Attention head 3, the goal is to fill the dimension uPI_SHIFTO at the i-th least significant digit of the response (when predicting
62

---

the (2 + 1)-th least significant digit of the response) with the (i + 1)-th least significant digit of the first operand. For our example, we want to fill OP ____SIFFT of the token = by 5. Here, ranges from O to la + 2, where the (-th least significant digit of the response denotes the equal token. In cases where 2 fa. we fill by 0.
Additionally, the third head has an extra objective: filling the dimension PRE_EOSI. This dimension utilized for EOS prediction in the subsequent FFN layer along with PRE_EOS2, which filled by the fifth head of the same layer. We observed that both objectives can be achieved by utilizing the same attention map. Thus, instead of implementing these objectives in separate heads, we can achieve them by utilizing the matrices V and U described below. Unlike previous heads, Vy²) and _ each have two elements, with each element contributing of the objectives.
Our specific construction as follows. With dQ = =
0!²²
(MMP
3x
OPxP
VMIP
OpxP
Opxp
OpxP
OpxP
Rddkkaxd
(e
(FULLLLOEES ONES
O 01xP
01xp
01xP
01xP
01xP
01xP
(106)
K2 (2)
Opx34 VMP (cis_Bos)
VMIP
PxxP
OpxP
0pxP
0pxP
0pxP
E I
Rddo
(4
O1xP
01xP
01xP
01xP
01xP
0 01xP
(107)
and with dv.23 = 2,
Vg (((eN))) TT Rdv.asxd
(eess )T
(108)
U (2) (CoP1
ePRE_EOSI)
aaddv,2
(109)
We provide the examples Tables 60 to 64. We note that within the dimension PRE_EOS] of the matrix U (22)) x A(") if we restrict our view to the equal symbol and the response sequence, 1 is only assigned to the first, second, and third most significant digits of the response (regardless of the query length).
Table 60: Example of Q, (1), continuing from Table 49.
I
7
5
9
5
9
5 0 o 0
I-P:
VMug VMol varu vMug vuruf vMug Varug vMuf vuug vMug vwanf vMug vuruf /Mof P+ P+1:MMB VMP VMP VMP VMP VMP VMP VMP VMP VMP /MP VMP VMP VMP VMP
Table 61: Example of K! X (1),
(2)
continuing from Table 49.
I
7
5
X 7
9
5
0
I-P:
V V C V
VM
0.
P+ VMF
63

---

Table 62: Example of A
(2)
(with explicit row/column indices and sufficiently large M), continuing
from Tables 60 and 61.
row col 2
4
6
8
9
10
11
12
13
14
15
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1
1
2
1/2
0
0
0
0
1/2
0
3
0
0
0
1/2
0
0
1/2
0
0
0
1/2
0
0
0
0
1/2
0
1/2
0
1/2
0
0
5
0
0
0
0
1/2
0
1/2
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
O
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
000002000
-0000
0
11
10
O
0
0
0
0
00000
0
0
0000
00000000000
0
00000
00000
0
0
-0000
0
0
0
0
00000
00000
0
00000
000002000
00000
00000
13
12
0
0
0
0000
000000000000
0
0
000
14
0
15
0
Table 63: Example of U V (2) for readability)
(1), continuing from Table 49. (Irrelevant dimensions are omitted
I
$
7
9
7
0
0
0
1: NUM
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_BOS
0
0
0
0
0
0
0
4: IS 3_MUL
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
5: IS_ S_EQUAL
0
0
0
0
0
0
0
10: SHIPOO
0
14
10
18
10
0
14
18
0
10
0
0
0
12
21: PRE BOS!
0
0
0
0
0
0
0
Table 64: Example of U V omitted for readability)
(2)
(1)
A
continuing from Tables 62 and 63.
(Irrelevant dimensions
I
$
7
5
9
X
7
9
=
5
0
0
0
6
1: NUM
0
0
0
0
0
0
0
0
0
0
0
0
0
2: FULL_ONES ONES
0
0
0
0
0
0
0
3: IS BOS
0
0
0
0
0
0
0
0
0
0
0
0
0
0
4: IS_MULL
0
0
0
0
0
0
0
5: IS_ EQUAL
0
0
0
0
0
0
0
0
0
0
0
0
0
10: OP1 SHIFO
0
0
5
5
5
5
7
0
21: PRE, PRE_EOS|
1 1
1/2
1/2
1/2
1
G.7.4 Attention Head 4: Copying the Appropriate Digit from the First Operand II
The objective of Attention head 4 is to fill the dimension uPI_SHIFTI at the i-th least significant digit the response (when predicting the (i + 1)-th least significant digit of the response) with the i-th least significant digit of the first operand. Similarly to the previous head, i ranges from 0 to la + 2. cases where the i-th least significant digit of the first operand not well-defined (i.e.,
i {0 la + 1, lo + 2)), assign 0.
64

---

The design of Attention head 4 is as follows. With dQK221 P + 1,
0
0x434
OpxP
OpxP
VMIP
OpxP
OpxP
OpxP 01xP
Rdaazux
<MP (PULLLONUSS
0IxP
01xP
01xP
0IxP
01xP
(110)
K!²)
Px334 VMP
VMIP
OpxP
Opxr
Opxp
Opxp
OpxP
E Reoxaaxd
(eis
(e34
T
01xP
01xP
01xP
01xP
01xP
0 0IxP
(111)
and with dy.24
Vi33 2(eNuM) T Rdvaadd
(112)
U d
(113)
We provide the examples in Tables 65 to 69.
Table 65: Example of Q (1), continuing from Table 49.
I
7
5
5
VM vMleg valog VMle Vilug VITe viluf vile vho VHof viluf VMlel V.t V.Tef P /MP VMP VMP JMP VMP VMP VMP JMP JMP VMP VMP JMP VMP VMP VMP
I-P:
Op
Table 66: Example of K! (1), continuing from Table 49.
7
I-P: OF VMvl VMug VMog VM P+1:/MP 0
0p 0p 0p
0p 00
02
0.
0.
0.
0p
0
0 0
Table 67: Example of (2)
(with explicit row/column indices and sufficiently large M), continuing
from Tables 65 and 66.
row col j
4
6
9
10
11
12
13
14
15
1/2
1/2
1/2
1/2
1
1/2
1/2
1/2
1/2
1/2
1/2
1
2
0
1/2
0
0
0
0
0
0
0
0
0
1/2
0
1/2
0
0
0
0
1/2
0
4
0
0
0
1/2
0
1/2
0
0
1/2
0
0
0
1/2
0
1/2
1/2
6
O
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
10
O
0
0
0
0
-0000
-
0
0
0
0
0
0
0
0
0
-
0
0
0
0
-0000
0
0
0
0
0
00000
0
0
0
00000
-000-
0
0
00000
0
0
0
0
0
12
13
0
0
0
0
0
0
0000
0000
00000
14
15
65

---

Table 68: Example of 313) V422 x for readability)
(2)
(11,, continuing from Table 49, (Irrelevant dimensions are omitted
$
7
5
9
X
7
I: NUM
0
0
0
0
0
0
0
0
0
2: TULL_ONES
0
0
0
0
0
0
0
0
0
3: IS_ S_MUL
0
0
0
0
0
0
0
0
0
0
0
0
4: IS_ EEUALL
0
0
0
0
0
0
0
0
0
11: OPI SHIFTI
0
14
10
18
10
0
14
18
0
10
0
0
0
12
Table 69: Example of U () are omitted for readability)
V1
2)
(2)
A
())
continuing from Tables 67 and 68. (Irrelevant dimensions
I
$
7
5
9
X
7
5
0
0
0
1: NUM
0
0
0
0
0
0
0
0
0
2: FULIL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_MUL.
0
0
0
0
0
0
4: S_EQUAL
0
0
0
0
0
0
0
0
0
0
0
0
0
0
1I: II_IIIIII
0 7
9
0
9 5
0
G.7.5
Attention Head 5: Copying the Appropriate Digit from the First Operand III
The main objective Attention head is to fill the dimension OP1 1P1_SHIF22 at the i-th least significant digit of the response (when predicting the (i 1)-th least significant digit of the response) with the
(i fi-1 1)-th least significant digit of the first operand. Similarly to the previous head, i ranges from 0 to la and in cases where the i-th least significant digit of the first operand is well-defined (i.e.,
i E {0, e, + 2}), we assign 0.
As mentioned in Attention head 3, assign extra goal to Attention head 5, which is to fill the dimension REs2.
The design of the fifth head is as follows. With dQx = P+I.
0p Oxx34
OPxP
OpxP
0pxP
VMIp
OpxP
OpxP
E Rdokaaxd
VMP
(c34 (CELLL_ONES))
0 0IXP
01XP
0IXP
01xP
0 01XP
0 01xp/
(114)
Kg (2)
Px34 VMP
VMIP
OpxP
0PxP
0pxp
0pxp
0pxp
E Rdok
Bei)_
(e34
01xP
01xP
0 OIXP
O1XP
01xP
0 01xP
(115)
and with dy.25 2,
Vi(0) (2 ) (ell ROS) )
E dv
(116)
U (edr1 [0mE2
ed
eoR2) E dxddv2s
(117)
We provide the examples in Tables 70 to 74. Note that within the dimension PRE_FOS2 of the matrix Ug V. (1() ( if we restrict our view to the equal symbol and the response sequence, is only assigned to the most and the least significant digit of the response, and the equal token. An important observation is that upon comparing PRE_EOS! EOST and PRE_EOS2, the most significant digit of the response is the only token that has a value of in both dimensions. This observation plays crucial role predicting EOS for the next token, and we will elaborate further the later section discussing the FFN layer.
66

---

Table 70: Example of Q (2)
_ continuing from Table 49.
I
7
5
9
5
5
0
6
I-P:
vMug Vivg Vo VM Ming Mf Vi VM VMug V V vMlg /M vM
P P+1VMB VMP VMP WMP VMP VMP VMP VMP VMP
VMP
VMP
vMP
VMP
VMP
VMP
Table 71: Example of K 5
(2)
x(),, continuing from Table 49,
I
$
7
5
7
9
5
0
I-P: VM V VM /
0g 0p
0p
P+1:/MP
Table 72: Example of A (with explicit row/column indices and sufficiently large M), continuing
(2)
from Tables 70 and 71.
row cool 2
3
4
6
7
8
9
10
11
12
13
14
15
1
1
1
1/2
1
1
1/2
1/2
1/2
1/2
0
0
0
0
1/2
0
0
0
0
0
0
0
0
0
0
0
1/2
0
0
0
0
0
1/2
0
5
0
0
0
0
0
1/2
0
0
0
1/2
0
0
0
0
0
O
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
10
9
O
0
0000020
0
0
0
00000
0
0
00000
0
0
0000000000
0
0
000002000
00000
-0000
0
0
00000
0
0
11
0
0
12
0
0
0
0
0
13
14
15
Table 73: Example of ( v.(2) (1), continuing from Table 49, (Irrelevant dimensions are omitted for readability)
$
9
0
0
0
I: NUM
0
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_MUL
0
0
0
0
0
4: IS EQUAL
0
0
0
0
0
0
0
0
0
0
0
0
0
12: 1 SHIFT2
0
14
10
18
10
0
14
18
0
10
0
0
12
22: PRE BOS2
1
0
0
0
0
0
0
0
0
0
0
0
67

---

Table 74: Example of U are omitted for readability)
2
(2)
X
A
(2)
continuing from Tables 72 and 73. (Irrelevant dimensions
$
9
X
0
I: NUM
0
0
0
0
0
0
0
0
2: ONES
0
0
0
0
0
0
0
0
3: IS_ IS_MUL
0
0
0
0
0
0
0
0
0
4: Si_ EQUAL
0
0
0
0
0
0
0
0
0
12: OPI SHIFT2
0 0
0
0
0
5
0
0
5
9
7
22: PRE BOS2
1
1/2
1
1/2
1/2
1/2
G.7.6 Attention Head 6: Copying the Appropriate Digit from the First Operand IV
The objective Attention head 6 to fill the dimension OPI_SIIFT3 at the i-th least significant digit of the response (when predicting the (i 1)-th least significant digit of the response) with the (i (--3 2)-th least sigmificant digit of the first operand. Similarly to the previous head, i ranges from to la 2. In cases where the i-th least significant digit of the first operand is not well-defined (i.e.,
i E {0, 1,2)), we assign
The design of Attention head 6 is as follows. With dQK P +1,
0px31 24 (EELLLLOEES)
OpxP
0pxP
0pxp
0pxP
VMIP
0pxP
RRdokaexd
VMP
OpxP
01xP
01xP
01xP
01xP
01xP
(118)
K.
Px33 VMP (e Ros)
VMIP
OpxP
0pxP
0pxP
0pxP
OpxP
E RRdxx
01xP
01xP
01xP
01xP
01xP
01xP
(119)
With dv.26
Vi20 2(eN)M) E Rddvaaxd
(120)
( ed
CoPI_SHIPT 2d
(121)
We provide the examples in Tables 75 to 79.
Table 75: Example of Q (1), continuing from Table 49.
I
7
5 9
X 7
9
0 0
0
I-P: Or VM vMof VMvl VMof VMuu vMug vTTef Vilufa vMof ViTuu vMof VT vMlo VM P+1.ve VMP VMP VMP VMP var VMP VMP VMP VMP VMT
VMP VMP VHP VMI VMP
Table 76: Example of K! (1), continuing from Table 49.
I
7
5
X
7
9
5
0
0
I-P:
V V E V VM
0.
0.
0.
0.
0
P+ VMF
68

---

Table 77: Example of A
(2)
(with explicit row/column indices and sufficiently large M), continuing
from Tables 75 and 76.
row col 2
4
6
8
9
10
11
12
13
14
15
1
1
1
1
1
1/2
1/2
1/2
1/2
0
0
0
0
0
1/2
3
0
0
0
0
0
0
0
0
0
0
0
0
1/2
0
0
0
0
0
1/2
0
5
0
0
0
0
0
0
0
0
0
0
1/2
0
0
0
0
0
0
0
0
0
0
0
0
O
0
0
0
0
0
0
0
0
0
-0000
0
0
000002000
00000
0
11
10
O
0
0
0000
00
0
00000
00000
00000
00000
00000
0
00000
0000
0
0
0
00000
12
13
14
0
15
Table 78: Example of U (2) x for readability)
(1),
continuing from Table 49. (Irrelevant dimensions omitted
I
$
9
X
0
1: NUM
0
0
0
0
0
0
0
2: FULL_ ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_ IS_MUL
0
0
0
0
0
0
0
0
0
4: IS_ EQUAL
0
0
0
0
0
0
0
0
0
13: UPI_SHIFTS
0
14
10
18
10
0
14
18
0
10
0
0
12
Table 79: Example of U V are omitted for readability)
(2)
(2)
X A
continuing from Tables 77 and 78. (Irrelevant dimensions
I
$
9
0
0
0
1: NUM
0
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
3: IS_MUL
0
0
0
0
0
0
4: EQUAL
0
0
0
0
0
0
0
0
0
0
0
0
0
13: OPI SHIFT3
0
0
0
5
5
G.7.7 Attention Head 7: Copying the Appropriate Digit from the First Operand V
The objective of Attention head 7 is to fill the dimension OP1
OPI__HIFTH SHIFT4 at the i-th least significant
digit the response (when predicting the (i 1)-th least significant digit of the response) with the (i 3)-th least significant digit of the first operand. Similarly to the previous head, i ranges from 0 f. + 2. cases where the i-th least significant digit of the first operand is not well-defined (i.e.,
i {0, 1, 3}), we assign 0.
69

---

The design of Attention head 7 is as follows. With dQK227 P + 1,
Q
Px33 34 VMP (FULLLOOESS)
OpxP
0pxP
PxxP
0pxP
0pxP
VMIp
Rdaaaaxd
01xP
O1xP
01xP
01xP
01xP
01xP
(122)
K2 (2)
O Ox331
VMIP
OpxP
OpxP
0pxP
0pxP
OpxP
E IRdoRaax
VMP (eis_oos)
01xP
01xP
01xP
01xP
01xP
01xP
(123)
With dv,27 =
V,80 2(eNUM) T Rddvardd
(124)
U ed S Rdxddvar,
(125)
We provide the examples in Tables 80 to 84.
Table 80: Example Q (1), continuing from Table 49.
I
7
5
9
5
6
I-P:
VMuu VMwg VMug VMol vMof VMug VM vMuli vMul Vinug vMog vaany vMw VMug P+1:VMP VMP VMP VMP VMP VMP VMP VMP VMP VMP VHP VMP VMP VMP VWP
P
Table 81: Example of K (2)
X continuing from Table 49.
I
7
5
X
7
9
5
0
0
I-P:
Vf vM oM V 37
0p 0p
0p
0p
P+1/MM
Table 82: Example of (2)
(with explicit row/column indices and sufficiently large M), continuing
from Tables 80 and 81.
row col
4
6
8
9
10
11
12
13
14
15
1
1
1
1/2
1/2
1/2
2
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
1/2
4
0
0
0
0
0
0
0
0
0
0
1/2
0
0
0
1/2
6
0
0
0
0
0
0
0
0
0
0
0
0
-
0
0
0
O
0
0
0
0
00000
0
0
0
0
0
00
0
00000
0
0
0
0000000000
00000
00000
0
0
00000
10
0
00000
0
12
0
0
13
0
0
0
14
15
70

---

Table 83: Example of U220) V. for readability)
(2)
x (1), continuing from Table 49, (Irrelevant dimensions are omitted (11,,
$
9
X
I: NUM
0
0
0
0
0
0
0
2: FULL_ONES
0
0
0
0
0
0
0
0
3: IS_ MMLUL
0
0
0
0
0
0
0
0
0
4: IS_ EQUAL
0
0
0
0
0
0
0
14: OP1 SHIFT4
0
14
10
18
10
0
14
18
0
10
0
0
0
12
Table 84: Example of U V are omitted for readability)
(2)
(2)
(1)
A continuing from Tables 82 and 83, (Irrelevant dimensions
$
9
X
0
1: NUM
0
0
0
0
0
0
0
2: FULL_ONES ONES
0
0
0
0
0
0
0
0
3: IS_MUL
0
0
0
0
0
0
0
0
0
4: IS_ B__QUULL
0
0
0
0
0
0
0
0
0
14: OP1 SHIFT4
0
0
0
0
0
0
0
0
9
5
G.7.8 Residual Connection
So we have computed the output of Attz operation. Passing through the residual connection, the output of the attention layer becomes the sum of (1) (the input the second Transformer block) and the output of Attz operation:
( X(), Σ (2) (2) (1) (2)
(126)
hE/7
A concrete example of the output of residual connection is presented in Table 85.
71

---

Table 85: Example output of residual connection, continuing from Tables 49, 54, 59, 64, 69, 74, 79 and 84, Uncolored rows represent the initial embedding. Gray rows indicate the rows filled by the first Transformer block. Yellow rows indicate the rows filled by the attention layers at the second Transformer block. Pink rows indicate the rows that will be filled by the subsequent FFN layer.
I
S
7
5
9
X 7
9
=
5
0
6
I: NUM
7
5
9
9
0
5
0
6
2: FULL_
1
1
1
1
1
1
1
1
1
1
1
3: IS_BOS
1
0
0
0
0
0
0
0
0
0
0
0
0
0
0
4: IS_MUL
0
0
0
0
0
0
0
0
0
0
0
0
5: IS_EQUAL
0
0
0
1
0
6: IS 1S_OP2_ _02__ONE ONE
0
0
1
0
0
0
7: 18 S__OP2_TTR TEN
O
0
0
0
0
1
0
0
0
0
0
0
8: OP2_oP000 OONE
0
7/2
4
2014
13/3
33/7
9
9
9
9
9
9
-
9: OP2 TEN
0
1/2
4
7
7
7
7
7
7
7
7
10: OP!. SHIFTO
7
9
5
0
0
0 00
11: OP1 SHIFT]
0
5
9
0
9
5
0
5
0
12: OP1_SHIFT2
0
O
0
0
5
0
13: OP1_SHIFT3
O
O
O
0
0
0
0
0
5
7
14: OP1_SHIFTT
0
0
0
9
5
15: RESULTI
0
0
0
0
0
0
0
16: RESULT2
0
0
0
0000000000
0
0
17: RESULT3
0
00
0
0000
0000
000
0
0000
18: RESULI4
0
0
0
0
PRE_PROD
20: PRE_CARRY
0
0
0
21: PRE_FOST
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1/2
1/2
22: PRE_FOS2
1
1/2
1
1
1/2
1/2
1/2
23-32: PROD
O1a
010
010
010
O10
010
O10
O10
O10
01d
010
OLD
O1a
O10
OLD
IS EOS
34: MASK
0
0
0
1
1
35-(P+34) POS 2 MASK
0p
0p
0P
0P
OP
0
0
0
0
0
0
(P435)-(PP+34) POS
Op
vg
vi
v2
of
(2P+35) pOs 2
0p U
07
(3P+35) POs
0p U
Ug
Up
0g
vi
ug
((P+35)66P433); POS_4
Op
vip
vip
m
(6P435)6)443) 5
Op
1gg
viu
Vi1
Vio
vi
via
ug
vi
G.8 Transformer Block 2 Token-wise Feed-forward Layer
Our ultimate goal is to fill the dimensions PROD and IS_ EOS with appropriate values. The dimensions RESULTI to RESULT4, PRE_ _ERD,, and PRE_ CARRY serve as temporary memories for storing inter- mediate values, which will help us achieve our ultimate goal. Our construction involves sequentially stacking the MLP networks step-by-step to generate each of these temporary values. (As mentioned in the theorem statement below, we allow FF, to be a multi-layer MLP.)
While our current construction for FF2 involves multiple hidden layers, we believe that our construc- tion can be improved to employ a single hidden layer. If employing multiple hidden layers in the FFN not feasible, this issue can be addressed by introducing additional Transformer blocks. Specifically, we can bypass the attention layer in these extra blocks by residual connection and only utilize their FFNs.
Step 1. Filling RESULT 1 to RESULT_ 4 Here, we first assume the existence of a single-hidden. layer MLP network, denoted f R - R. such that given any integers a, b E {0, ........90. 9), f(a, b) equals to their multiplication, i.e., ab. Such network can be implemented with 100 hidden nodes (Zhang al., 2021).
Recalling Appendix G.4, we construct the first MLP network by utilizing eight instances of the function f in parallel as follows:
1. RESULTI ([OP1_SHIFTI op2_ONE) + f(ool_ SHIFTI, oP2_ oP2_TEN)= (0, 1, ..., 162).
72

---

2. RESULT2 =
+ [(OPI_SHIFT2, OP2_TEN] E {0, 1, 162).
3. RESULT3 f(OP1_ SOOP__SHIFT2,PP ONE) + SSI___SSIII,,2111111 {0, 1, 162).
4. RESULT4 [(OPI_SHIPT3, ) + SFPP__SSHH4,,,_11111 {0, (0.1..... 162).
Step Filling PROD and FRE_CARRY
Here, we assume the existence of the following three single-hidden-late -ayer MLP networks, denoted 91 92.93 93 - IR such that given any most 3-digit integer {0, 1..... 162), 91(),, 92 92(a) and 9a (a) output the ones, tens, and hundreds digit of a, respectively. Similarly to the previous step, each network can be implemented with 163 hidden nodes (Zhang al., 2021).
Recalling Appendix G.4, we construct the second MLP network on top of the first MLP network, by utilizing 2 instances of each of the function 91, 92, and gs in parallel follows:
PRE_PROD 9i (RESULTI) + 92 (RESULT))) + g3 (RESULT3) E (0, 27}.
PRE_CARRY 91 (RESULT2) + ga (RESULT3) + (RESULL)))) E 10.
Step 3. Filling PROD Here, we assume the existence of a single-hidden-lag MLP network, denoted h IR2 R. such that given any integers {0, (0... 27). b {0, ......9.) satisfying b b E 11,001,0,8,0 10, 18, 19, 20J, h satisfies
0, i - be (-2, -1, 0},
1. a E {8, 9, 10). h(,, b)
2, a E (18, 19, 20).
We also assume the existence of a single- hidden-layer MLP network, denoted as h' R - R, such that given any integer (0, 1, .., 19), h'(a) equals (mod 10).
We finally assume the existence of single-hidden-layer MLP network qi R RR for each
i (0.1. such that given any integers a E (0, (........9) 9}, satisfies
gi(a)
(i a).
Similarly the previous step, each network can be implemented with 280, 20, and 10 hidden nodes. Recalling Appendix G.4, we construct the third MLP network, on top of the second MLP network, by
9o (PRE PROD + [PRE_PRODD________ 2 (PRE_PROO + (PRR__CARRY, NUM)))
h(RRE_ [PRE_CARRR,, NUM)))
PROD
E IR100
29
+ ((PRE_CARR) NUM))) (ARRY)
One can easily check that (PRE BLREE__RRDD_- + h(PRE_ CARRY, NUM)) yields an element of 0, 1.....9. 9. and thus PROD is one-hot column vector. Specifically, if (wF.pppP__RP CARRR/REEEAARREEE_E NUM)) i, then PROD becomes e
Step 4. Filling Is EOS We construct single- MLP network r: IR2 TR. by
r(a, b) 20(a b.5).
We then can fill the dimension i$_BOS by
EOS (PPR__EOS1, PRE [RR_))
Since PRE_ E 1oslamm and PRE, PRE_EOS2 have either 1/2 1, si__so_________ equals only when both PRE 1 and PRE_EOS are 1. Additionally, we note that PRE_ EOS and PRE_ are the direct outputs from the attention layer. Therefore, the network network and does not require an additional FFN layer.
be deployed in parallel with the first MI.
The example output resulting from passing through all these steps is presented in Table 86.
73

---

Table 86: Example output of FFN layer in the second Transformer block, continuing from Table 85. Here, we mark for the entries before the equal token, as these entries do not affect the next-token prediction our construction and are thus not important.
I
S
7
5
5
5
0
0
6
15: RESULTI
45
116
108
98
49
0
16: RESULT2
0
45
116
108
98
49
0
RESULT3
0
116
108
49
18: RESULT4
0
0
0
45
116
108
98
PRE_PROD
5
9
0
20: RRE_ARRRY
0
5
10
9
9
19
4
23-32: PROD
C6
10
6
e
10
e
el
33: IS_BOS
0
G.8.1 Residual Connection
The last task of the feed-forward layer is pass FF2 (Y through the residual connection. As result, we have
x (2) FFF2 (1(2)).
(127)
This is the end of the second Transformer block, and an example of X is illustrated in Table 87.
G.9 Decoding Function
As mentioned in Appendix D, the decoding function performs a linear readout (with weight matrix Wout E IR M/xd) and a (token-wise) arg-max operation. That is,
Dec (x()) = (Vk.) E V.N.
(128)
where V4 is the k-th element of V and
k; arg max for Wout x(1)= - [o
qiee 7
(129)
kE[|v]]
The objective of the decoding function is to perform proper next-token prediction for N x 2 multiplication, especially utilizing the dimensions PROD and _S_FOS of (2)
We now construct the weight matrix Wout- For a token d;, if the value of dimension EOS of (2) is O. then the linear readout output the dimensions PROD as it is to return one of number token (0-9). On the other hand, the value of dimension IS_FOS 1, then large number (like 9 for example) for the token '8' to return EOS (8). This can be implemented by the weight matrix W Wont described Table 88, Also, an example of applying the linear transform showcased in Tables 89 and 90.
linear readout outputs
74

---

Table 87: Example embedding matrix after the second Transformer block. The yellow rows represent the results introduced during the second block, while the gray rows indicate the results from the first block. Similarly to Table 85, we mark for the entries before the equal token, as these entries do not affect the next-token prediction in our construction and are thus not important.
$ 7
9
5
X
7
9
0
6
NUM
0
9
5
9
0
2: FULL_
1
1
1
1
1
1
1
1
1
3: IS_BOS
1
0
0
0
0
0
0
0
0
0
0
0
0
0
4: IS [SMMUL
0
0
0
0
0
0
0
0
0
0
0
0
[S_EQUAL
0
0
0
0
6: IS_OP2_ OP2_ONE ONE
0
0
7: 2S_OP2 OP2__EN TEN
0
0
0
0
0
0
0
0
0
0
0
8: OP2_ONE
9
9
9
9
9
9
9: OP2 TEN
7
7
7
10: OP SHIFTO
5
7
11: OP! SHIFT
9
9
5
7
3
0
0
12: OP1 OPI_SHIFT2
O
5
7
0
5
13: OPI_SHIFT3
O
5
7
14: OP OPI_SHIFT4
0
15: RESULTI
45
108
98
49
16: RESULT2
O
45
116
108
98
49
0
17: RESULT3
O
45
116
108
98
49
18: RESULT4
O
45
116
108
98
19: PRE_PROD
9
20: FRE_CARRY
10
9
9
19
4
21: PRE_FOSI
1/2
1/2
1/2
1/2
22: PRE_FOSE FOS2
1/2
1/2
1/2
23-32: PROD
10
61
E1
61
e7
33: IS EOS
0
0
34: MASK
0
0
0
1
35-(P+34); POS 2 MASK
0p 0p
Op
0p
0p
0p
0p
0g
0p
0p
0P
(P+35) POS_
0y 0i
0
of
U6
o
0
U2
0
POS 2
0,
0g
(8P435)-(44+ (4P+34)) POS
0y ve
0g
0
(1P435)-6PP434; POS
0p vi
thg
p
whg
(6P4-5)-66433))) POS
0p
0g
0g
via
07
Table 88: The transposed weight matrix WT out of the linear readout in decoding function. p represents 6P + 1.
V
0
2
4
5
6
7
9
S
1-22: NUM-PRE EOS 2
Oz
O22
O22
O22
Og2
Ouz
02
O22
Og2
O
0
0 22
O22
23: PROD1
0
0
0
24: PROD2
0
0
0
0
0
0
0
0
0
0
0
25: PROD3
0
0
0
0
0
26: PROD4
0
0
0
0
0
0
0
0
0
0
0
27: PRODS
0
0
0
0
0
0
28: PRODG
0
0
0
0
0
0
0
0
0
0
0
29: PROD7
0
0
0
0
0
0
0
30: PRODs
0
0
0
0
0
0
0
0
0
0
0
31: PROD
0
0
0
0
0
0
32: PROD10
0
0
0
0
0
0
0
0
33: IS EOS
0
0
0
0
0
0
100
e4d
pr
0
P
P
0
P
P
0 pr
0
0Pr
0
Pr
0
0P
0
PP
0
P
Pr
0
75

---

Table 89: Example output of linear readout ( out (2),, continuing from Tables 87 and 88. The yellow cells represent the maximum value of each column, from the '=' token's column to the rightmost column (which are used for next-token prediction).
$
7
5
9
5
X
6
0
1
0
0
2
0
0
0
0
3
0
4
0
0
0
0
0
5
0
6
0
0
7
0
8
0
0
9
X
0
0
0
$
0
100
Table 90: Example output sequence O = Dec (2),, continuing from Table 89, The yellow cells in the bottom row exactly predict the next tokens.
I
S
5
0
0
0
0
$
76

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We summarize our contributions in Section 1.1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation of our work in conclusion (Section 7).
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
77

---

Justification:
We provide a detailed construction to prove Theorem 5.1 in Appendix E. We prove Proposition 5.2 in Appendix F. LLsslyy we provide a constructive proof Theorem 6. in Appendix G.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide the training method in Sections 3 and 4 and comprehensive details on the hyperparameters used for our experiments in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
78

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We open our codebase: github. cow/HanseulJo/ wttp:///oppsooiooooppooee ing. contains fully reproducible Python codes for data generation, model training, and model evaluation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
//nips. sttpss.cc/uublccguuiceee CCodeSummideeebia ssiiiioicyy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We organize the experimental details in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We provide 95% confidence intervals in result plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
79

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide our computation resources and exccution time Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: httt::::/errrps...ppplpc.l.co http///llcc/thiic///eeeeeee csGuidel ines?
Answer: [Yes]
Justification: We are sufficiently aware of the NeurIPS Code of Ethics and have ensured that our research complies with all its principles.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper focuses exclusively on arilhnetic/algorithmie tasks, which do not have direct societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
80

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Our paper focuses exclusively arithmetic/algorithmic tasks, which inher- ently do not pose high risks for misuse.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We clarify the reference of our code base in Appendix C.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
81

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwperswithcae thcode. com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not propose any new datasets or models and does not release any code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
82

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
83


Reviews:
[{"review_id": "lFeZZay3gh", "reviewer": "Reviewer_ow6J", "summary": "This work considers the problem of length generalization of Transformers and proposes injecting the task structure through positional embeddings for improving length generalization. Task structures are known and therefore, the authors come up with a (relatively) general heuristic to leverage this structure.\n\nThe paper relies on the observation that for tasks like addition, there are “groups” of tokens that should be treated similarly as they carry the same semantics. I.e., the digits of the summands from least significant to most significant should be embedded with the same positional embedding (called *position coupling*) so that the model can take advantage of this structure and carry out the sum correctly.\n\nPosition coupling is used in conjunction with a number of other tricks such as reversing the sum, zero padding, and using BOS/EOS.\n\nThe proposed method is then evaluated comprehensively on the addition task (for which it was proposed), and strong length generalization is observed. Ablations on the number of layers, and different positional embeddings are carried out to further emphasize the importance of position coupling.\n\nThe results are also backed by theory, and interestingly the attention patterns predicted by theory are observed in the experiments.\n\nOther than the addition task, Multiplication ($N\\times2$) and a 2D task are considered in the experiments as well.", "questions": "None.", "limitations": "See weaknesses.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The empirical results are quite strong and proper ablations and baselines are considered.\n- The definitions, presentation, and method are quite clear.\n- The method is backed by theory.\n- The method is relatively general, however, it is only applicable to a class of tasks where there is a clear structure to be exploited.\n- The predictions of the theory are further supported in the experiments (the attention patterns for carry detection)\n- The experiments go beyond addition to multiplication and a 2D game (where more explanation is required)", "weaknesses": "The proposed method, though general for some tasks, seems to be very much geared towards a specific class of tasks of interest in length generalization. In particular, tasks like sorting, mode, and parity seem to be automatically out of reach for positional coupling and limit the generality of its applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive review and valuable comments. Below, we address the reviewer’s concern.\n\n> **W1. Position coupling is geared towards specific tasks in length generalization.**\n\nFirst, as mentioned in the conclusion of our paper, we focus on certain tasks with a handy structure between specific positions of tokens. Since vanilla Transformers often fail to learn the true structure of the tasks—even for simple algorithmic tasks—as observed by plenty of previous works [1–10], we aim to guide the model to properly learn the structure by coupling the positions, thereby improving length generalization. The structure between positions is important for employing our method, so we admit that it is not easy to apply our method to some tasks. **However**, we address the three tasks raised by the reviewer.\n\n**Parity**. \n- Given a binary query sequence, the objective is to output 1 if there is an odd number of 1s in the query and to output 0 otherwise. \n- Without any data processing, it is a difficult task for Transformers to achieve length generalization [1,4,5,7,9]. However, with the help of a scratchpad generated by leveraging the idea of unrolling the query, Transformers can achieve length generalization [9]. To utilize the power of position coupling, we opt for a different scratchpad from that used in [9].\n- By unrolling the query, we can generate a sequence of partial results. Let’s put the partial result from the 1st to n-th token of the query into the n-th token of the response sequence. For example, if the query is given as 010011, the response with scratchpad would be 011101. Then, the last token of the response immediately becomes the answer for the original task. We train the model to generate the whole response (including scratchpad) for each query to solve the task step-by-step with next-token prediction. Now, we can naturally couple the n-th positions of the query and the response when we apply our method. For example, given an input sequence “010011=011101”, we can assign (3,4,5,6,7,8,2,3,4,5,6,7,8) as position IDs (when the starting ID is randomly chosen as 3).\n- To showcase the efficacy of position coupling on the parity task with a proper input format, we compare 4 different settings: position coupling with/without scratchpad and NoPE with/without scratchpad. For the “position coupling without scratchpad” setting, we naively couple all the positions (except for the ‘=’ token) with the same position ID (e.g., “010011=1” can get (3,3,3,3,3,3,4,3)). We train the models on the queries of lengths 1–20 and test the lengths up to 100. We measure exact-match accuracy (including scratchpad if applicable) as well as the accuracy only for the single token at the position of the last token of the response except for the EOS token (called “parity accuracy”). The result of the experiments is shown in the PDF file attached in our General Response. Without scratchpads, both NoPE and position coupling cannot even achieve good in-distribution performances. Even with our scratchpads, without any position embeddings (NoPE), a 6-layer 8-head model showcases a very restricted length generalization capability up to the length of ~30. The model performs worse than random from the length of 40 because the model sometimes outputs tokens other than 0 or 1. Most importantly, our 1-layer 4-head model with position coupling and scratchpad achieves perfect length generalization up to length 100! We strongly believe that thanks to the combination of coupled position IDs and scratchpad enabling Transformers to learn a simple algorithm for solving the task with next-token prediction, we could achieve another remarkable length generalization result.\n- We will add this result to our final manuscript with more ablations.\n\n**Sorting** and **Mode**. \n- The sorting task aims to generate a response equal to the sorted sequence of the given query; the mode task aims to find the most frequent token appearing in the given query.\n- In both tasks, there is no exploitable structure between the positions of tokens. Thus, it is not straightforward to apply position coupling to solve these tasks, thereby we did not test our method on these tasks.\n- Not only is it unnatural to couple the positions, but it is also unnecessary to do so. Vanilla Transformers already length-generalize well on these tasks [9].\n\nIn short, position coupling is an effective method if we can create a clear structure between positions with proper usage of input format; however, it is inapplicable if there is no such structure. Nonetheless, there are a lot of real-world tasks whose underlying structure between positions is vague or unavailable. We leave the research direction of extending our idea to such tasks by automatically discovering appropriate couplings of the positions as interesting future work.\n\n Please let us know without hesitation if you have further questions or comments.\n\n--- \n**References**\n\n[1] Bhattamishra et al., On the ability and limitations of transformers to recognize formal languages. arXiv preprint, 2020.  \n[2] Kim et al., Have you seen that number? investigating extrapolation in question answering models. NeurIPS 2021.  \n[3] Nye et al., Show your work: Scratchpads for intermediate computation with language models. arXiv preprint, 2021.  \n[4] Chiang and Cholak., Overcoming a theoretical limitation of self-attention, arXiv preprint, 2022.  \n[5] Delétang et al., Neural networks and the Chomsky hierarchy, ICLR 2023.  \n[6] Kazemnejad et al.. The impact of positional encoding on length generalization in transformers. NeurIPS 2023.  \n[7] Ruoss et al., Randomized positional encodings boost length generalization of transformers, ACL 2023.  \n[8] Lee et al., Teaching arithmetic to small transformers. ICLR 2024.  \n[9] Zhou et al., What algorithms can transformers learn? a study in length generalization. ICLR 2024.  \n[10] Zhou et al., Transformers can achieve length generalization but not robustly. arXiv preprint, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ght6EDsyrt", "reviewer": "Reviewer_CVfs", "summary": "This paper proposes \"position coupling\", a novel technique to improve the length generalization ability of decoder-only Transformers. Unlike standard positional embeddings, position coupling assigns the same position ID to semantically related tokens across the input sequence, directly embedding task structure within the model. This approach achieves near-perfect length generalization on integer addition, extrapolating from training on up to 30-digit sums to successfully solving 200-digit additions. The authors theoretically prove the capability of a 1-layer Transformer with coupled positions to perform additions with exponentially long operands. They further demonstrate the effectiveness of position coupling on tasks like addition with multiple summands, N×2 multiplication, and copy/reverse operations, providing a theoretical construction for the multiplication task as well. The paper explores the application of position coupling in 2D tasks, showcasing its potential beyond 1D sequences.", "questions": "- The paper acknowledges the reliance on a specific input format for optimal performance. Could you elaborate on the sensitivity of position coupling to variations in input format? Have you experimented with alternative formats and if so, what were the outcomes?\n- While the paper explores several arithmetic and algorithmic tasks, it would be beneficial to understand the limitations of position coupling. Are there specific task characteristics or structures that might render position coupling less effective or even inapplicable?\n- The design of the position coupling scheme relies on an intuitive understanding of the task structure. Could you formalize the notion of task structure and provide guidelines for designing appropriate position coupling schemes for different tasks?\n- For the 2D task, why does using the same embedding layer for both position coupling modules perform better than separate layers? Is this specific to the task or a general observation?\n- Would it be possible to extend the evaluation of position coupling to more complex 2D tasks, such as image-related tasks or tasks involving graphs or other non-sequential structures?", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- While inspired by index hinting, position coupling offers a novel and more elegant solution for incorporating task structure into Transformers. It directly embeds this information within the positional encoding, eliminating the need for augmenting the input sequence and simplifying model training.\n- The paper demonstrates a high level of technical rigor. The proposed method is well-motivated and thoroughly evaluated on a variety of tasks, including both empirical analyses and theoretical constructions. The experimental design is comprehensive, with thorough comparisons against relevant baselines and ablations on various architectural choices.\n- The paper is well-written and easy to follow. The authors clearly articulate the problem, their proposed solution, and the key contributions. The use of figures and examples effectively illustrates the concepts and makes the theoretical constructions more accessible.\n- This work addresses a crucial challenge in Transformer-based learning: length generalization. The impressive results on arithmetic tasks, particularly the significant extrapolation achieved in addition, highlight the potential of position coupling for enabling Transformers to learn algorithms and generalize far beyond their training data. The theoretical analyses provide valuable insights into the mechanism of position coupling and its role in achieving length generalization. The extension to 2D tasks further broadens the applicability and impact of this work.", "weaknesses": "- While Theorem 5.1 provides a strong theoretical foundation for the capabilities of 1-layer Transformers with position coupling, the paper lacks a theoretical understanding of why deeper models might perform worse despite their greater expressivity. Further theoretical analysis on the interaction between position coupling and depth, especially on tasks like N×2 multiplication where deeper models are necessary, would significantly strengthen the work.\n- The success of position coupling heavily relies on the specific input format (reversed response, zero-padding, etc.). It remains unclear how robust the method is to variations in input format and whether it can be applied to tasks where such specific formatting is not possible or desirable. Exploring alternative position coupling schemes that are less sensitive to the input format or evaluating the method on tasks with diverse input structures would strengthen the claims of generalizability.\n- The paper primarily compares position coupling against basic positional embedding techniques. Including a broader range of recent length generalization techniques in the experimental comparison, such as those based on relative positional encodings would provide a more comprehensive understanding of the method's effectiveness and potential advantages.\n- The minesweeper generator task serves as a preliminary investigation into the potential of position coupling for multi-dimensional tasks. However, exploring the applicability and effectiveness of position coupling on a wider range of 2D or even higher-dimensional tasks, potentially with more complex structures, would further highlight the significance and generalizability of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments. We provide our response to the reviewer's concerns.\n\n> **W1. Theoretically, why do deeper models perform worse despite their expressivity?**\n\n- For a broader answer to the question, see our General Response. \n- We hypothesize that the performance degradation is due to the bad implicit bias of deep models (learning shortcuts to only achieve in-distribution generalization) when learning a simple algorithm to solve the task. We believe exploring a theoretical explanation for the bad implicit bias of large models on low-complexity tasks is a promising research direction.\n\n> **W2+Q1. How robust is position coupling to input formats? Can it be applied to tasks when specific formatting is undesirable?**\n\n- See our General Response for details on the robustness to input formats.\n- As we emphasized there, proper use of the input format is crucial for solving the task even with position coupling. Also, the model’s performance depends on the choice of input format. Thus, if we cannot apply any formatting, we should not expect significant success in solving the given task, not just in terms of length generalization.\n\n> **W3. Compare with recent techniques for length generalization.**\n\n- Thank you for your suggestion. One notable length generalization result for the addition task before us appears in [1], combining FIRE (a relative PE method) and index hinting. They achieve near-perfect length generalization up to operand length 100 by training on up to 40-digit additions. (Recall that we achieve 30-to-200 generalization with a 1-layer model.) Despite their great performance, they require doubling the input length because of index hinting. In contrast, our method doesn't require doubling the input sequence, so we believe our method is more efficient to run. \n- We also conducted experiments combining RoPE and our method and achieved an improved length generalization compared to vanilla RoPE: see Fig. 4 of PDF in our General Response.\n- We will add these comparisons in our final manuscript.\n- [1] Zhou et al., Transformers can achieve length generalization but not robustly. arXiv, 2024.\n\n> **W4. Can position coupling be applied to a wider range of 2D or higher-dim tasks?**\n\n- As demonstrated in the paper, our method significantly helps Transformers to length-generalize on the tasks with a clear structure between token positions. Extending this to various multi-dim tasks is an interesting future work.\n- A challenge in multi-dim tasks (not only for length generalization and for applying our method) is the exponential growth (in task dimension) of the number of tokens, which makes it difficult for the model to analyze queries and generate responses. Overcoming this dimensionality problem is an interesting future direction.\n\n> **Q2. Are there specific task structures that render position coupling less effective or inapplicable?**\n\n- Let us give you some examples. Our preliminary experiments weren’t very successful for some tasks including (1) addition with a varying number of operands and (2) multiplication with varying lengths of both operands. We tried couplings similar to the ones applied to simple addition and Nx2 multiplication, respectively, but they weren’t effective (although we did not invest much effort in making it work).\n  - The algorithm for solving (1) needs to attend to a varying (over problem instances) number of positions for generating tokens.\n  - The algorithm for solving (2) needs to attend to positions in varying relative distances in terms of the “coupled” position IDs.\nOn the contrary, our theoretical construction for simple addition and Nx2 tasks does not suffer from these difficulties, which is the key reason for successful length generalization; thus, we do not expect length generalization on tasks (1) and (2) without any advance in input formats.\n- Some tasks don’t have any structure between specific positions (e.g., sorting and mode), thereby we cannot directly apply position coupling. (See our response to Reviewer ow6J.)\n\n> **Q3. Can you formalize the notion of task structure and provide guidelines for designing proper coupling for different tasks?**\n\n- Defining task structure involves the relationship between the query and response, focusing on which tokens in the query influence the determination of each token in the response. Designing position coupling relies on human intuition, making it challenging to provide concrete guidelines. However, for tasks with unclear coupling structures, assigning position IDs (piece-wise) consecutively may be worth trying. It is shown to be empirically effective in the Nx2 multiplication task (which is later theoretically backed by Thm 6.1), although it's not intuitive how the coupled positions capture the task structure.\n\n> **Q4. Why does using the same embedding layer for both position coupling modules perform better than separate layers?**\n\n- As mentioned in lines 655–657, we do not have a clear explanation for why sharing the embedding layer performs better. We conjecture it is due to the row-column symmetry of the Minesweeper generator task, meaning that transposing the board still results in a valid problem. We believe that tasks where rows and columns have different semantic structures might not benefit from using the same embedding layer.\n\n> **Q5. Can we extend the evaluation of position coupling to more complex 2D tasks?**\n\n- Although our study focuses on simple arithmetic/algorithmic tasks, we are eager to explore more complex multi-dimensional tasks.\n- For image-related tasks (e.g., involving ViT), we could apply position coupling similar to what we did for the Minesweeper generator task. For graphs and other non-sequential tasks, the position coupling scheme may need refinement, but extending it to these complex tasks is promising future work.\n\nWe hope our response has adequately addressed the reviewer's concerns, and we would appreciate it if you could reconsider your assessment.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IFtGobNh6v", "reviewer": "Reviewer_ksWg", "summary": "This paper proposes a new way to bake in the positional structure of problems for transformers. The authors also analyze the potential for models with and without their proposed positional coupling to solve problems of arbitrary size. They also show empirically that their method helps a small transformer learn addition.", "questions": "1. Is my understanding of the test set construction accurate? All addition problems are of the form A + B where A and B have the same number of digits?  What happens when we test these models on things like \"20345 + 34 = \" for example?\n1. Why is only the answer reversed? In related work, reversal is often discussed, but this particular scheme requires even more structure-aware processing. I think blindly reversing all numbers that a model encounters would be more compelling than knowing to only reverse the answer. Do the methods in this paper perform better/worse with more consistent data processing?\n1. In Section 4, deeper models are shown to perform worse than the shallower models, which the authors attribute to the difficulty of training deep networks. Do the deeper models fit the training data as well as the shallow ones? It seems this way from the plot and if it is the case, then I think another explanation is needed here as the optimization does not seem to be the problem. Perhaps there is an argument to be made about over-fitting, but the rate of degradation seems to vary somewhat nicely with added depth raising more questions about what could be happening. Can the authors offer any other insights here? Maybe some understanding of what is happening at operand lengths where accuracy is ~80%, i.e. which examples are correct/wrong would lend some clarity here.\n\nI'm excited to discuss my review with the authors during the rebuttal period and remain open to improving my score with answers to my questions.", "limitations": "The authors adequately addressed all limitations in the limitations section. The significance is limited in ways I addressed above, but these are relatively minor.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: To the best of my knowledge the theoretical analyses are novel and the methods are novel (up to concurrent works).\n1. Quality: The work is detailed and thorough.\n1. Clarity: The paper is well written and clear.", "weaknesses": "1. The significance is limited.  \n  i. I think often algorithm learning papers can come across as limited in impact and I don't mean to bring this up. Rather, this paper specifically makes a portion of its contributions around 1-layer transformers where some of the claims seem particularly limited. For example, the proof of impossibility of a 1-layer decoder only transformer is interesting, but I don't think it applies to even a 2-layer model weakening the motivation for the use of this positional coupling in practical settings.  \n  ii. The authors state the limitation that this method requires a priori understanding of the problem structure. I appreciate that this is acknowledged as it feels important to me. With *enough* a priori knowledge of problem structure, one can often solve the problem without ML. This paper is a cool demonstration of learning addition from data, but it is limited in this sense.  \n  iii. The fact that operands are zero padded to match in length is also a limitation here. This is another instance of requiring problem-aware data processing. The test set is only made up of operands of the same size (clarifying questions below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We deeply appreciate the reviewer's valuable questions and rich comments, and we hope our response relevantly addresses all points raised in the review.\n\n> **W1. Some of the contributions are limited to 1-layer Transformers.**\n\n- We first highlight that our theoretical construction (Thm 5.1) is not limited to 1-layer models because it naturally extends to multi-layer cases. \n- As you pointed out, our impossibility result on NoPE may not hold for a multi-layer model. However, the main purpose of this result is to show a provable gap in the expressive power between position coupling and NoPE. Therefore, we believe that the inapplicability of the impossibility result in multi-layer models does not weaken the motivation for using position coupling.\n\n> **W2. With enough knowledge of the problem, we don’t need ML.**\n\n- We highlight that length generalization is a crucial issue in LLM. Plenty of existing research [1–4] (including us) regards arithmetic/algorithmic tasks as manageable but interesting test beds for studying length generalization because LLMs fail to length-generalize even on such simple tasks.\n- The main message from our work is that incorporating task structure when designing positional encodings can effectively improve the length generalization capabilities. We believe our findings can serve as a stepping stone for future research in this area.\n\n> **W3. Zero-padding is a problem-aware data processing.**\n\n- Note that zero-padding to match the operand lengths is prevalent in this research area ([3–5]).\n- One can implement our method without relying on zero-padding. We present experimental results using a no-padding scheme in the attached PDF file (Fig. 1) in our General Response. While the no-padding position coupling approach fails in in-distribution generalization when trained on a 1-layer model (due to the increased complexity of the algorithm that the model should learn), combined with proper reversing of the number(s), it functionally extrapolates the performance of deeper models.\n\n> **Q1. Should the operands in every testing example have the same length?**\n\n- You’re correct: we sampled the operands to be the same length while testing. To test the model (trained on zero-padding format) with the sample “20345+34=”, we could input it as “20345+00034=”.\n- To address the concern, we tested on operands sampled with different lengths. See Fig. 3 of our PDF attached to the General Response. Each axis corresponds to the lengths of each operand. The results show that the model is capable of solving tasks even when the two operands have different lengths, although zero-padding is applied to ensure consistency in the input format. We will add this result to our revision.\n\n> **Q2. Reversing the answer solely is a problem-aware data processing.**\n\n- Note that solely reversing the answer is also a common practice in this research area [3–5]. From now on, let us compare two different formattings: (a) solely reversing the answer and (b) reversing all the numbers.\n- We first empirically compare (a) and (b) while applying position coupling (see Fig. 1 of the PDF in our General Response). For 1-layer models, both formats exhibit near-perfect length extrapolation and show little difference. Conversely, for 6-layer models, a noticeable performance gap emerges. With zero-paddings, (b) performs better than (a); but if there’s no padding, (a) performs much better than (b). There seems to be no clear winner between the two.\n- The similarity between (a) and (b) for 1-layer models is expected. If we assign the position IDs based on the significance of the digits as usual, there is NO effective difference between (a) and (b) for a 1-layer model in terms of its prediction, which can be deduced from Prop 5.2. Accordingly, our Theorem 5.1 based on (a) can also be applied to (b) without any modification.\n- The difference between (a) and (b) for deeper models is also expected. In multi-layer models, the causal attention mask causes each token embedding to depend on the embeddings of preceding tokens after the first layer. Thus, unlike in the previous case, the predictions of the model may differ depending on the input format.\n\n> **Q3. The optimization might not be the reason why deeper models perform worse.**\n\n- There are two key aspects of optimization: convergence and implicit bias. Our explanation primarily concerns implicit bias. To answer your first question, deeper models do fit the training samples just as well as the shallower models. Thus, convergence is not the issue. We also think that overfitting is not the case, as the trained models only struggle with longer sequences while achieving perfect accuracy on in-distribution samples.\n- Therefore, we believe that the issue lies in the implicit bias of the models. Among the infinite number of solutions that achieve zero training loss, shallower models seem to possess a better implicit bias, which allows them to find solutions that generalize better for longer sequences.\n- Specifically, we conjecture that the outstanding performance of the 1-layer model is due to its relatively restricted expressivity and that the model has no way to fit the training samples other than learning the true algorithm. However, for deeper models, the model can still fit the training samples without necessarily learning the true algorithm due to greater expressivity.\n- For a detailed discussion on this, please refer to our General Response.\n\n---\n**References**\n\n[1] Jelassi et al. Length generalization in arithmetic transformers. arXiv preprint, 2023.  \n[2] Kazemnejad et al.. The impact of positional encoding on length generalization in transformers. NeurIPS 2023.  \n[3] Zhou et al., What algorithms can transformers learn? a study in length generalization. ICLR 2024.  \n[4] Zhou et al. Transformers can achieve length generalization but not robustly. arXiv preprint, 2024.  \n[5] Lee et al. Teaching arithmetic to small transformers. ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R8IgeSqD81", "reviewer": "Reviewer_U8WH", "summary": "This paper proposes a method called \"position coupling\" to enhance the length generalization of Transformer models, specifically targeting arithmetic tasks such as integer addition. The authors claim both empirical success and theoretical guarantees for their approach. The method involves assigning the same position IDs to semantically related tokens to better embed the task structure into the Transformer’s positional encoding.", "questions": "1. Can you provide insights on whether \"position coupling\" can be generalized to other positional encoding schemes beyond Absolute Positional Encoding? It would be beneficial to understand if and how this method can be adapted or integrated with other popular positional encodings, such as Rotary Positional Embedding (RoPE).", "limitations": "As mentioned above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. **Novelty**: The idea of coupling positional embeddings to reflect semantic relationships is interesting and novel.\n2. **Empirical Results**: The paper provides detailed empirical results showing improved performance on length generalization for integer addition tasks.\n3. **Theoretical Analysis**: The authors offer theoretical insights into why position coupling should work, which is a positive aspect of the paper.", "weaknesses": "1. **Scope of Application**: While the paper demonstrates the effectiveness of position coupling on integer addition tasks, the generalizability to broader and more complex tasks is not convincingly shown. The examples provided are limited and do not cover a wide range of real-world applications.\n2. **Complexity and Practicality**: The proposed method introduces additional complexity in the positional encoding process. This complexity may limit the practical applicability of the method, especially for larger and more diverse datasets.\n3. **Proper Citation of Related Work**: In section 3, lines 113-115, the randpos method for length extrapolation, introduced by \"Randomized Positional Encodings Boost Length Generalization of Transformers\" (ACL 2023), is not properly cited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the reviewer’s effort and constructive feedback. Below we summarize your feedback/questions and address these one by one.\n\n> **W1. Scope of Application: the generalizability of position coupling to broader and more complex tasks is not convincingly shown.**\n\n- We agree that the applicability of position coupling to broader tasks has not yet been explored in this paper. However, we want to note that enhancing the length generalization ability of models, even for the addition tasks, is considered an important problem as addressing and improving the arithmetic abilities of models can lead to a better understanding of the model’s capabilities.\n- Definitely, the next step of our work is to extend the application of position coupling to complex, real-world tasks. To do so, we plan to develop a method called “automatic coupling”. Unlike position coupling which requires manual assignment of position ID for each token, it will automatically learn the task structure and assign appropriate position IDs without human intervention. This advancement will enable broader applications beyond simple arithmetic/algorithmic tasks. We believe our findings can serve as a stepping stone towards this future development.\n\n> **W2. Complexity and Practicality: Position coupling introduces additional complexities in PE, which may limit its practical applicability.**\n\n- The introduction of position coupling does add complexity, particularly in designing task-specific couplings. However, once the design is determined, the additional overhead due to its implementation is negligible, as explained in lines 659-662.\n- Furthermore, in contrast to approaches such as [1] and [2] that employ index hinting which requires doubling the input sequence length, position coupling does not increase the sequence length. As a result, we believe that there are no significant additional complexities (e.g., time, memory) introduced during the training phase.\n- Additionally, we believe that the development of “automatic coupling”, as mentioned earlier, has the potential to fully address the complexity issue. \n- If there are remaining concerns regarding additional complexity introduced by our method, please let us know without hesitation.\n\n> **W3. Proper Citation of Related Work: “Randomized Position Encodings Boost Length Generalization of Transformers” (ACL 2023) is not properly cited.**\n\n- We agree that there are some similarities in the underlying concepts, and therefore we will add the citation and provide a more detailed comparison.\n- We note that our method (described in lines 113-115) differs from that of Ruoss et al., 2023. Our method randomly selects the start position ID and assigns a sequence of consecutive numbers. In contrast, Ruoss et al., 2023 assigns a sequence of increasing integers, which are generally not consecutive.\n\n> **Q1. Can position coupling be generalized to other PE schemes than APE?**\n\n- Thank you for an insightful question. During the rebuttal period, we realized that our proposed coupling method could be extended to Relative PEs such as RoPE and T5’s relative bias. The relative position between the query and key is determined by the difference in position IDs that were assigned by our position coupling method. Specifically for RoPE, we conducted experiments and observed that position coupling enhances the length generalization capability of RoPE (See Fig. 4 of the PDF file in our Global Response). We believe that this approach has significant potential for the research of length generalization and we will add experiments in our next revision.\n- Instead of adapting our method to other PE methods than APE, it is also possible to integrate position coupling with existing RPE methods such as standard RPEs, RoPE, or FIRE [3]. In this approach, RPE methods are used independently alongside position coupling, which may provide hope for application to general LLM models. Thus, we think this is another promising direction to further improve the applicability of our proposed methodology, so we will conduct some experiments and consider adding the results to our final manuscript. \n\nWe hope our response has adequately addressed the reviewer's concerns, and we would appreciate it if you could reconsider your assessment.\n\n---\n**References**\n\n[1] Zhou et al., What algorithms can transformers learn? a study in length generalization. ICLR 2024.  \n[2] Zhou et al., Transformers can achieve length generalization but not robustly. arXiv preprint, 2024.  \n[3] Li et al., Functional interpolation for relative positions improves long context transformers. arXiv preprint, 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We deeply appreciate all reviewers for their insightful and detailed reviews, questions, and comments on our work. We assure the reviewers that all the answers and discussions will be incorporated into our final manuscript. 

We are encouraged to see that the reviewers recognized that our method is novel and significant (CVfs, ksWg, U8WH), our experiments are thorough and detailed (ow6J, CVfs, ksWg, U8WH), and our theoretical analysis is insightful and well-motivates our approach (ow6J, CVfs, ksWg, U8WH).

Please check out our PDF file containing:
* Fig. 1: ablations on the input formats and the model size.
* Fig. 2: Parity task solved with position coupling + scratchpad.
* Fig. 3: Testing on separate operand lengths.
* Fig. 4: RoPE + Position Coupling.

Now, we will provide our response to two commonly raised questions.

## **1. Ablations on Input Formats**
* Reviewers raised a question on the robustness of our method to input formats (CVfs) and a concern that input formatting is problem-aware processing (ksWg).
* We would like to clarify that our input format is primarily selected to simplify the algorithm of solving the addition task, not through extensive ablation studies. Thus, we are not arguing that our choice of input format is empirically optimal for training Transformers.
* However, we note that applying proper input formatting is crucial and natural in general ML. Even in a simple image classification task, appropriate standardization and augmentation are often useful. When applying these techniques, we often utilize the fact that typical image pixels range in [0, 255] (to apply standardization) and that appropriate augmentation methods may differ by the image type. Hence, enough understanding of the task leads to proper input processing that helps the model to effectively solve a given task.
* Our additional experiments on position coupling with various input formats show that the model’s performance varies with different input formats (refer to the attachment). It is expected, as the complexity of the algorithm that the model should learn changes according to the input format.
  * Small models (1-layer 4-head) achieve near-perfect generalization when the numbers are zero-padded and the answer or all numbers are reversed. We believe this is because the combination of zero-padding and reversing enabled a small Transformer to learn a simple length-generalizing algorithm. If we flip the answer or all the numbers without zero-padding, small models exhibit a bit worse in-distribution performance and a restricted length generalization capability. Without reversing, the models perform poorly.
  * Larger models (6-layer 16-head) perform better than the small model when the numbers are no longer zero-padded (especially when the answer is reversed). We believe this is because the task-solving algorithm with reversing and without zero-padding that the model should learn is more sophisticated, which larger models can learn more easily. Contrarily, we observe a degradation in performance when we add zero-padding in the larger model, which suggests that the model may have learned a "shortcut" due to its (overly) strong expressive power relative to the problem's complexity. (Refer to the next question for more details on this matter.)

## **2. Deeper models seem to perform worse. Why?**
* Some reviewers expressed concerns that position coupling performs worse when applied to deeper networks. Here, we provide our thoughts on this phenomenon.
* The first thing to note is that position coupling enables the model to solve the addition problem using a much simpler algorithm. As proven in Theorem 5.1, even a 1-layer model is sufficiently expressive. We hypothesize that the reason for the outstanding performance of a 1-layer model is that the architecture is simple so the model has no way to fit the training samples other than by learning the true (or length-generalizable) function.
* In contrast, we believe that deeper models perform poorly because their larger expressive capacity allows them to learn shortcuts to fit the training distribution that may not generalize well across different lengths. This phenomenon differs from classical overfitting in that these models still generalize well to in-distribution samples.
* A similar observation was made in [1]: there exist non-extrapolating models that generalize well to in-distribution samples but struggle with longer ones. The authors interpreted this phenomenon as indicating that there might be unexpected ways to solve the problem and the model may rely on shortcut-like solutions that work for in-distribution samples but fail on longer samples.
* The superior performance of the position coupling scheme without zero-padding in deeper models also aligns with our interpretation. In this scenario, the query sequence becomes less consistent: the lengths of two operands may differ and some position IDs appear only once. This makes the functions the model needs to learn more complex. This complexity is evident from 1-layer experiments that the no-padding scheme fails for even in-distribution generalization. We believe that such a complex structure of the target function compels the model to learn the true function rather than a shortcut-like solution, resulting in strong length generalization performance.
* Combining the answers for the first two questions, we do not believe that position coupling itself is bad for deep models. The performance degradation on deep models can also be attributed to the input formats, which may make the task-solving algorithm way simpler, and other training details such as small dataset size, which can make the model easily overfit to training distribution.

Again, we deeply thank all reviewers for their time and effort in reviewing our work. We are excited to hear more feedback.

Warm regards,  
Authors

---
**References**

[1] Zhang et al., Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint, 2022.

Meta-review: The paper studies the issue that Transformers are not able to generalize algorithms learned on sequences of length $n$ to longer sequences $n+k$ for even moderate values of $k$. This is demonstrated on many algorithmic toy tasks like addition where the algorithm is known, relatively simple, and yet Transformers still fail. Theoretical results for the 1-layer cases are presented to show that the proposed solution of "Position Coupling" theoretically (and empirically) improves the representational power of the Transformers. 

All reviewers are generally positive on the paper. The work is significantly improved in rebuttal by showing the parity problem can also be tackled with their approach as can other base positional embeddings. This gives the work a broader potential scope and applicability. 

There are still weaknesses that could be better incorporated into the manuscript to set appropriate expectations. In particular, it is not made obvious that the position coupling strategy is single-purpose in that it must be specified for each unique problem, and does not act as a more generic "plug-and-play" later/encoding that can be used for arbitrary problems. This reduces the scope to a theoretical inspection of the limits of Transformers and potential paths forward rather than a generally useful tool. Also noted by the reviewers is the drop in performance as more layers are added. 

However, immediate practical utility is not a prerequisite for publication. The work is overall interesting, sheds light on a new approach to talking a real limitation, and is relatively easy to implement. As such, I recommend the paper for acceptance but note that the reviewers should temper the reader's expectation to the manual work needed to apply position coupling to new tasks.
Decision: Accept (poster)