Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure

Authors: hanseul  cho
jaeyoung  cha
pranjal  awasthi
anupam  gupta
chulhee  yun
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:55:15

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).

Main paper: 
 introduction: 
 Since the appearance of a sequence-to-sequence 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-distribution 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 a 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 et 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-generalize?In this paper, we propose position coupling, a 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 to 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).• 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 a 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 a striking similarity between the attention patterns from our theoretical construction and those extracted from a Transformer trained with a standard optimizer (Section 5. 1.1) . As a complementary result, we also prove that any 1-layer Transformer without positional information cannot fully solve any permutation-sensitive 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 6.1, in which we also provide a theoretical construction of a 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-dimensional structures (Section 6.2). Moreover, we evaluate position coupling on some other tasks (addition with multiple operands, copy/reverse) in Appendix B.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_ONES, 3=IS_BOS, 4=IS_MUL, 5=IS_EQUAL, 6=IS_OP2_ONE, 7=IS_OP2_TEN, 8=OP2_ONE, 9=OP2_TEN, 10=OP1_SHIFT0, 11=OP1_SHIFT1, 12=OP1_SHIFT2, 13=OP1_SHIFT3, 14=OP1_SHIFT4, 15=RESULT1, 16=RESULT2, 17=RESULT3, 18=RESULT4, 19=PRE_PROD, 20=PRE_CARRY, 21=PRE_EOS1, 22=PRE_EOS2{23,...,32}=PROD, 33=IS_EOS, 34=MASK, and for the last P dimensions ({35, ..., 34 + P }), we named them as POS_2_MASK.The initial token embedding fills only NUM, FULL_ONES, IS_BOS, IS_MUL, and IS_EQUAL, 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 we fill the first 5 dimensions.Dimension 1 (NUM). For a number token (0, . . . , 9), we put itself into the dimension NUM. For the other tokens (×, =, $), we put 0.Dimension 2 (FULL_ONES). We put 1 everywhere in this dimension.Dimension 3 (IS_BOS). For a special token '$', we put 1 into the dimension IS_BOS. Otherwise, we put 0.Dimension 4 (IS_MUL). For a special token '×', we put 1 into the dimension IS_MUL. Otherwise, we put 0.Dimension 5 (IS_EQUAL). For a special token '=', we put 1 into the dimension IS_EQUAL.Otherwise, we put 0. 
preliminaries: 
 We focus on decoder-only Transformers that solve the tasks using next-token prediction (See Appendix A for a brief background on it). Since we study deterministic tasks with a unique answer, we consider greedy decoding throughout the paper.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 a 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 a reversed format. This enables a 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 a 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 -digit + N -digit addition with 0's, the input sequence becomes a max{M, N }-digit addition with (max{M, N } + 1)-digit 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-ofsequence) tokens at the beginning/end of the sequence. Lee et al. (2024) use the same token '$' for BOS and EOS tokens and observe that it is 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 Embeddings/Encodings (PE) Vaswani et 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. It 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 a Transformer decoder can 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 et al., 2023) .The goal of the second head is to fill in the blanks of the encoding matrix at dimensions PRE_CARRY and PRE_EOS. At dimension PRE_EOS, we will put (approximately) 1 if the next token would be the EOS token ('$'), otherwise, we will put strictly smaller numbers like (approximately) 2/3 and 1/2. What we will put at dimension PRE_CARRY 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 3 + 9 = 12. Since it is greater than or equal to 10, the least significant digits in the operands generate a carry 1. But in some cases, a pair of digits with a sum less than 10 can make a carry. Next, consider an addition 53 + 49 = 102. In the second least significant digits, An addition of 5 and 4 occurs. However, a carry is already produced in the least significant digits (3 + 9 = 12), so the total sum including the carry is 10, not 9. Thus, it also produces a carry. But how can we know the presence of a carry while only looking at the second least significant digits? The answer is to observe the second least significant digit in the sum, 0 of 102. Somehow, the consequence of adding 5 and 4 is 0, (or 10, implicitly) so it makes a carry.To generalize this explanation, let a and b be digits of the operands in the same significance, and c be a digit of the sum in the same significance as a and b. We find that the rule of recognizing that the addition of a and b generates a carry is that If a + b -c ∈ {9, 10}, then a carry is generated, Otherwise, then the carry is not generated. (29) Thus, it is crucial to store the information of a + b -c or any related one somewhere. In fact, we can store a + b + c at dimension PRE_CARRY of the encoding matrix, and it can be transformed into a + b -c and used later in the feed-forward layer. Formally, we aim to perform σ i + σ i+ℓ+1 + σ 3ℓ+5-i for each i ∈ {2, ..., ℓ + 1} and put its result at the (3ℓ + 5 -i)-th position (column) of the dimension PRE_CARRY (row). To this end, we again utilize our position embedding.Recall that d = 2P + 17 and let d QK,2 = P + 1. LetEQUATIONEQUATION)The linear transformations with matrices Q 2 and K 2 do two different jobs at once. ( 1) they take the dimensions POS_1 from the input encoding matrix and scale them up by√ M ; (2) Q 2 (K 2 , resp.) takes the dimension FULL_ONES (IS_BOS, resp.) and scale it up by √ M P . For concrete examples, refer to Tables 18 and 19 . By these, the attention score matrix C 2 := (K 2 X (0) ) ⊤ Q 2 X (0) 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 softmax operator; the dots represent the hidden M P 's. Let R 2 = A 2 -T 2 ∈ R N ×N be the error matrix, which is upper triangular as well. Its exact form is messy as well, but we can obtain the bounds of their entries. Consider a pair of indices (i, j)\ col j = 1 2 • • • ℓ + 1 ℓ + 2 ℓ + 3 • • • 2ℓ + 2 2ℓ + 3 2ℓ + 4 • • • 3ℓ + 3 3ℓ + 4 i = 1 M P M P • • • M P M P M P∈ [N ] 2 such that i ≤ j. Let x j = 1/[T 2 ] 1j ∈ {1, 2, 3, 4}. If [T 2 ] ij = 1 xj , [R 2 ] ij < 0 and -[R 2 ] ij ≤ 1 x j - e M Px j e M P + (j -x j )e M (P -2) = j -x j x j (x j e 2M + (j -x j )).(32) \ col j = 1 2 • • • ℓ + 1 ℓ + 2 ℓ + 3 • • • 2ℓ + 2 2ℓ + 3 2ℓ + 4 • • • 3ℓ + 3 3ℓ + 4 i = 1 1 1/2 • • • 1/2 1/2 1/3 • • • 1/3 1/3 1/4 • • • 1/4 1/2 2 *EQUATION)Now let d V,2 = 2 andEQUATIONEQUATIONThe linear combination with matrix U 2 V 2 does two jobs at once. First, it takes the dimension NUM from the encoding matrix, scales it up by 4, and puts it to the dimension PRE_CARRY. Second, it takes the dimension IS_BOS from the encoding matrix, scales it up by 2, and puts it to the dimension PRE_EOS. A concrete example is provided in Table 20 .Obtaining U 2 V 2 X (0) A 2 , its every entry is zero except at the dimensions PRE_CARRY and PRE_EOS.Observe that [U 2 V 2 X (0) ] (PRE_CARRY)1 = 0, because in the input encoding matrix, the dimension NUM starts with 0. Also, note again that it is enough to focus on the columns j ∈ {2ℓ + 3, . . . , 3ℓ + 4}, since we only care about the next-token prediction of the tokens after σ 2ℓ+3 ='='. Specifying the dimensions (i.e., the particular rows) for these columns, we haveEQUATIONEQUATIONEQUATIONRefer to Table 21 for a concrete example of computing U 2 V 2 X (0) T 2 . Also, for the softmax errors,EQUATIONEQUATIONLet us first obtain a bound of the softmax error term at dimension PRE_CARRY. If j = 2ℓ + 3, sinceEQUATIONandEQUATIONEQUATIONwhereEQUATIONandEQUATIONAndEQUATIONwhereEQUATIONandEQUATION)Next, we obtain a bound of the softmax error term at dimension PRE_EOS. SinceEQUATIONthe error term can be bounded asEQUATIONWe then can ensure that both [U 2 V 2 X (0) R 2 ] (PRE_SUM)j and [U 2 V 2 X (0) R 2 ] (PRE_EOS)j smaller than 0.1 for each j ∈ {2ℓ + 3, . . . , 3ℓ + 4}, by letting M ≥ 1 2 log(N ) + 3. The proof is similar to the one that is presented for head 1. 0) , continuing from Table 9 .0 P √ M v P 3 √ M v P 4 √ M v P 5 √ M v P 6 √ M v P 3 √ M v P 4 √ M v P 5 √ M v P 6 √ M v P 5 √ M v P √ M v P 3 √ M v P 2 P + 1: √ M P 0 0 0 0 0 0 0 0 0 0 0 0 Table 20: Example of U 2 V 2 X (I $ 6 5 3 + 0 4 9 = 2 0 0 1: NUM 0 0 0 0 0 0 0 0 0 0 0 0 2: IS_BOS 0 0 0 0 0 0 0 0 0 0 0 0 3: FULL_ONES 0 0 0 0 0 0 0 0 0 0 0 0 4: PRE_SUM 0 0 0 0 0 0 0 0 0 0 0 0 5: PRE_CARRY 024 20 12 0 0 16 36 0 8 0 28 0 6: PRE_EOS 2 0 0 0 0 0 0 0 0 0 0 0 7-16: SUM 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 17: IS_EOS 0 0 0 0 0 0 0 0 0 0 0 0 18-end: 20 . See Table 17 for 2/3 2/3 2/3 2/3 1/2 1/2 1/2 1 7-16: SUM 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 0 10 17POS_1 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 2P 0 2P Table 21: Example of U 2 V 2 X (0) T 2 , continuing from Table: IS_EOS 0 0 0 0 0 0 0 0 0 0 0 0 18-end: POS_1 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 2P 0 2P 
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 a group must have a unique semantic meaning so that a one-to-one correspondence between tokens in different groups can be made.Next, for each group of tokens, we assign a sequence of consecutive numbers (usually, positive integers) as position IDs, starting from a random number (at training time) or a fixed predetermined number (at evaluation time). We use random position IDs at training time for inducing length generalization by enabling all position embedding vectors to be trained, up to a pre-defined hyperparameter of maximum position ID (max_pos). 1 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 0 as the position IDs of special tokens like BOS/EOS tokens and the PAD token (padding for minibatch training and evaluation).Position Couplings reversed sum zero-paddings 6 5 3 + 0 4 9 = [BOS] 2 0 7 0 [EOS] 6 7 8 9 6 7 8 9 0 8 7 6 5 0 Position IDs 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) as a 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 '$' at the same time. For example, '653 + 49 = 702' becomes '$653 + 049 = 2070$'.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 us say that the random starting number is 6. In our example, we assign 6, 7, and 8 to the tokens in the operands, and assign 5, 6, 7, and 8 to the tokens in the sum in a reversed order: see Figure 2 . We remark that, first, we assign 9 as position IDs of '+' 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 a position ID of the most significant digit of the sum (which may be '0' due to the zero-padding) just because it is next to the number token with position ID 6, even though there are 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 is to prevent unnecessary coupling between '+' and the most significant digit of the sum. McLeish et al. (2024) proposes an analogous approach for solving arithmetic tasks, while they employ a different input format. We provide a 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 is an existing approach named index hinting (Zhou et al., 2024a) that applies a similar idea but to the input sequence. Index hinting is an input augmentation technique that places position markers in 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-paddings, where a, b, c, and d are consecutive index hints. Here, the starting hint character a is randomly selected during training, similar to our method of choosing the starting position ID. The reversed format and BOS/EOS 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 an index hint and a 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.1, 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 generalization of Transformers on 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 et al., 2024a) . Data Sampling. We opt for the balanced sampling in terms of the number of digits (Nogueira et al., 2021) . Given the maximum number of digits D max , we do balanced sampling for each operand in two steps. First, we sample the number of digits D ∈ [1, D max ] uniformly at random. Next, we sample an operand from [10 D-1 , 10 D -1] uniformly at random, except for D = 1 where we sample from [0, 9] . This procedure addresses the imbalance problem in the number of digits of operands.Training. We train decoder-only Transformer models from scratch. We properly choose max_pos so that the maximum testable length of summands is 200. We do 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. Longer Trained Sequences Lead to Longer Generalizable Lengths. We train 1-layer 4-head models with D max ∈ {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 is showcased in Figure 3 . We decide that a model successfully generalizes to a certain length of operands (referred to 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 generalizable length for the models trained on 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) 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 to the bad implicit bias of deeper models (learning shortcuts only to achieve in-distribution generalization) when learning a simple algorithm to solve the task. Since the theoretical construction of a 1-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 a 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, we also report the maximum accuracies (for the experiments in Figures 3 and 4 ) in Appendix B.3, showing that our 6-layer models can achieve near-perfect generalization for up to 200-digit additions as well. 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-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-padded and the response or all the numbers are reversed. We believe this is because the combination of zeropadding and reversing enabled a small Transformer to learn a simple length-generalizing algorithm. Zero-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.Larger models (6-layer 8-head) perform better than small models when the numbers are no longer zero-padded or 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 easily. 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 a "shortcut" due to its (overly) strong expressive power relative to the problem's complexity.Comparison with NoPE and APE (with random starting position ID). Our experiments demonstrate that simple PE methods like NoPE and random-start APE cannot length-generalize 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 1 imply that naively training all position embeddings does not necessarily help produce a strictly 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 we implemented ourselves since the original code is unavailable. From Figure 1 , we observe that index hinting indeed helps the model to 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-distribution validation accuracies.Particularly, the training accuracies of 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. 
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 a theoretical explanation by carefully constructing a 1-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 a 1-layer Transformer to solve the addition task in Section 5.2.For the sake of simplicity of presentation, we consider a Transformer without any normalization layers, as conventionally done in theoretical constructions by previous works (Awasthi and Gupta, 2023; Yun et al., 2020a,b) . For the sake of completeness, readers can find a 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 a depth-1 two-head decoderonly Transformer with coupled positions that solves the addition task with next-token prediction.Here, the operand length is at most 2 ⌊(d-17)/2⌋ -2, where the embedding dimension is d ≥ 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.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 11 and 17 in Appendix E 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: (1) 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 a carry 1 to the next position (with ID 6). 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 on 5-digit addition. They also observe that each head handles different subtasks of addition, such as digit-wise summation and carry detection.In Section 4.1, we observed that 1-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 a general proposition: a 1-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-1 finite-head decoder-only Transformer model T without positional encoding (NoPE). Given an input sequence I and its arbitrary permutation I ′ , if the last tokens of I and I ′ are identical, then the next tokens predicted by T will also be identical for both sequences when applying a 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 is identical, even if their order varies. However, the addition task is permutation-sensitive, meaning that the permuted queries may result in different responses. Therefore, the 1-layer Transformer cannot completely solve the task without positional encoding. It 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 as 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.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. 
applying position coupling beyond addition task: 
 To demonstrate the versatility of position coupling, we consider two other tasks in this section: N × 2 multiplication and a two-dimensional (2D) task. Other example tasks (e.g., addition with multiple summands, copy/reverse allowing duplicates) can be found in Appendix B.Here, we study length generalization on the N -digit × 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 as 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: see Figure 17 in Appendix B.5. The experimental results showcased in Figure 8 verify the efficacy of position coupling compared to NoPE and random-start APE. We observe that a 1-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. Unlike addition, position coupling for N × 2 multiplication is less intuitive, as 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, we can still construct a 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 × 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 ⌊(d-34)/6⌋ -3, where we denote the token embedding dimension by d ≥ 46.We defer the proof to Appendix G. This result suggests that the proposed position coupling scheme for the N × 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-1 model exhibits total failure even for in-distribution samples in our experiment.Now, we investigate the extension of position coupling for handling a 2D task, where the query and the response are 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 to generate a 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 in 8 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 a 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. 5 5 5   6 6 6   2 3 4   2 3 4   2 3 4   2 3 4   2 3 4   2 3 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 to be different for training samples. We evaluate the test performance on a square board with a width between 5 and 14 inclusively. We also employ a 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 distinct 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 at 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 is a 2D board, therefore the actual length generalization is from 81 to 144. Train widths/heights (5-9) NoPE (6-layer 8-head) (EM=99.50% at 5) Position Coupling (4-layer 8-head) (EM=98.07% at 12)  
conclusion: 
 Achieving length generalization of Transformers even in the simple case of the addition task has been a challenge that received a lot of attention. We propose position coupling, a variant of learned APE, which enables capturing task structure to improve the length generalization performance of Transformers for addition. We show that a Transformer trained on 1-30 digit addition can 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.Limitations & Future Directions. We intentionally limited ourselves to the tasks with an 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 a Transformer by training on short sequences. Designing the coupling of positions for tasks whose structure is implicit or black-box (e.g., for general NLP tasks) remains a fascinating next step: we leave the methodology for uncovering hidden structures and autonomously creating appropriate couplings (without manually designing them) for future work.We also leave two challenging arithmetic tasks to length-generalize for future work. One is the 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 a recent work (Cho et al., 2024) . Here we assume a basic input format (plain, no zero-padding), which is different from that we used in our experiment.A decoder-only Transformer returns an output sequence of the same length as the input sequence.One difference from a Transformer encoder is that the attention mechanism in a 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. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper focuses exclusively on arithmetic/algorithmic tasks, which inherently do not pose high risks for misuse.Guidelines:• The answer NA means that the paper poses no 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 a 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 a 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
: 
 Figure 1 : Methods for Length Generalization in the Integer Addition Task. We report exactmatch (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) ). 
 
 Appendix: Length Generalization in the Addition Tasks. Lee et al. (2024) observe that reversing the output in the addition task enables the model to learn a simple function. Shen et al. (2023) propose "Random Spacing" and "Recursive Scratchpad", achieving near-perfect generalization from 10-digits to 12digits addition. Zhou et al. (2024a) introduce "index hints", position markers placed in front of each token, in both the input and output of addition tasks. Most recently, Zhou et al. (2024b) demonstrate a 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 in the N × M Multiplication Task (M is fixed). Jelassi et al. (2023) investigate N × 3 using an encoder-only model and Duan and Shi (2023) study N × 1 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 a small number of long samples to the training set. Furthermore, although Duan and Shi (2023) provide perfect length generalization results even for test samples 10× longer than those observed during training, their approach requires a retraining step with hand-crafted 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. In Abbe et 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.proposes a new position embedding method called "Abacus". From a 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.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 performance. 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 a 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 × 2 task, fixing the length of the second operand by 2. While we achieve length generalization up to 90-digit multiplication by training the model on up to 40-digit multiplication, they report near-perfect in-distribution performance but poor length generalization.Finally and notably, we provide novel theoretical analyses, including (1) the constructive proof that a depth-1 Transformer equipped with position coupling can completely solve the addition task for exponentially long digits and ( 2) the impossibility of the same model being capable of the addition task. We also present theoretical results for the N × 2 multiplication task.B.1 Decimal Integer Addition Task: Scale-up to Length of 500Here, we 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 a train length of 160, we achieve excellent length generalization for 500-digit addition.On the other hand, training on sequences of length up to 40 or 80 is insufficient for extreme length generalization. The results demonstrate that position coupling, as an approach, is highly scalable. We see that while training with sequences of length up to 40 and 80 is insufficient for generalization to large lengths, at training length 160 we achieve strong performance for lengths up to 500. The experimental details can be found in Table 3 .In the main text, we mainly focus on 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 can be different (although the zero-padding is applied to ensure the consistency of the data format). See Figure 13 . A 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 ≥98% EM accuracy for 100-digit additions with a 6-layer 8-head model trained on 1-40. Moreover, they achieve a generalizable length of 45 for a model trained on 1-30, 25 for 1-20, and 10 for 1-10 (no length generalization). One big difference between their analysis and ours is they report the maximum accuracy for each testing length over trials, while we report the medians. Thus, we choose a bit lower threshold (95%) for generalizability than theirs. For a better comparison with Zhou et al. (2024b) , we report the maximum exact-match (EM) accuracies. See Figures 14 and 15 Figure 15 : Ablation on the number of layers (trained with position coupling). Here, we report maximum EM accuracies over 8 runs for each tested length.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 a position ID m is assigned to a key vector k, for every two consecutive entries of k (i.e., (k 2i-1 , k 2i )), we rotate by a predefined angle θ i multiplied by m:k 2i-1 k 2i → k 2i-1 cos mθ i -k 2i sin mθ i k 2i-1 sin mθ i + k 2i cos mθ i .We also apply a similar rotation to the query vectors (say, q with position ID n). As a result, the attention score for this key-query pair becomes a function of k, q, and the relative distance n -m.Unlike the original implementation of RoPE, we apply rotations based on the re-assigned position IDs according to our position coupling method. We incorporate this RoPE variant into every attention head of every layer. One more difference in implementation is that, during training, we randomly sampled an integer scaler ℓ ∈ {1, 2, 3, 4, 5} and multiplied it by the rotation angle. By such random re-scaling of rotation angles, we expect the model could handle unseen large rotation angles at test time.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 , our combination of RoPE and position coupling achieves a much better generalization up to operand lengths 100. We present an example of the position coupling method for N × 2 (N -digit by 2-digit) multiplication, which is omitted from the main text. See Section 6.1 for the experimental results. 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. 6 .Experiments. We train on sequences with operands of 1-40 digits. Our choice of max_pos is 102, so we test the operands of up to 100 digits. We investigate the performance of 3 different architectures, each with a different depth. The experimental results are 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 of Theorem 5.1 can be extended to addition tasks with multiple summands with slight adjustments to the feed-forward layer in the construction.Data Format & Position Coupling. Each token of the query sequence is a digit (10 distinct characters). We couple the positions in the query and the response by their correspondence. Note that the position ID assigned to the equal token is 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.a b a c = [BOS]a b a c [EOS] 3 4 5 6 2 0 3 4 5 6 0a b a c = [BOS]c a b a [EOS] 3 4 5 6 7 0 3 4 5 6 0Figure 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 5 to 300, at intervals of 5. While a 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 a 6-layer 8-head model for the latter cases as it is commonly used in the literature (Zhou et al., 2024a) . 7 .The experimental results are described in Figure 20 . For both copy and reverse tasks, position coupling exhibits near-perfect accuracy across the entire test length (7.5× for the trained length). In contrast, NoPE and random-start APE immediately fail to length-generalize.Here, we present the extra experimental results for training the minesweeper generator task with position coupling. Specifically, we compare the performance of two configurations: one 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 in Figure 21 . When sharing the same positional embedding layer, position coupling achieves over 98% accuracy on a 12×12 board, and maintains near 90% accuracy on a 14×14 board. However, with distinct positional embedding layers, position coupling only successfully generalizes to a 10×10 board. We currently do not have a 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. Train widths/heights (5-9) Position Coupling with Shared Embedding (4-layer 8-head) (EM=98.07% at 12) Position Coupling with Distinct Embedding (4-layer 8-head) (EM=99.68% at 10) Figure 21 : Exact-match accuracy (median over 4 runs) for minesweeper generator task, trained on sequences of length (5-9)×(5-9) with position coupling. For further experiment details, see Table 8 .Position coupling can be easily implemented on top of usual libraries of training transformer models like HuggingFace (Wolf et al., 2019) and Flaxformer4 since these libraries support an arbitrary array of position IDs (in the case of using APE). All we need is to build up a 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. 5It contains a custom implementation of decoder-only T5 (Raffel et al., 2020) established on top of PyTorch (Paszke et al., 2019) and Huggingface, including several PE methods. We additionally implement a 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), to follow the implementation details of Zhou et al. (2024b) . Here we detail the architecture of a depth-L, H-head decoder-only Transformer (Vaswani et al., 2017) . For a simple presentation, we ignore the normalization layers, as in Yun et al. (2020a) .Let V be the (ordered) vocabulary, a set of all tokens. Given an input sequence I ∈ V N and its length N , the encoding function Enc : V N → R d×N maps it toEQUATIONIt 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 l : R d×N → R d×N the operation of the l-th block (l ∈ [L]), so that 1) .X (l) := Tf l X (l-(EQUATIONThe block Tf l consists of a (causal) attention layer Att l : R d×N → R d×N and a (token-wise) feed-forward layer FF l : R d×N → R d×N , each of which contains a residual connection:EQUATIONwhere we denote by id : R d×N → R d×N an identity map.Each attention layer Att l consists of H attention heads. Its h-th head (h ∈ [H]) has matricesQ (l) h , K (l) h ∈ R d (l) QK,h ×d , V (l) h ∈ R d (l) V,h ×d and U (l) h ∈ R d×d (l)V,h as its parameters. 6 With these matrices, borrowing the notation from Yun et al. (2020a) , the attention layer with an input X ∈ R d×N can be written asEQUATIONHere the softmax operator takes a square matrix M ∈ R N ×N and outputs an N ×N upper-triangular column-stochastic 7 matrixEQUATIONwhere 1 {E} is an indicator function for a predicate E: it equals 1 if E 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 Y (l) := X (l-1) + Att l (X (l-1) ) ∈ R d×N which is a consequence of residual connection right after the attention layer.Each feed-forward layer FF l is a two-layer perceptron havingW (l) 1 ∈ R d F ×d , b (l) 1 ∈ R d F , W (l) 2 ∈ R d×d F , b (l)2 ∈ R d as its parameters. It applies the following map to each column y of an input Y :EQUATION)where ϕ is a component-wise activation function. That is, the feed-forward layer is defined asEQUATIONwhere 1 d is the d-dimensional vectors filled with 1's. Here we mainly use 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 et al., 2017) , ReGLU, and GEGLU (Shazeer, 2020) .The final component of the Transformer model is the decoding function Dec : R d×N → V N , which is composed of a linear readout and a (token-wise) arg-max operation. Here, the linear readout is simply a linear layer having W out ∈ R |V|×d as its parameter. The decoding function produces the output sequenceEQUATION)6 One can let dH = max l,h max{dEQUATIONV,h } as an inner dimension of each head. This makes our formal constructions a bit messier with redundant entries 0.7 Every entry is non-negative and the sum of entries in each column is 1.Here we show how to implement the addition by employing a 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 decoderonly Transformer with coupled positions that solves the addition task with next-token prediction.Here, the operand length is at most 2 ⌊(d-17)/2⌋ -2, where the embedding dimension is d ≥ 21.Organization of the Proof. A whole section is 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 a concrete example (Appendix E.3), where Enc maps an input sequence of length N to a d × N encoding matrix X (0) . • We devote a lot of pages to the detailed construction of the parameters of a causal attention layer Att 1 to generate desired attention patterns (Appendix E.4). The attention layer has two attention heads playing distinct roles: (1) preparing for a sum without considering carries; and(2) preparing for the carry prediction & EOS detection.• We provide a construction of a token-wise feed-forward neural network FF 1 which is a two-layer ReLU network (Appendix E.5). It consists of two subnetworks playing different roles: (1) producing one-hot vectors, each of which indicates a 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 .For the architecture of the decoder-only Transformer, we follow the notation introduced in Appendix D.Let e d i denote the i-th standard basis vector of R d . For example, e 3 1 = [1 0 0] ⊤ . Let I m be the m × m identity matrix. Let 0 p and 1 p denote the p-dimensional vectors filled with 0's and 1's, respectively. Similarly, let 0 m×n denote the m × 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 i• and A •j , respectively. Given two non-negative integers a and b, let ℓ(a, b) be the length of a longer one between a and b. For example, ℓ(12, 3456) = 4.Consider an ordered vocabulary V = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, +, =, $). We include a special token '$' that plays the role of both the beginning-of-sequence (BOS) token and the end-of-sequence (EOS) token. 8 We denote V k as k-th element of V. For instance, V 4 = 3 and V 13 = $. Lastly, since we employ only one Transformer block, we omit the superscripts (l) in the parameter matrices/vectors and the size of dimensions d We seek to perform an addition a + b = c using next-token prediction. To this end, we want to transform it into an input sequence I = $A + B = C of an appropriate format. Note that the EOS token is the last token that needs to be predicted, so we exclude EOS in the input sequence. Let ℓ := ℓ(a, b).We first zero-pad the shorter one between a and b to match the length of the part A and part B as ℓ. Sometimes, the sum c might be longer than a or b due to a carry. To make the length of the part C consistent, we also put a zero-pad in front of c to set its length as ℓ + 1. Also, to ease calculating the addition with next-token prediction, we reverse the sum c to make the part C. For example, if 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 are zero-paddings, and the blue digits are the reversed sum.To recap, the input sequence I = σ 1 σ 2 . . . σ N ∈ V N of length N = 3ℓ + 4 consists of six parts:1. the BOS token σ 1 = '$' 2. the first operand A = σ 2 . . . σ ℓ+1 where σ i ∈ {0, . . . , 9};3. the addition symbol σ ℓ+2 = '+'; 4. the second operand B = σ ℓ+3 . . . σ 2ℓ+2 where σ i ∈ {0, . . . , 9}; 5. the equality symbol σ 2ℓ+3 = '='; 6. the (reversed) sum C = σ 2ℓ+4 . . . σ 3ℓ+4 where σ i ∈ {0, . . . , 9}.Note that the part C might be incomplete (i.e., N < 3ℓ + 4) 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 = σ 1 . . . σ N to produce an output sequenceO = σ ′ 1 . . . σ ′ N where σ ′ 2ℓ+3 . . . σ ′ N -1 = C = σ 2ℓ+4 . . . σ N and σ ′ N = '$' (EOS).We plan to produce an input encoding, given an input sequence I designed as above. The encoding matrix X (0) is of size d × N : each column represents an embedding vector for a token, while each row represents a particular named dimension. What we mean by named dimension is that we give a name to each dimension for a clear description of our formal construction.We construct an input encoding by concatenating the token embedding and the position embedding, which can be viewed as a sum of two different embedding matrices of the same size.Table 9 : Example initial encoding. Here we consider the input sequence $653 + 049 = 2070 and the starting position ID is chosen as s = 2. The vectors v P □ are defined in Equation ( 11). The gray rows will be filled in later. 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.EQUATIONDimension 1 (NUM). For a number token (0, . . . , 9), we put itself in the dimension NUM. For the other tokens (+, =, $), we put 0.Dimension 2 (IS_BOS). For a special token '$', we put 1 in the dimension IS_BOS. Otherwise, we put 0.Dimension 3 (FULL_ONES). We put 1 everywhere in this dimension.Before constructing a position embedding, we specify the coupled position IDs for the addition task.Let max_pos be a hyperparameter of the maximum position IDs, where position IDs are non-negative integers. Basically, we match the significance of the digits: e.g., a least significant digit is always coupled to the other least significant digits. To this end, we first randomly choose a starting position ID s ∈ [max_pos -ℓ -1]. (For that, max_pos ≥ ℓ + 2 must hold.) Then we allocate the position IDs of token σ i in the input sequence I = σ 1 . . . σ N asp(i) =        0, i = 1, s + i -1, i = 2, . . . , ℓ + 2, s + i -(ℓ + 2), i = ℓ + 3, . . . , 2ℓ + 3, s + (3ℓ + 4) -i i = 2ℓ + 4, . . . , 3ℓ + 4. (9)Recall that N = 3ℓ + 4. Also, observe that for i ∈ {2, . . . , ℓ + 1}, p(i) = p(i + ℓ + 1) = p(3ℓ + 5 -i) = s + i, (10) which couples the position of (ℓ -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 '$'.Before moving on to the positional embedding, we definev D k (k ∈ [2 D ]) as v D k = (-1) b (D,k) i D i=1 ∈ R D (11)where bEQUATION)We remark that the points v D k are the vertices of D-dimensional hypercube with side length 2, centered at the origin. 9 Note that for k ̸ = l,EQUATION)Now we explain the position embedding. It consists of 2P dimensions, which eventually become from 18-th to (2P + 17)-th dimension after concatenation. If p(i) = 0, we let 0 2P as a position embedding vector. For the positive position IDs p(i) ≥ 1, we let a concatenationEQUATIONas a position embedding vector of a token σ i . (In case of p(i) = 2 P , we use v P 1 instead of v P p(i)+1 .) We call the former P dimensions for the position embedding as POS_1 and the latter P dimensions as POS_2.Concatenating the token embedding and the position embedding, we 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 v P k 's is max_pos = 2 P = 2 ⌊(d-17)/2⌋ . Therefore, the length of an operand must be ℓ ≤ max_pos -2 = 2 ⌊(d-17)/2⌋ -2.The goal of the causal attention layer is to fill in the zero-blanks 10 of the encoding matrix at dimensions PRE_SUM, PRE_CARRY, and PRE_EOS. We divide the roles into two different heads.The goal of the first head is to perform a digit-wise addition and to fill in the blanks of the encoding matrix at dimension PRE_SUM. Later, using this dimension, combined with the dimension PRE_CARRY, we will be able to perform the next-token prediction for addition. For now, we do not care about the carries, which will be dealt with in a later section. Formally, we aim to perform σ i + σ i+ℓ+1 for each i ∈ {2, • • • , ℓ + 1} and put its result at the (3ℓ + 4 -i)-th position (column) of the dimension PRE_SUM (row). To this end, we utilize our position embedding.Recall that d = 2P + 17 and let d QK,1 = P + 1. Let M > 0 be a number determined later. LetEQUATIONEQUATION)9 The choice of the vectors v D k is not strict. They only need to have the same length and be distinguishable (for at least a constant order) in terms of inner products. That is, there should be a noticeable difference betweenv D k 2 and v D k , v D lfor k ̸ = l. 10 Such 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.The linear transformations with matrices Q 1 and K 1 do two different jobs at once. (1) Q 1 (K 1 , resp.) takes the dimensions POS_1 (POS_2, resp.) from the input encoding matrix and scale them up by√ M ; (2) Q 1 (K 1 , resp.) takes the dimension FULL_ONES (IS_BOS, resp.) and scale it up by √ M P . For concrete examples, please refer to Tables 12 and 13 . By these, the attention score matrix C 1 := (K 1 X (0) ) ⊤ Q 1 X (0) becomes as in Table 10 . The blanks in Table 10 are the numbers smaller than M (P -2); the asterisks ('*') are the entries (or lower triangular submatrices) ignored by the causal softmax operator; the dots represents the hidden M P 's. \ col j = 1 2 3 • • • ℓ + 2 ℓ + 3 ℓ + 4 • • • 2ℓ + 3 • • • 3ℓ + 2 3ℓ + 3 3ℓ + 4 i = 1 M P M P M P • • • M P M P M P • • • M P • •\ col j = 1 2 3 • • • ℓ + 2 ℓ + 3 ℓ + 4 • • • 2ℓ + 3 • • • 3ℓ + 2 3ℓ + 3 3ℓ + 4 i = 1 1 1 1/2 • • • 1/2 1 1/3 • • • 1/3 • • • 1/3 1 1 2 0 0 1/2 0 0 1/3 0 1/3 0 0 . . . . . . . . . . . . ℓ + 1 0 0 0 1/2 0 0 1/3 0 0 0 ℓ + 2 0 0 0 0 0 0 0 0 0 0 ℓ + 3 0 0 0 0 0 1/3 0 1/3 0 0 . . . . . . . . . 2ℓ + 2 0 0 0 0 0 0 1/3 0 0 0 2ℓ + 3 0 0 0 0 0 0 0 0 0 0 . . . 3ℓ + 4 0 0 0 0 0 0 0 0 0 0 Let R 1 = A 1 -T 1 ∈ R N ×Nbe the error matrix, which is upper triangular. Its exact form is messy as well, but we can obtain the bounds of their entries. Consider a pair of indices (i, j)EQUATION)On the other hand, if -2) x j e M P + (j -x j )e M (P -2) = 1[T 1 ] ij = 0, [R 1 ] ij > 0 and [R 1 ] ij ≤ e M (Px j e 2M + (j -x j ).Now let d V,1 = 1 andEQUATIONEQUATIONThe linear transformation with matrix U 1 V 1 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 U 1 V 1 X (0) A 1 , its every entry is zero except at the dimension PRE_SUM. Observe that [U 1 V 1 X (0) ] (PRE_SUM)1 = 0, because in the input encoding matrix, the dimension NUM starts with 0. Also, note that it is enough to focus on the columns j ∈ {2ℓ + 3, . . . , 3ℓ + 4} since we only care about the next-token prediction of the tokens after σ 2ℓ+3 ='='. Specifying the dimension (i.e., the particular row) for these columns, we haveEQUATIONEQUATION)Refer to Table 15 for a concrete example of computingEQUATIONSpecifically, if j ∈ {2ℓ + 3, . . . , 3ℓ + 2} (thusEQUATIONwhereEQUATIONholds by Equation ( 17), andEQUATIONholds by Equation ( 18). On the other hand, if j ∈ {3ℓ + 3, 3ℓ + 4},EQUATION)One can easily prove these inequalities by using the bounds of [R 1 ] ij 's and the fact that the entries inX (0) (NUM)• lie in the interval [0, 9].If we let M ≥ 1 2 log(N -1) + 3, we can ensure that [U 1 V 1 X (0) R 1 ] (PRE_SUM)j smaller than 0.1 for each j ∈ {2ℓ + 3, . . . , 3ℓ + 4}. The proof is simple: it is enough to checkEQUATION)Table 12 : Example of Q 1 X (0) , continuing from Table 9 .I $ 6 5 3 + 0 4 9 = 2 0 7 0 1-P : 0P √ M P √ M PTable 13 : Example of K 1 X (0) , continuing from Table 9 . 0) , continuing from Table 9 . 14 . See Table 11 for I $ 6 5 3 + 0 4 9 = 2 0 7 0 1-P : 0P √ M v P 4 √ M v P 5 √ M v P 6 √ M v P 7 √ M v P 4 √ M v P 5 √ M v P 6 √ M v P 7 √ M v P 6 √ M v P 5 √ M v P 4 √ M v P 3 P + 1: √ M P 0 0 0 0 0 0 0 0 0 0 0 0 Table 14: Example of U 1 V 1 X (EQUATION0 0 0 0 0 0 0 0 0 0 0 0 18-end: POS_1,POS_2 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P Table 15: Example of U 1 V 1 X (0) T 1 , continuing from Table0 0 0 0 0 0 0 0 0 0 0 0 18-end: POS_1,POS_2 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P E.4.So far we have computed the output of Att 1 operation. Passing through the residual connection, the output of the attention layer is the sum of the original input encoding matrix and the output of Att operation:EQUATION)Since the term h∈{1,2} U h V h X (0) T h has nonzero entries only at dimensions PRE_SUM, PRE_CARRY, and PRE_EOS, the residual connection plays a role of "filling in some blanks" in the input encoding matrix. A concrete example of the output of residual connection is 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 in the orange rows. The gray rows will be filled in later. The goal of the feed-forward layer is to fill in the blanks of the encoding matrix at dimensions SUM and IS_EOS. Be careful that the feed-forward layer can only implement token-wise mappings; a tokenwise 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 a one-hidden-layer ReLU network.For a token σ i for i ∈ {2ℓ + 3, . . . , 3ℓ + 3} (from '=' token to the second ), we will put a standard unit vector e 10 k+1 to dimensions SUM if the next token is k ∈ {0, . . . , 9}. Recall from the discussion in Appendix E.4.2 that we can judge whether a carry 1 is generated at a certain position by exploiting only the digits (of the operands and the sum) in the same significance. Bringing the notation, let a and b be digits of the operands in the same significance, and c be a digit of the sum in the same significance as a and b. Then the rule of recognizing that the addition of a and b generates a carry is that If a + b -c ∈ {9, 10}, then a carry is generated, Otherwise: if a + b -c ∈ {-1, 0}, then the carry is not generated.(52)A simple case analysis shows that the value of a + b -c must be one of -1, 0, 9, and 10. Let us briefly check this claim in our example:6 + 0 -7 = -1; no carry from 6+0 (53) 5 + 4 -0 = 9;there is a carry from 5+4 (54)3 + 9 -2 = 10.there is a carry from 3+9 (55)Recall that a noisy version of a + b + c is already stored at dimension PRE_CARRY of Y (1) , and c is exactly at dimension NUM. Thus, we can (approximately) implement a + b -c for a token σ j byEQUATION)This is a kind of token-wise linear transform, so we do not need to consume any hidden layer (with ReLU activation ϕ) to implement it.Combining with Y (0) (PRE_SUM)j , a noisy version of 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.We first describe the formal construction of feed-forward network FF 1 for dimensions SUM and IS_EOS and then explain the intuition behind the construction. For the example result of applying the feed-forward network is presented in Table 23 . and consider a one-hidden-layer ReLU network f k : R → R (k = 0, 1, . . . , 9) defined asEQUATION)Then we construct a subnetwork of our feed-forward network for a token σ j by k for the answer of next-token prediction 'k'. There are two cases where we need to predict the next token as 'k':EQUATION• Case 1: (Addition without carry) = k mod 10 and no carry propagates.• Case 2: (Addition without carry) = k -1 mod 10 and there is a propagating carry 1.In the first case, due to the softmax error (with magnitude at most 0.1),EQUATIONEQUATIONIn the second case, again due to the softmax error (with magnitude at most 0.1),EQUATIONEQUATIONIn both cases,EQUATION)We can map the column Y Recall that our objective is to output 1 to the dimension k + 6 (among the dimensions 7, 8, . . . , 16 in SUM) and to output 0 to the other dimensions in SUM if we need to predict 'k' as the next token. To this end, it is enough to map the set [k, k + 0.5] ∩ [k + 10, k + 10.5] to 1 and to map the other sets (for different k's) to 0. This can be done by a ReLU network f k (x) is a ReLU network having two bumps at intervals [k -0.5, k + 1] and [k + 9.5, k + 11]. In particular,f k (x) = 1 if x ∈ [k, k + 0.5] ∪ [k + 10, k + 10.5]: see Figure 23 for an illustration.Lastly, we have a desired one-hot vector output for each j by taking a composition between g and [f 0 (•), . . . , f 9 (•)] ⊤ as written in Equation ( 60).We move on to the dimension IS_EOS. For a token σ j for j ∈ {2ℓ + 3, . . . , 3ℓ + 4}, if k is the next token, we will put 1 {k=$} to dimension IS_EOS: 1 if k is the special token '$' and 0 otherwise. To this end, we define a ReLU network h : R → R asEQUATIONThen, we can construct a subnetwork of our feed-forward network for a token σ j byEQUATION(PRE_EOS)j .(68)Explanation. Note that for columns j ∈ {2ℓ + 3, . . . , 3ℓ + 4}, if we consider the presence of softmax errors with magnitude at most 0.1, the values that Y(PRE_EOS)j can have lie in the set [0.4, 0.6] ∩ [2/3 -0.1, 2/3 + 0.1] ∩ [0.9, 1.1] ⊂ (-∞, 0.8) ∩ [0.9, ∞). We want to output 1 if Y(1) (PRE_EOS)j ≥ 0.9 and 0 otherwise: this can be done with the ReLU network h with two neurons.• In total, we consume 8 × 10 + 2 = 82 ReLU neurons in our feed-forward network FF 1 . However, it is possible to construct the addition Transformer with a smaller number of neurons, with a 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 feed-forward ReLU network plays the role of denoiser.Table 23 : Example output after applying the feed-forward network. The last task of the feed-forward layer is to pass FF 1 Y (1) through the residual connection. As a result, we have X (1) = Y (1) + FF 1 Y (1) . (69)A concrete example of the output of the second residual connection is showcased in Table 24 . As mentioned in Appendix D, the decoding function performs a linear readout (with a weight matrix W out ∈ R |V|×d ) and a (token-wise) arg-max operation. That is,EQUATION)where V k is the k-th element of V andEQUATIONEQUATION)The objective of the decoding function is to perform a proper next-token prediction for addition, especially utilizing the dimensions SUM and IS_EOS of X (1) .We now construct the weight matrix W out . For a token σ i , if the value of dimension IS_EOS of X (1) is 0, then the linear readout output the dimensions SUM as it is to return one of a number token (0-9). On the other hand, if the value of dimension IS_EOS is 1, then the linear readout outputs a large number (like 100 for example) for the token '$' to return EOS ($). This can be implemented by the weight matrix W out described in Table 25. Also, an example of applying the linear transform is showcased in Table 26 . 1-6: NUM-PRE_EOS 0 6 0 6 0 6 0 6 0 6 0 6 0 6 0 6 0 6 0 6 0 6 0 6 0 67: SUM 1 1 0 0 0 0 0 0 0 0 0 0 0 8: SUM 2 0 1 0 0 0 0 0 0 0 0 0 0 9: SUM 3 0 0 1 0 0 0 0 0 0 0 0 0 10: SUM 4 0 0 0 1 0 0 0 0 0 0 0 0 11: SUM 5 0 0 0 0 1 0 0 0 0 0 0 0 12: SUM 6 0 0 0 0 0 1 0 0 0 0 0 0 13: SUM 7 0 0 0 0 0 0 1 0 0 0 0 0 14: SUM 8 0 0 0 0 0 0 0 1 0 0 0 0 15: SUM 9 0 0 0 0 0 0 0 0 1 0 0 0 16: SUM 10 0 0 0 0 0 0 0 0 0 1 0 0 0 17: IS_EOS 0 0 0 0 0 0 0 0 0 0 0 100 18-end: POS_1, POS_2 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2P 0 2PTable 26 : Example output of linear readout (W out X (1) ), 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).EQUATION0 0 0 0 0 0 5 0 0 0 0 0 0 0 0 0 0 0 0 6 0 0 0 0 0 0 1 0 0 0 0 0 7 0 0 0 1 0 0 0 0 0 0 1 0 8 0 0 0 0 0 0 0 0 0 0 0 0 9 0 0 1 0 0 0 0 1 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 = 0 0 0 0 0 0 0 0 0 0 0 0 $ 100 100 100 100 100 0 0 0 0 0 0 0 100 Table 27 : Example output sequence O = Dec X (1) , continuing from Table 26 . The yellow cells in the bottom row exactly predict the next tokens. I $ 6 5 3 + 0 4 9 = 2 0 7 0O $ $ $ $ $ 1 6 9 2 0 7 0 $For the sake of readability, we 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 I and its arbitrary permutation I ′ , if the last tokens of I and I ′ are identical, then the next tokens predicted by T will also be identical for both sequences when applying a greedy decoding scheme.Remark. We assume the 1-layer (L = 1) H-head Transformer achitecture specified in Appendix D.Although it omits normalization layers, we remark that Proposition 5.2 remains valid even for the architecture with a 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.1. Throughout the proof, we denote the value/vector/matrix related to I ′ by appending ' ′ ' to it.Let encoding matrices generated from the input sequences I, I ′ ∈ V N asEQUATION)Since there is no positional encoding, the encoding function Enc(•) maps the same tokens to the same columns. In particular,I i = I ′ j implies X •i = X ′ •j . Since we assume that I ′ is a permutation of I such that I N = I ′ N , there exists a bijection π : [N ] → [N ] such that I ′ i = I π(i) for each i ∈ [N ] and π(N ) = N . Then, it follows that X ′ •i = X •(π(i))for each i and, specifically, X ′ •N = X •N . Recall that the single H-head attention layer Att : R d×N → R d×N operates as Att(X) = H h=1 Head h (X) where the attention head h is defined asHead h (X) := U h V h X • softmax (K h X) ⊤ Q h X ∈ R d×N , where Q h , K h ∈ R d QK ×d , V h ∈ R d V ×d and U h ∈ R d×d V . Claim: [Head h (X)] •N = [Head h (X ′ )] •N for all h ∈ [H].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)]•N = [Att(X ′ )] •N .Note that the operations after the attention layer-residual connections, FF, and Dec-all operate in a 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 ∈EQUATIONwhich are both stochastic (sum to 1) column vectors. Considering that we are taking the last column of the softmax output, it follows that s ′ i = s π(i) for each i ∈ [N ]: this can be proved by applying the definition of the softmax operation and the factEQUATION)Consequently, sinceEQUATIONwe haveEQUATION)Therefore, the claim holds. This concludes the proof.Here, we provide the Python code that calculates the maximum possible exact-match accuracy that a 1-layer Transformer with NoPE can achieve for the m-digit addition problem. Here we show how to implement the N × 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 × 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 ⌊(d-34)/6⌋ -3, where we denote the token embedding dimension by d ≥ 46.We note that our construction for the N × 2 multiplication task permits the use of multiple FFN layers at the second decoder block. However, we believe that there exists a potential improvement in our construction, wherein a 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 can reduce small attention scores to zero values, thereby clearly revealing the desired attention patterns.Consider an ordered vocabulary V = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ×, =, $). We include a special token '$' that plays the role of both the beginning-of-sequence (BOS) token and the end-of-sequence (EOS) token. We denote V k as k-th element of V. For instance, V 4 = 3 and V 13 = $. Unlike the addition task, our construction for the multiplication involves multiple layers and hence we do not omit the superscripts (l) in the parameter matrices/vectors and the size of dimensions.Our objective is to use next-token prediction for implementing a × b = c. To this end, we want to transform it into an input sequence I = $A × B = C of an appropriate format. Let ℓ a and ℓ b represent the lengths of a and b, respectively, and we denote their sum as ℓ = ℓ a + ℓ b . While our immediate focus is on the case where ℓ b = 2, it is worth noting that our approach can be extended to the case where ℓ b > 2, as the key insight for the construction does not rely on ℓ b . Thus, we present the input sequence and encoding function in a more general form applicable to ℓ b ≥ 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 c may either equal the sum of the lengths of a and b, or be less than the sum of their lengths by 1. Hence, we zero-pad in front of c for the latter case to fix the length of c by ℓ. We also reverse the response c to make the part C. For instance, if we have 312 × 24 = 7488, the input sequence transforms to $312 × 24 = 88470. If we have 589 × 62 = 36518, then the input sequence would be $589 × 62 = 81563. The red digit is a zero-padding, and the blue digits are the reversed product.To recap, the input sequenceI = σ 1 σ 2 . . . σ N ∈ V N of length N = 2ℓ + 3 consists of six parts:1. the BOS token σ 1 = '$' 2. the first operand A = σ 2 . . . σ ℓa+1 where σ i ∈ {0, . . . , 9}; 3. the multiplication symbol σ ℓa+2 = '×'; 4. the second operand B = σ ℓa+3 . . . σ ℓ+2 (note that ℓ = ℓ a + ℓ b ) where σ i ∈ {0, . . . , 9}; 5. the equality symbol σ ℓ+3 = '='; 6. the (reversed) product C = σ ℓ+4 . . . σ 2ℓ+3 where σ i ∈ {0, . . . , 9}.Note that the part C might be incomplete (i.e., N < 2ℓ + 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 = σ 1 . . . σ N to produce an output sequenceO = σ ′ 1 . . . σ ′ N where σ ′ ℓ+3 . . . σ ′ N -1 = C = σ ℓ+4 . . . σ N and σ ′ N = '$' (EOS).We now explain the input embedding for given an input sequence I designed as above. The embedding matrix X (0) is of size d × N : each column represents an embedding vector for a token, while each row represents a particular named dimension. We concatenate the token embedding and the position embedding, which can be viewed as a sum of two different embedding matrices of the same size.Table 28 : Example initial encoding. Here we consider the input sequence $7595 × 79 = 500006 and the starting position ID is chosen as s = 1. The vectors v P □ are defined in Equation ( 79). The gray rows will be filled in later. EQUATION: POS_2_MASK 0P 0P 0P 0P 0P 0P 0P 0P 0P 0P 0P 0P 0P 0P 0P (P +35)-(2P +34): POS_1 0P v P 3 v P 4 v P 5 v P 6 v P 7 v P 5 v P 6 v P 7 v P 6 v P 5 v P 4 v P 3 v P 2 v P 1 (2P +35)-(3P +34): POS_2 0P v P 4 v P 5 v P 6 v P 7 v P 8 v P 6 v P 7 v P 8 v P 7 v P 6 v P 5 v P 4 v P 3 v P 2 (3P +35)-(4P +34): POS_3 0P v P 5 v P 6 v P 7 v P 8 v P 9 v P 7 v P 8 v P 9 v P 8 v P 7 v P 6 v P 5 v P 4 v P 3 (4P +35)-(5P +34): POS_4 0P v P 6 v P 7 v P 8 v P 9 v P 10 v P 8 v P 9 v P 10 v P 9 v P 8 v P 7 v P 6 v P 5 v P 4 (5P +35)-(6P +34): POS_5 0P v P 7 v P 8 v P 9 v P 10 v P 11 v P 9 v P 10 v P 11 v P 10 v P 9 v P 8 v P 7 v P 6 v P 5 G.3.We now specify the allocation of coupled position IDs for the N × M multiplication task as the following: given an input sequenceEQUATIONCompared 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 we explain the position embedding. We utilize the same v D k (k ∈ [2 D ]) defined for the addition task, specificallyEQUATIONwhere b(D,k) i is defined as the i-th (from left) digit of D-digit binary representation of k -1. Using v Dk , we design the position embedding for each position ID p(i) byEQUATION)The first P dimensions of the position embedding are named as POS_1, and subsequent sets of P dimensions are named as POS_2, POS_3, POS_4, and POS_5, respectively. Thus, the position embedding is a 5P -dimensional vector. In case of p(i) + j (j ∈ [4]) exceeding 2 P , we use v P p(i)+j-2 P instead of v P p(i)+j . If p(i) = 0, we let 0 5P as a position embedding vector. By concatenating the token embedding and the position embedding, we get the input embedding X (0) . Specifically, the position embedding is placed under the token embedding ((P + 35)-th to (6P + 34)-th dimension). See Table 9 for an example. As a result, the total embedding dimension is d = 6P + 34. Note the maximum possible position ID that can be represented with v P k 's is max_pos = 2 P = 2 ⌊(d-34)/6⌋ . Therefore, the length of the first operand must be ℓ a ≤ max_posℓ b -1 = 2 ⌊(d-34)/6⌋ -ℓ b -1. For the case when ℓ b = 2, this inequality becomes ℓ a ≤ 2 ⌊(d-34)/6⌋ -3.Here, we provide an example that demonstrates how we construct the N × 2 multiplication. Consider the calculation 7595 × 79 = 600005. While a 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 is the same, and then calculate the sum of the pairwise products. For example, the number 116 in Table 30 is generated by 9 × 9 + 5 × 7, and the number 108 is generated by 5 × 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 a method is to provide a clearer intuition for determining which numbers from each operand we should attend to when predicting the next token.Table 29 : Multiplication I 7 5 9 5 × 7 9 6 8 3 5 5 5 3 1 6 5 6 0 0 0 0 5 Table 30: Multiplication II 7 5 9 5 × 7 9 4 5 1 1 6 1 0 8 9 8 4 9 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 sum of 9, 9, 1, and an additional 1, representing the carry from the previous step. Similar to the explanation in E.4.2, we highlight that the carry 1 can be detected by computing 8 (ones digit of 98) + 0 (tens digit of 108) + 1 (hundreds digit of 116) -0 (current token): yielding a result of 9, indicating the occurrence of a carry 1.In summary, 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 × 9 + 7 × 7, • 98 = 7 × 9 + 5 + 7, • 108 = 5 × 9 + 9 + 7, • 116 = 9 × 9 + 5 + 7.Thus, for predicting the next token, we need 0, 7, 5, 9, 5, 9, 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 we will see in the later subsection, a depth-2 decoder-only Transformer model can be constructed to fill OP2_ONE by 9, OP2_TEN by 7, and OP1_SHIFT0 to OP1_SHIFT4 by 0, 7, 5, 9, and 5, respectively. One may be concerned that 0 is not given in the first operand at 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 0 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.Consequently, the embedding vector of the current token 0 (the 0 preceding 60) will be structured as the left-most table in Table 31 , with some irrelevant dimensions omitted for readability. We then utilize a feed-forward layer to fill• RESULT1 with OP1_SHIFT0 × OP2_ONE + OP1_SHIFT1 × OP2_TEN, • RESULT2 with OP1_SHIFT1 × OP2_ONE + OP1_SHIFT2 × OP2_TEN, • RESULT3 with OP1_SHIFT2 × OP2_ONE + OP1_SHIFT3 × OP2_TEN, • RESULT4 with OP1_SHIFT3 × OP2_ONE + OP1_SHIFT4 × 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 RESULT1 + 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 in the right-most table of Table 31 . Once this process is done, we can finally predict the next token by the following two steps:• CARRY =    0, if PRE_CARRY -NUM ∈ {-2, -1, 0}, 1, if PRE_CARRY -NUM ∈ {8, 9, 10}, 2, if PRE_CARRY -NUM ∈ {18, 19, 20},• NEXT_TOKEN = PRE_PROD + CARRY (mod 10). 31 . The goal of the first Transformer block is to fill IS_OP2_ONE (6-th dimension) and IS_OP2_TEN (7-th dimension) by 1 if the token corresponds to the ones or tens digit of the second operand, respectively, and 0 otherwise. These two dimensions enable the filling of OP2_ONE (8-th dimension) and OP2_TEN (9-th dimension) at the second Transformer block. Furthermore, we will fill MASK (34-th dimension) in the first block, which will serve as a base for filling OP1_SHIFT0 to OP1_SHIFT4 in the second block. Thus, we currently have 3 objectives(IS_OP2_ONE, IS_OP2_TEN, MASK), each of which will be addressed by an individual head.The goal of the first head is to make the dimension IS_OP2_ONE as a one-hot row vector, where 1 is placed only at the token corresponding to the ones digit of the second operand.Recall that d = 6P + 34 and let d QK,11 = P + 1. Let M > 0 be a sufficiently large positive real number. LetEQUATIONEQUATION)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, we provide an illustrative example for each step. We will also simply regard M as a sufficiently large real scalar and thus the attention values can be clearly separated after going through the softmax operation.The matrix Q(1)1 maps the embedding matrix X (0) into a query matrix Q(1)1 X (0) ∈ R (P +1)×N , where the first P rows are obtained by copying from the dimensions POS_2 and scaling by √ M , while the last row is the copy of the dimension FULL_ONES scaled by √ M P . Similarly, the matrixK (1)1 maps the embedding matrix to a key matrix K(1) 1 X (0) ∈ R (P +1)×N . In this case, the first P rows are obtained by copying from the dimensions POS_1 and scaled by √ M , with the last row being the dimension IS_BOS, scaled by √ M P . For concrete examples, refer to Tables 32 and 33 .Table 32: Example of Q (1)1 X (0) , continuing from Table 28 .I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1-P : 0 P √ M v P 4 √ M v P 5 √ M v P 6 √ M v P 7 √ M v P 8 √ M v P 6 √ M v P 7 √ M v P 8 √ M v P 7 √ M v P 6 √ M v P 5 √ M v P 4 √ M v P 3 √ M v P 2 Table 34: Example of A (1)1 (with explicit row/column indices and sufficiently large M ), continuing from Tables 32 and 33 (1)1 X (0) A1 , continuing from Tables 34 and 35 In the previous head, we set the dimension IS_OP2_ONE (6-th dimension) to a one-hot row vector, where 1 is placed only in the token corresponding to the ones digit of the second operand. The objective of Attention head 2 is to fill the dimension IS_OP2_TEN (7-th dimension) similarly to IS_OP2_ONE, but with 1 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 to the construction of attention head 1, we only push √ M I P to the next block for designingQ (1)2 . Specifically, Q2 and K(1)2 are defined asEQUATIONEQUATION)where d QK,12 is set to P + 1. We refer to Tables 37 and 38 for specific examples.Table 37 :Example of Q (1)2 X (0) , continuing from Table 28 .I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1-P : 0 P √ M v P 5 √ M v P 6 √ M v P 7 √ M v P 8 √ M v P 9 √ M v P 7 √ M v P 8 √ M v P 9 √ M v P 8 √ M v P 7 √ M v P 6 √ M v P 5 √ M v P 4 M P √ M P √ M P Table 38: Example of K (1)2 X (0) , continuing from Table 28 .I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1-P : 0 P √ M v P 3 √ M v P 4 √ M v P 5 √ M v P 6 √ M v P 7 √ M v P 5 √ M v P 6 √ M v P 7 √ M v P 6 √ M v P 5 √ M v P 4 √ M v P 3 √ M v P 2 √ M v P 1 P + 1: √ M P 0 0 0 0 0 0 0 0 0 0 0 0 0 0By these, the attention score matrix C(1) 2:= (K (1) 2 X (0) ) ⊤ Q (1)2 X (0) and the attention matrixA (1) 2 := softmax(C (1)2 ) ∈ R N ×N can be obtained, and the example ofA (1)2 is provided in Table 39 . and U(1)2 the same to that of the previous head. That is, with d V,12 = 1,EQUATIONEQUATIONand the example of U(1) 2 X (0) and U(1)2 X (0) A (1)2 is provided in Tables 40 and 41 . Consequently, the matrixU (1) 2 V (1) 2 X (0) A (1)2 is a matrix that matches the size of the input embedding matrix and is filled with zeroes, except for a unique 1 located at the tens place of the second operand in the input sequence, with the dimension IS_OP2_TEN (7-th dimension). (1)2 X (0) A (1)2 , continuing from Tables 39 and 40 The goal of Attention head 3 is to generate a binary mask at the dimension MASK (34-th dimension), with '0' placed before the multiplication symbol (×) and '1' placed starting from the multiplication symbol to the end.To this end, we set d QK,13 = 1 and design query and key weights byEQUATIONEQUATION)The matrices Q3 X (0) and K3 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 (1)3 X (0) , continuing from Table 28 . I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 Table 43 : Example of K (1) 3 X (0) , continuing from Table 28 . I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1: 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 By these, the attention score matrix C(1) 3:= (K (1) 3 X (0) ) ⊤ Q (1)3 X (0) and the attention matrix A (1) 3 := softmax(C (1) 3 ) ∈ R N ×N can be obtained and the example of A (1) 3 is provided in Table 44 . 42 and 43 . row \ col j = 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 and U1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 i = 1 0 1 1/2 1/(1)3 by d V,13 = 1 and V(1) 3EQUATIONEQUATIONThe example of U3 V3 X (0) and UEQUATION3 is provided in Tables 45 and 46 . Consequently, the matrix U(1) 3 V (1) 3 X (0) A (1)3 is a matrix that matches the size of the input embedding matrix and is filled with 1 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 seem somewhat unclear. We note that the output of Attention head 3 will be utilized to fill the dimensions POS_2_MASK in the subsequent FFN layer, and this POS_2_MASK plays a crucial role in designing the key matrices in the Attention heads 3 to 7 at the second Transformer block. EQUATION3 , continuing from Tables 44 and 45 EQUATIONh .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 5 9 5 × 7 9 = 5 0 0 0 0 6 35-(P +34): POS_2_MASK 0 P v P 4 v P 5 v P 6 v P 7 0 P 0 P 0 P 0 P 0 P 0 P 0 P 0 P 0 P 0 PThe last task of the feed-forward layer is to pass FF 1 Y (1) through the residual connection. As a result, we have 1) .EQUATION(97) This is the end of the first Transformer block, and a concrete example of X (1) is illustrated in Table 49 . Consider a 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, fill OP2_ONE, OP2_TEN, and OP1_SHIFT0 to OP1_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), (i -4)-th least significant digit of the first operand, respectively. Our construction assigns each head to each dimension. The second goal is to fill PRE_EOS1 and PRE_EOS2 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 2 dimensions can be implemented within the same head for OP1_SHIFT0 and OP1_SHIFT2 respectively, thus requiring a total of seven heads.EQUATIONEQUATIONG.7.1 Attention Head 1: Copying the Ones Digit of the Second OperandThe objective of Attention head 1 is to fill the dimension OP2_ONE with the ones digit of the second operand. To do so, we design the weights by defining d QK,21 = 1 andEQUATIONEQUATION)We also define d V,21 = 1 andEQUATIONEQUATIONA concrete example of Q1 X (1) , K1 X (1) , A 2 1 , U(2)1 X (1) , and UEQUATION1 is provided in Tables 50 to 54. One might be concerned that in Table 54 , the dimension OP2_ONE is not 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. (2)1 X (1) , continuing from Table 49 . I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1   Table 51 : Example of K (2)1 X (1) , continuing from Table 49 . I $ 7 5 9 5 × 7 9 = 5 0 0 0 0 6 1: 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0Table 53: Example of U (2) 1 V (2)1 X (1) , continuing from Table 49 1 (with explicit row/column indices), continuing from Tables 50 and 51 . row \ col j = 1 2 3 4 5 6 7 8 9 10 11 12 13 14 151 1 1 1 1 1 1 1 1 1 1 1 1 1 i = 1 0 1 1/2 1/3 1/4 1/5 1/6 1/7 0 0 0 0 0 0 0 2 0 0 1/2 1/3 1/4 1/5 1/6 1/7 0 0 0 0 0 0 0 3 0 0 0 1/3 1/4 1/5 1/6 1/7 0 0 0 0 0 0 0 4 0 0 0 0 1/4 1/5 1/6 1/7 0 0 0 0 0 0 0 5 0 0 0 0 0 1/5 1/6 1/7 0 0 0 0 0 0 0 6 0 0 0 0 0 0 1/6 1/7 0 0 0 0 0 0 0 7 0 0 0 0 0 0 0 1/7 0 0 0 0 0 0 0 8 1 0 0 0 0 0 0 0 1 1 1 1 1 1 1 9 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 10 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 11 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 12 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 13 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 14 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 15 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 Table 54: Example of U (2) 1 V (2) 1 X (1) A (2)1 , continuing from Tables 52 and 53 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_OP2_TEN instead of IS_OP2_ONE for generating the key weight. We design the weights by defining d QK,22 = 1 andEQUATIONEQUATIONWe also define d V,22 = 1 andEQUATIONEQUATIONA concrete example of QEQUATION(2) 2 X (1) , andU (2) 2 V (2) 2 X (1) A (2)2 is provided 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. (2)2 X (1) , continuing from Table 49 . I $ 7 5 9 5 × 7 9 = 5 0 0 0 6 1:1 1 1 1 1 1 1 1 1 1 1 1 1 1 Table 56: Example of K (2)2 X (1) , continuing from Table 49 . I $ 7 5 9 5 × 7 9 = 5 0 0 0 6 1: 0 0 0 0 0 0 1 0 0 0 0 0 0 0 55 and 56 . row \ col j = 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15EQUATION(2) 2 X (1) , continuing from Table 49 . EQUATION2 , continuing from Tables 57 and 58 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 and OP2_TEN, respectively. For Attention head 3 to 7, we mainly focus on the first operand. Specifically, in Attention head 3, the goal is to fill the dimension OP1_SHIFT0 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 + 1)-th least significant digit of the first operand. For our example, we want to fill OP1_SHIFT0 of the token = by 5. Here, i ranges from 0 to ℓ a + 2, where the 0-th least significant digit of the response denotes the equal token. In cases where i ≥ ℓ a , we fill by 0.Additionally, the third head has an extra objective: filling the dimension PRE_EOS1. This dimension is utilized for EOS prediction in the subsequent FFN layer along with PRE_EOS2, which is 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(2) 3and U(2)3 described below. Unlike previous heads, Vand U(2)3 each have two elements, with each element contributing to one of the objectives. Our specific construction is as follows. With d QK,23 = P + 1,EQUATIONEQUATION)and with d V,23 = 2, VEQUATIONEQUATIONWe provide the examples in Tables 60 to 64 . We note that within the dimension PRE_EOS1 of the matrix UEQUATION3 , 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). (2)3 X (1) , continuing from Table 49 . 3 X (1) , continuing from Table 49 . 4 X (1) , continuing from Table 49 . The main objective of Attention head 5 is to fill the dimension OP1_SHIFT2 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 -1)-th least significant digit of the first operand. Similarly to the previous head, i ranges from 0 to ℓ a + 2, and in cases where the i-th least significant digit of the first operand is not well-defined (i.e., i ∈ {0, 1, ℓ a + 2}), we assign 0.As mentioned in Attention head 3, we assign an extra goal to Attention head 5, which is to fill the dimension PRE_EOS2.The design of the fifth head is as follows. With d QK,25 = P + 1, Q We provide the examples in Tables 70 to 74 . Note that within the dimension PRE_EOS2 of the matrix U(2) 5 V(2) 5 X (1) A(2)5 , if we restrict our view to the equal symbol = and the response sequence, 1 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_EOS1 and PRE_EOS2, the most significant digit of the response is the only token that has a value of 1 in both dimensions. This observation plays a crucial role in predicting EOS for the next token, and we will elaborate further in the later section discussing the FFN layer. 5 X (1) , continuing from Table 49 . The objective of Attention head 6 is to fill the dimension OP1_SHIFT3 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 -2)-th least significant digit of the first operand. Similarly to the previous head, i ranges from 0 to ℓ a + 2. In cases where the i-th least significant digit of the first operand is not well-defined (i.e., i ∈ {0, 1, 2}), we assign 0.The design of Attention head 6 is as follows. With d QK,26 = P + 1, Q = 2(e d NUM ) ⊤ ∈ R d V,26 ×d , (120)EQUATIONWe provide the examples in Tables 75 to 79 . 6 X (1) , continuing from Table 49 . 6 X (1) , continuing from Table 49 . 7 X (1) , continuing from Table 49 . (2) 7 X (1) , continuing from Table 49 So far we have computed the output of Att 2 operation. Passing through the residual connection, the output of the attention layer becomes the sum of X (1) (the input to the second Transformer block) and the output of Att 2 operation:EQUATIONh X (1) A(2)h .(126)A concrete example of the output of residual connection is presented in Table 85 .• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 a 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).The hyperparameter max_pos 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 it assigns a sequence of increasing integers which are generally not consecutive.Note that they match up to matrix transpose, which is due to the difference in the formulations.The attention matrices depicted in Figure7are square, lower-triangular (due to causal attention pattern), and row-stochastic (all entries are nonnegative and the sum of each row equals 1).github.com/google/flaxformergithub.com/McGill-NLP/length-generalizationBOS and EOS tokens do not need to be identical. We regard them as the same token just for the simplicity of the presentation.* * * * * * * * * 1/4 0 3ℓ + 4 * * * * * * * * * * * * 1/2 On the other hand, if [T 2] ij = 0, [R 2 ] ij > 0 and [R 2 ] ij ≤ e M (P-2) x j e M P + (j -x j )e M (P -2) = 


Decision: 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.

Reviews: review: paper_summary: paper_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.

The 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.

Position coupling is used in conjunction with a number of other tricks such as reversing the sum, zero padding, and using BOS/EOS.

The 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.

The results are also backed by theory, and interestingly the attention patterns predicted by theory are observed in the experiments.

Other than the addition task, Multiplication ($N\times2$) and a 2D task are considered in the experiments as well.
strength_weakness: strength_weakness: 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.
questions: questions: None.
limitations: limitations: See weaknesses.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.
- 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.
questions: 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?
- 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?
- 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?
- 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?
- 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: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. The significance is limited.  
  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.  
  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.  
  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).
questions: 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?
1. 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?
1. 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.

I'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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. **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.
3. **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.
questions: 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: limitations: As mentioned above.

Comments: title: Reviewer Response
comment: Dear Reviewer CVfs,

Thank you for taking the time to review our work and for providing such insightful and constructive feedback. We understand that you may have a busy schedule, but we wanted to follow up to ensure that our responses have sufficiently addressed your concerns. If you have any further questions or comments, we would be glad to hear them.
Thank you for taking the time to provide valuable feedback. We appreciate your suggestions on clarifying the 2D experiment, and will certainly consider these points in our future work. We also appreciate the reference you provided and will look into the potential connection.
Thank you for your feedback and for reconsidering the score. We are glad to hear that our response addressed your concerns and that you view the paper as above the acceptance threshold. We would also be happy to hear if you have any additional thoughts or suggestions.
Thank you for your feedback and for reconsidering the score. We will be sure to incorporate your suggestions in our next revision. If you have any further questions or comments, please feel free to share them.
I have read the detailed reply from the authors. The additional clarity and the details in the general response and the PDF have addressed my concerns. I have changed my score from a 6 to a 7. At this point, I think the paper is clearly above the acceptance threshold.
Thank you for your efforts and your detailed rebuttal, especially for coming up with how to apply position coupling to the parity task and showing successful results on it (and thanks for acknowledging the limitation of the method w.r.t. tasks like sorting and mode). 

I'd add that the 2D experiment could benefit from a bit more explanation and how position coupling is applied to it. I am happy with the rebuttal and like the paper, thus I'll maintain my score, and would just encourage the authors to see if they can extend their image experiment to other tasks with more structure, but regardless I find the work solid.

As a side note, the discussion on length generalization getting worse with deeper models seems to be somewhat connected to a concurrent work [1], it might be worth seeing if there's indeed any relation.

[1] https://arxiv.org/pdf/2402.04875
Thank you for your feedback. I will raise my score from 5 to 6.