Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count
Authors: Hanseul Cho
Jaeyoung Cha
Srinadh Bhojanapalli
Chulhee Yun
Venue: ICLR 2025 Poster (2025)
Keywords: Length Generalization
Transformers
Scratchpad
Position Coupling
Positional Encoding
Out-of-distribution Generalization
Arithmetic Tasks
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose combining scratchpad with position coupling, and demonstrate that Transformers can achieve length generalization in both operand length and count for addition problems.

Abstract: Transformers often struggle with *length generalization*, meaning they fail to generalize to sequences longer than those encountered during training. While arithmetic tasks are commonly used to study length generalization, certain tasks are considered notoriously difficult, e.g., multi-operand addition (requiring generalization over both the number of operands and their lengths) and multiplication (requiring generalization over both operand lengths). In this work, we achieve approximately 2–3× length generalization on both tasks, which is the first such achievement in arithmetic Transformers. We design task-specific scratchpads enabling the model to focus on a fixed number of tokens per each next-token prediction step, and apply multi-level versions of *Position Coupling* (Cho et al., 2024; McLeish et al., 2024) to let Transformers know the right position to attend to. On the theory side, we prove that a 1-layer Transformer using our method can solve multi-operand addition, up to operand length and operand count that are exponential in embedding dimension.

Published as a conference paper at ICLR 2025
ARITHMETIC TRANSFORMERS CAN LENGTH- GENERALIZE IN BOTH OPERAND LENGTH AND COUNT
Hanseul Cho, Jaeyoung Cha" Kim Jaechul Graduate School of AI, KAIST (jhs4015, cha jaeyoung @kai staistcoe st.ac.ck
Srinadh Bhojanapalli
Google Research
bsr nadh @ e com
Chulhee Yun
Kim Jaechul Graduate School of AI, KAIST chulhee, yun@kaist kr
ABSTRACT
Transformers often struggle with length generalization, meaning they fail to gener- alize to sequences longer than those encountered during training. While arithmetic tasks commonly used to study length generalization, certain tasks consid- ered notoriously difficult, e.g., multi-operand addition (requiring generalization over both the number of operands and their lengths) and multiplication (requiring generalization over both operand lengths). In this work, we achieve approximately 2-3x length generalization on both tasks, which the first such achievement arithmetic Transformers. We design task-specific scratchpads enabling the model to focus on fixed number of tokens per each next-token prediction step, and apply multi-level versions of Position Coupling (Cho et al., 2024; McI.esh et al., 2024) to let Transformers know the right position to attend to. On the theory side, we prove that I-layer Transformer using method can solve multi-operand addition, up to operand length and operand count that are exponential in embedding dimension."
NoPE
FIRE
Position Coupling
(-Scratchpad) (8 runs)
(-Scratchpad) (8 runs)
(-Scratchpad) (8 runs)
2
15 20 25
30
100
5
10
15
20
25
30
100
10
15
20
25
100
Longth of Each Operand
1.
Leneth of Lach Operand
Longth of Operand
80
80
10.
10
10
60
60
60
15-
15
15
20
20
20.
40
25-
25-
20
30
30.
30-
Number Operands
Number of Operands
Number of Operands
NoPE
FIRE
Bi-level Position Coupling (+Scratchpad) (8 runs)
(+Scratchpad) (8 runs)
(Sseratchpad) (8 runs)
30
100
20
100
100
Longth of Each Operand
Length Each Operand
Langth of Operand
80
80
10.
60
15.
20-
20
20
40
25.
25
30.
30.
30.
Number of Operands
Number of Operands
Number of Operands
Figure 1: Unlocking Length Generalization Multi-Operand Addition Task. We present median exact-match accuracies for 6-layer 8-head decoder-only Transformers trained on multi- operand additions of 2-10 operands, each having 1-10 digits (red boxes represent the scope of trained lengths). We compare three state-of-the-art position embedding (PE) methods for length generalization: NoPE (Kazemnejad al., 2023), FIRE (Li et al., 2024), and Position Coupling (Cho et al., 2024; McLeish et al., 2024). With a proper scratchpad enabling Transformers to do extrinsic multi-step reasoning (described Section 4), all three PE methods can extend their generalization scope (blue area of heatmaps). Remarkably. with our proposed multi-level Pasition Coupling with scratchpad, we achieve a significant length generalization superior all other methods.
Authors contributed equally to this paper. All experiments with code gi thub. htt.c/co//aaaeee Jc/ wtp.piiiiooooool11111

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Transformer- based language models (Vaswani et al., 2017) have become a comeestone of modern deep learning in recent years (Chowdhery et al., 2023; Gemini et al., 2023; OpenAI, 2023; Thoppilan et al., 2022). Despite their seemingly limitless capabilities, they often struggle with a critical limitation known as the length generalization problem, meaning that the model does not perform well on input sequences longer than those encountered during training (Anil et al., 2022; Deletang et al., 2023; Press et al., 2022; Wu et al., 2023; Zhang et al., 2023). Length generalization has recently dragged the attention of many researchers because of the following two aspects: (1) the failure in length generalization corroborates the models' fundamental limitation that they may not genuinely understand the task-solving algorithm but may rely on short-cut learning that is only applicable to sequences of trained lengths; (2) improving length generalization can automatically extend the applicability of the models in both memory-efficient and computation-efficenn way.
As manageable and intriguing test beds, arithmetic and algorithmic tasks are commonly used to study the capabilities (including length generalization) of Transformers (Cho et al., 2024; Fan et al., 2024; Kazemnejad et al., 2023; Kim et al., 2021; Lee et al., 2024; McLeish et al., 2024; Nogueira et al., 2021; Qian al., 2023; Sabbaghi et al., 2024; Zhou et al., 2024a;b). In this paper, we mainly focus on arithmetic tasks, specifically integer addition and multiplication. While humans can easily generalize longer examples of these tasks, recent works have shown that Transformers often struggle in length generalization, and various approaches (Cho et al., 2024; McLeish et al., 2024) have been proposed to help Transformers learn the true underlying mechanisms that solve addition and multiplication.
While recent work has made significant progress on the addition task, the scope has largely been limited to cases of two operands. Similarly, studies on multiplication have achieved length gener- alization for just one operand, while the other is kept fixed at small length. Notably, as far as we know, no research has demonstrated significant generalization in terms of the number of operands for the addition task; e.g., training on problems with up to four operands and successfully extending to problems with more than four. Likewise, for multiplication, none has achieved length generalization for both operands simultaneously (e.g., see Figure 5 of McLeish et al., 2024).
In this paper, we address these challenges by proposing a nontrivial combination of two techniques: scratchpad (Nye et al., 2021) and Position Coupling (Cho et al., 2024; McLeish et al., 2024). By equipping decoder-only Transformers with a scratchpad (an appended sequence that contains multi- step reasoning) and integrating a bi-level extension of Position Coupling, we demonstrate that Transformers can learn to solve multi-operand addition, generalizing in terms of both the number of operands and their lengths. Similarly, for multiplication, we employ a scratchpad and tri-level Position Coupling to train Transformers that Sintt--generalize in terms of both operand lengths.
Admittedly, the two key components-ssratchhaaal and Position Coupling-are not entirely new, and they have been adopted in existing approaches to improve length generalization in arithmetic tasks. Here, our key contribution is to combine them in a complementary way that empowers Transformers to solve the tasks that were considered very challenging. We propose novel form of scratchpads to eliminate the need attend to an increasing number of positions as the number of operands increases. By spelling out the intermediate outcomes on scratchpads, it now suffices for the model to attend to only a constant number of tokens at each inference step. Position Coupling then offers the model information about the "correct" position(s) to attend to, thereby assisting the model to quickly learn the correct inference mechanism from data.
1.1 SUMMARY OF CONTRIBUTIONS
We first tackle the multi-operand addition task (Section 4). By devising and employing scratchpad with bi-level Position Coupling, we achieve a significant length generalization not only in the length of each operand but also in greater numbers of operands. Our model substantially improves the generalization performance (with median exact-match accuracy 2 90.0%) for the integer addition task involving up to 30 operands of lengths up to 30, even though it was trained on samples with a maximum of 10 digits and 10 operands. In contrast, models trained with either NoPE (Kazemnejad et al., 2023) or FIRE (Li et al., 2024) completely fail to solve for 13 operands of length 13, even with the help of the same scratchpad (Figure l).
By refining the scratchpad and employing trielevel Position Coupling, we achieve a significant length generalization for multiplication tasks where both operands vary in length (Section 5).
2

---

Published as a conference paper at ICLR 2025
Our models are capable of multiplying (up to) 20-digit integers times (up to) 15-digit integers with median exact-match accuracy 2 78.55%, even though it was trained on samples with a maximum of 10 digits for each operand (Figure 9).
We develop a theoretical construction of a small (1-layer. 4-head) Transformer model, equipped with a scratchpad and bi-level Position Coupling. capable of solving multi-operand addition (Theorem 4.1). Our construction can handle problems involving both exponentially long operands and exponentially many operands. The scratchpad crucial, as it enables this shallow architecture to accurately predict the next tokens by ensuring that the model only needs to attend to constant number of tokens at each inference step, which we also verify in trained Transformers (Figure 7).
PRELIMINARIES
We use next-token prediction (NTP) with decoder-only Transformers to solve every task. Each task can be represented as set of sequences of the form " "[query] (response)]]]]] where the goal is to correctly infer the response sequence from a given query sequence via NTP, starting from a sequence " "[query] =", Since we are mostly studying length generalization on arithmetic tasks, we treat a single digit (between 0-9) as a single token, but there are other non-digit tokens such as '+', 'x' (or interchangeably **'), '=', '-" (or interchangeably '>'), and special tokens like laa-miinngeeddoo--ssuueeeeeeeeee
(BOS/EOS) and padding (PAD) tokens. 2 Moreover, because of the deterministic nature of arithmetic/algorithmm tasks, we only use greedy decoding for every NTP step. In the following subsections, we provide an explanation of the background underlying our approach. For additional related works on length generalization, we refer the readers to Appendix A.
2.1 RELATED WORKS
Position Embedding Methods for Length Generalization. Various position embedding (PE) methods have been explored to enhance the Transformers' length generalization capability. Kazem- nejad et al. (2023) claim that, in some downstream tasks, decoder-only model without PE (NoPE) can achieve a length generalization performance comparable those of widely-used PE techniques including ALiBi (Press et al., 2022), Rotary (Su et al., 2024), and T5's Relative PE (Raffel et al., 2020). However, is still a promising research direction to enhance length generalizability of Transformers with more appropriate choice of PEs (Jelassi al., 2023; Ruoss et al., 2023; Shen et al., 2023; Zhou et al., 2024b).
Position Coupling. Independent works by Cho et al. (2024) and McLeish et al. (2024) propose Position Coupling (also called Abacus), a structured position ID assignment estabistallsheed on a learned absolute PE (Gehring al., 2017)-177000 that captures the inherent positional symmetry of the target task. In this approach, tokens in an input sequence are grouped and each group of tokens is assigned sequence of consecutive integers task, the identical position IDs are assigned to the digits at the same significance across numbers in both the query (i.e., operands) and the response (or, answer). During training, the starting position ID is randomized to mitigate the problem of encountering unseen position IDs for longer sequences. At test time, the starting position ID) is arbitrarily fixed (e.g.. 1). Position Coupling demonstrates a remarkable length generalization performance several arithmetic and algorithmic tasks such as two-operand addition and N N-digit x 2-digit multiplication.
position IDs. For example, in the integer addition
Scratchpad. To enhance the reasoning capabilities of Transformer models, several heuristic-driven methods for data formatting have been introduced. One such technique scratchpad (Nye et al., 2021), an auxiliary intermediate sequence of tokens before arriving at the final answer. Training the model with a scratchpad allows the model to store and refer to intermediate task-solving states when predicting subsequent tokens. This approach has been shown to enhance both in-distribution and out-of- dit-of--iittrbutionn performance in tasks such as integer addition and code execution (Anil et al., 2022; Nye al., 2021). Zhou et al. (2024a) further validate these findings in the parity task, explaining that the effectiveness of the scratchpad lies in its ability to simplify next-token prediction.
3 WARM-UP: LENGTH GENERALIZATION ON PARITY TASK
Before moving on to our findings about addition and multiplication tasks, we begin with a warm-up example: the parity task. Given a binary sequence as a query, the goal of the parity task is to output
'If possible, we ignore the details about the BOS/EOS/PAD tokens for simple presentation.
3

---

Published as a conference paper at ICLR 2025
if the query contains odd number of I's, and 0 otherwise. Despite its simple description, it is well known that Transformers struggle with achieving length generalization for the parity task (Anil et al., 2022; Deletang al., 2023; Hahn & Rofin, 2024; Kazemnejad et al., 2023; Zhou et al., 2024a). In this section, we will demonstrate an enhanced length generalization performance on the parity task by applying Position Coupling on top of the input sequence with properly designed seratchpad.
3.1 METHOD: SCRATCHPAD & POSITION COUPLING
Intermacitd Shep
0
0
1
0
Simple Algortti
Figure 2: An illustration of scratchpad for a parity problem (with query 0101).
We follow the scratchpad format proposed by Anil et al. (2022). Figure 2 illustrates an example of an input sequence consisting of 4-bit parity problem (with query 0101) and its scratchpad (0110) 0). The idea of this scratchpad is to record every intermediate parity state as given binary query is processed starting from the leftmost token. That is, the k-th bit in the scratchpad is the parity of the subsequence containing the first & bits of the query. Then, the final rightmost token of the scratchpad the desired answer for the parity task. Recording the process of solving the parity task is beneficial in the following two points:
1. It makes the task-solving algorithm simpler. The first token in the scratchpad is just copy of the first token in the query sequence. Also, we do not need to directly solve the intermediate parity task once in order infer the k-th (k: 1) intermediate token in the scratchpad; instead, is enough to focus on the k-th token of the query and the (k- - 1)-th token in the scratchpad, and then sum them up modulo (see Figure 2).
2. It is straightforward to apply Position Coupling onto the input sequence with the scratchpad. The scratchpad generates a natural positional correspondence between the query and the response. Thus, we can assign the same position ID to the k-th query token and k-tth scratchpad token, for example, 23452 2.345 5 in the example depicted in Figure 2.
To up, such scratchpad simplifies the task by allowing the model perform step-by-stey reasoning without need of attending to linearly increasing number of query tokens until it generates the answer, while Position Coupling (applied
top of the scratchpad) can explicitly let the model know "where it should focus on" to perform every step of the reasoning. We thus can expect synergetic effect of such combination, and our experiments do align with this expectation.
3.2 EXPERIMENTS & DISCUSSION
Parity Taikk Answer-Token Acc
MIP
52
Scrvtripa
Stralooadd
FortsinCouning Senbehpod
2.
200
Quayy Lengi
Figure 3: Parity task. report the accuracies only for the answer token (i.e., the token before BOS) (light area: 95% confidence intervals). The gray region indicates the query lengths in our training data. A complete failure indicated by the accuracy ~50%: a random guess between and 1.
To test the efficacy of the combination of scratchpad and Position Coupling, we compare its perfor- mance against six other configurations: models trained with NoPE, RoPE, and FIRE, each tested with and without the scratchpad. The result is presented in Figure 3. Please refer to Table for detailed experimental details.
We observe that, without the scratchpad, all three PE methods generalize well to in-distribution samples but struggle with out-of-distribution queries. Their performance sharply drops to 50%
4

---

Published as a conference paper at ICLR 2025
accuracy for query lengths slightly exceeding the training samples, indicating no better than random guessing. When the scratchpad is employed, NoPE and FIRE demonstrate improved performance, achieving strong generalization up to 35-bit queries. However, when a scratchpad is combined with Position Coupling. the model demonstrates outstanding length extrapolation, achieving near-pes per-eeefect generalization queries up to 100 bits. From these observations, we conclude that for the parity task, (1) the scratchpad helps length generalization, albeit not significantly, and (2) the combination of the scratchpad with Position Coupling substantially boosts the model's length extrapolation capability.
4 LENGTH GENERALIZATION ON MULTI-OPERAND ADDITION TASK
In the previous section, we demonstrated the potential of integrating the seratchpad with Position Coupling for better length generalization. We now aim evaluate the effectiveness of this approach on a more challenging task. Specifically, we address the problem of achieving length generalization in the integer addition task, where operand length and count can both get longer at test time. Indeed, there has already been a length extrapolation result on additions with more than two operands: e.g.. the "triple addition" task presented by Cho et al. (2024). However, their approach still requires the number of operands to be fixed (e.g.. 3), leaving the challenge of generalizing the problem setting to varying operand counts as an open question. Here, we demonstrate that it can be overcome by employing scratchpad in conjunction with carefully designed multi-level Position Coupling.
4.1 METHOD: SCRATCHPAD BI-LEVEI. POSITION COUPLING
Token 0S740088006-00007755500111100110 PosID1: 4321492143213411341332222
PosID2: 011122233911122223333333
Figure 4: An example input sequence equipped with scratchpad and bi-level Position Coupling. The original example was "57+48+96=201"" the query is inside blue box and the response inside a red box. All numbers in the scratchpad (i.e., intermediate steps) are reversed; all numbers in the whole sequence are minimally zero- sero-padded to match their length.
Scratchpad for Multi-operand Addition. Similar to the parity task, we store the intermediate cumulative sums in the scratchpad. It makes the algorithm of solving the multi-operand addition task easier. This is because a model can obtain the k-th intermediate cumulative sum by adding exactly two numbers: the most recently generated ((k 1)) intermediate sum and the k-th number/operand in the query. As minor detail, we start the scratchpad with zeros in order to make the task-solving rule more clear and consistent."
Li-Level Position coupling. Now we motivate the usage of the bi-level Position Coupling: each token has two levels of position IDs, whose couplings happen only in each level.* As observed in prior works, adding two numbers can be successfully done by coupling the digits of the same significance in every number by assigning the identical position ID. The resulting position IDs (level 1) are as in Pos ID1 in Figure 4. However, this is not enough because we want to know which specific numbers we should add together, while the numbers are not very distinguishable solely with level-1 position IDs. Our mitigation is to add another level of position IDs that can distinguish properly between numbers. Combining the idea of Position Coupling again, we naturally couple two numbers of the same order in query and response: the resulting position IDs (level 2) is in PosID2 in Figure 4. In short, the model will choose which numbers it should add based on level-2 position IDs, and perform two-operand addition with the help of level-1 position IDs. Lastly, when we implement multi-level position IDs for each token, we use separate PE modules for different levels of position IDs to map them to PE vectors, and then add them all to the token embedding vector.
"This is not must. We empirically observed that directly starting the scratchpad with the (zero- pero-p reversed) first operand does not hurt the performance of moderate-sized models. We choose to start the scratchpads with all zeros for the sake of simplicity in our theoretical construction, later presented in Section 4.4. In fact, for multiplication (Section 5) we adopt a scratchpad that does not start with all zeros.
'It means that even if token has position ID) p at level and another token has the same position ID) p level 2, these two tokens are not necessarily coupled by these position IDs.
'It is worth mentioning that multi-level position ID has already heen studied (He et al., 2024; Zhang et al., 2024). The implementation detail is very different because their approaches involve relative PEs.
5

---

Published as a conference paper at ICLR 2025
Input Formats. There are additional design choices regarding the input sequence format: zero- padding and number reversing. We apply zero- prro-padiing to every number in both query and response to ensure that the length of every number is identical to the maximum possible length of the final answer, depending on the operand count. For example, if we add 11 operands of length at most have a length + 2 (because 99 x 11 1089), so we match the length
2, the final answer of the numbers to be + 2 with zero-padding, In addition, we reverse all numbers in the scratchpad (i.e., intermediate cumulative sums). This a quite natural choice since even humans do additions starting from the least significant digit (Lee et al., 2024; Nogueira et al., 2021).
4.2 EXPERIMENTAL SETUP
Given two integers S b, we denote by [a: b- := {a,a+1,..... .6} set of consecutive integers.
Data Sampling. For the sake of simplicity of notation, we denote a training set by SA (n, m). consisting of addition problems where each operand can have at most digits and the operand count can range from to m. The dataset consists of two equally sized chunks. In the first chunk, for each sample, the number of operands is uniformly sampled from [2: m). and the length of each operand is independently sampled from [1 rd]; this means that the operands within single sample can differ in length. Then, based on the chosen lengths (e.g., 4), each of the operands are uniformly randomly generated (e.g., from [1000 : 9999]). In the second chunk, for each sample, the number of operands is still uniformly sampled from [2 ml. but this time, the lengths of all operands are identical. The operand length is sampled from [1 n). and all operands are randomly chosen to be of the chosen length. We use SA (10, 10) with size 500,000 the baseline training set.
We also denote a test set by Ta (n, m), consisting of a single component. In each sample, both the number of operands and the length of each operand are fixed specifically at m and n, respectively. For model evaluation, we draw a 30 x 29 grid heatmap and assess the performance of the model on the test set (A ji) of size 1,000 for every entry (i,j) E [1 30) x [2 : 30).
Model and Training. Our baseline model is a 6-layer 8-head decoder-only Transformer with embedding dimension 1024 (with approximately 63M parameters), trained from scratch. We do not incorporate weight decay or dropout. Further details can be found in Table 2.
Random Offseet of Position IDs During Training. An important detail about the training procedure is that we randomly choose offsets (for each level of position IDs) and add them to every position ID in each training sample. This is to promote learning all the position embedding vectors as evenly as possible, and this training- time random assignment of position IDs is already used in prior works for similar reasons (Cho et al., 2024; McLeish et al., 2024; Ruoss al., 2023). As Cho et al. (2024) and McLeish et al. (2024) did, we pre-define the maximum position IDs (for each level) as hyperparameters, which naturally determine maximum testable ranges operand lengths and count. See Table 2 for our choice of the maximum position IDs.
4.3 EXPERIMENTAL RESULTS
Position Coupling and Scratchpad Together Enable Powerful Length Generalization. We trained models using 3 different PE methods- Coupling. NoPE, and FIRE-both with and without scratchpad. The implementation details of Position Coupling differ by the presence/absence of scratchpad: if we use scratchpads, we apply the bi-level Position Coupling explained in Section 4.1; if not, we use simple single-level Position Coupling that matches the position IDs for digits the same significance in all numbers, which is also used for "triple addition" task in Cho et al. (2024). The results are showcased in Figure 1. We measure the exact-match accuracy for correct inference of the whole response, including scratchpad if it is used. The top 3 heatmaps in the figure, which are the results without scratchpads, indicate that these models can only generalize to in-distribution samples and exhibit near-zero exact-match accuracy on out-of-distribution samples. Next, when either NoPE or FIRE is combined with the scratchpad, the models show a slight improvement in terms of generalizable operand lengths and counts. They barely solve problems involving 13 numbers, each 12 digits long. In contrast, the combination of the scratchpad and Position Coupling enables much stronger length generalization, achieving triviaivial accuracy even test samples involving 30 numbers, each with 30 digits. We emphasize that such problems are extremely difficult, as the model must accurately predict a total of 1023 consecutive tokens to solve the problem.

---

Published as a conference paper at ICLR 2025
Effect of Trained Length. We compare the models trained on different lengths: S7 (, 7), SA 10), and S131 13). The results are presented in Figure As expected, the model's ability generalize longer sequences improves the training data covers wider range of lengths.
Trained on 5x(7.7 7) runs)
Trained on 54/20, 10) (6 runs)
Trainnd on 5x113,13) 13))( (4 runs)
5
10
20 25
15 20
30
130
2
10
15
25
30
100
Opurand
10
60
Opeand Langth
60
opsond Lorgan
10
15-
15-
40
30
30
30
Figure 5: Comparison of training lengths in the integer addition task. We report exact-match accuracies (median over at least 4 runs) for the addition task. The red box indicates the training distribution: from the left plot, we trained on S S7 7), (10,, 10), and (1(13, 13).
Effect Zero-padding Figure exhibits the exact-match accuracies for models trained on input sequences with scratchpad without zero-padding both the query and the response. Al- though there is moderate degradation in overall performance, the models still generalize well with respect to an increased num- ber of operands. Also, they maintain reasonable accuracy for operand lengths below 25 digits. implies that, although zero- padding aids in enhancing length extrapolation capability, it is not an absolute necessity.
No Zero- Parro-Padinng Coupling
Londh Ea Eachepetee
Effect of Architecture. To study whether approach be applied across various depth/width configurations, we explore vary the number of layers and heads.
Nurber of Opecrds
the performance Specifically, we tested configurations with 1, 2, 4, and 6 layers, and 2, 4, and heads, resulting in 12 distinct configurations. The results are illustrated in Figure 10 (see Appendix C).
Figure 6: Training without zero- padding in input sequences.
turns out that I-layer 2-head models perform the worst. While they could generalize across operand counts, they immediately fail for longer digits. The remaining 1l configurations exhibit significantly better performance than the 1-layer 2-head models. It is noteworthy that networks with only four heads in total (1 (1-layer 4-head or 2-layer 2-head) show surprisingly high accuracy, even outperforming our baseline (6-layer 8-head). In particular, the 2-layer 2-head models achieve least 90.0% exact-match accuracy (median over 6 trials) across every combination (n, m)) E [1 30] [2: 30), whereas the 6-layer 8-head models achieve 67.8% or above. However, we do not observe an overall trend between the accuracy and the number of layers/heads, the number of layers, and vice versa. We suspect that this is due to the stochasticity of the training process.
the optimal number of heads varied depending
For broader ablation results including head/embedding dimensions, training set sizes, and input formats, refer the reader to Appendix C.
4.4 THEORETICAL ANALYSIS ON 1-LAYER TRANSFORMER
In this section, explain the success of approach by providing theoretical analysis in Theo- rem 4.1. Specifically, we construct a Transformer model (whose normalization layer is omitted for simplicity) that capable of solving the addition task involving both exponentially long operands and exponentially large number of operands when our approach is applied. Furthermore, the constructed model is I-layer 4-head Transformer, which supports the experimental results presented in Figure 10: the 1 1-layer 4 head model succeeds, but the 1 lllayer 2-head model fails.
Theorem 4.1. With a proper input format, scratchpad, and Position Coupling, there exists a 1-layer 4- head decoder-only Transformer that solves the multi-operand integer addition task involving up to m operands each with up to n digits. Here, a sufficient choice of the embedding dimension is d Ollog (m + 1) log, 1)).
7

---

Published as a conference paper at ICLR 2025
Theorem 4.1 says that llaayer 4-head model is enough to solve integer additions involving exponen- tially many and exponentially long operands (in embedding dimension) when a proper scratchpad and applied. As being sufficiency result, it implies that larger architectures (with more layers and attention heads) can solve the same task as well. We put its detailed and constructive proof in Appendix D. We also remark that our theorem extends Theorem 5.1 of Cho et al. (2024), which can only handle addition problems with fixed number of operands.
Position Coupling
To illustrate our key idea, consider an example problem 057+048-0096 000 750 501 102 (with the output reversed). First, consider the case without scratchpad: 057 + 048 + 096 = 102. To predict the least significant digit of the final answer, 1, the model must attend the least significant digits of all the operands, which are 7, 8, and 6, in the query sequence. In a scenario with m operands, the model would need to attend to m digits. This property-the number of tokens the model has to attend to increases with the number of operands-
the construction difficult.
With a scratchpad, the process becomes lot simpler. Instead of attending to every least significant digit of all operands, the model only needs to attend to two tokens: e.g., 6 from 096 and 5 from 501. Importantly, by utilizing the intermediate states stored in the scratchpad, the number of tokens the model needs to attend to remains fixed, even if the number of operands gets larger.
Scrutinizing the Attention Patterns of Trained Transformer. Surprisingly, our insight into the advantage of our scratchpad and its synergy with Position Coupling can be visually verified, supporting the significance of our method and theoretical finding. We probe the attention matrices of transformer models trained with a practical Adam optimizer. We visualize the lower-triangular row-stochastic matrix soft Q K as a heatmap in Figure 7. Thus, if you want to know the distribution of softmax logits over key tokens for an NTP the g-th query token, you should look g-th row of the heatmap; the brighter the point, the more the model attends to that key token position.
(Postion Coopling Scriichood
InSC. Scriicppat
(a) 1-layer 4-head model trained with bi-level Position Coupling and Scratchpad.
(b) 6-layer 8-head model trained with FIRE but no scratchpad.
Figure 7: Extracted attention matrices from certain attention heads of trained Transformers on an addition dataset (10, 10). We ran forward passes on 1,000 test samples from Ta 11) and obtained average attention matrices. The softmax values below 0.01 are clipped out (black cells). We compared with the model without Position Coupling nor scratchpad to demonstrate that, without scratchpad, model often try to "look at" every relevant positions at once, even without help of Position Coupling. We put more examples of attention patterns in Appendix
Now first look the attention pattern extracted from model trained with scratchpad and bi-level Position Coupling (Figure 7a). Since the scratchpad takes up about half of the total input sequence length, we may focus on the bottom half of the heatmap. The attention pattern tells us that, for most of the NTP step, each attention head focuses on at most fixed number of previous tokens: one on the short anti- diagona line (which corresponds the token in the query sequence) and one on the long diagonal line (which corresponds to the token in scratchpad). This observation is strikingly similar to our theoretical construction of the attention pattern.
On the other hand, let us move on to the attention pattern extracted from model trained with FIRE but without scratchpad (Figure 7b). In this case, the length of the response is the same scale as a single operand, so you may focus the last few rows of the heatmap. Unlike the previous case, and as we expected, the eleven anti-diagona lines the attention pattern reveal that the model actually focuses as many previous token positions the number of operands every time it performs NTP!
8

---

Published as a conference paper at ICLR 2025
5 LENGTH GENERALIZATION ON INTEGER MULTIPLICATION TASK
Achieving length generalization for multiplication when both multiplier and multiplicand can vary in length has long been recognized as a challenging problem. To our knowledge, prior works on the multiplication task (Duan & Shi, 2023; Fan al., 2024; Jelassi et al., 2023; McLeish et al., 2024) have never successfully addressed this issue. In this section, we demonstrate the combination of Position Coupling and the scratchpad can serve as a powerful solution to this obstacle.
5.1
METHOD: TWO-STAGE SCRATCHPAD & TRI-LEVEI. POSITION COUPLING
37
925
185
00185
Token 97998--8884400333355100333332223
074
00925
PosID1: $2000123422613340000000000000000
+333
34225
PosID2: 000211122333111122222333
34225
PosID3: 00000122443345445123441122223223
(a)
(b)
Figure 8: (a) A usual computation of integer multiplication, displaying the decomposition of multipli- cation task into two stages. (b) An example input sequence with a two-stage scratchpad and tri-level Position Coupling. The original example was "37x x925=34225""; the query is inside a blue box and the response is inside a red box. All numbers in the scratchpad are reversed; all numbers in the whole sequence are minimally aero-padded to match the lengths of numbers in the same stage.
Different from the addition tasks, every digit of first operand interacts with every digit of the second operand. This makes it difficult to come up with single-stage cumulative scratchpad. To achieve strong length generalization in multiplication, we take a step beyond the simple cumulative scratchpads. We propose two-stage scratchpad, motivated by observation that the usual (human) computation of integer multiplication can be decomposed into two stages: (i) a series of M -igitt 1-digit multiplications and (ii) (linearly shifted variant of) multi-operand addition. Figure 8b illustrates a concrete example of our two-stage scratchpad and tri-level position ID assignment rule, and Figure 8a explains the intuitive motivation of the scratchpad. Note that we concatenate two stages of scratchpad with a '=' token, which the model is required to infer as well as other digit/symbol tokens. For simplicity, let us write two operands A (with M digits) and B (with N digits).
Stage 1: M- 1erixx11ddd 1-digit Multiplications. The first stage of the scratchpad consists of N numbers: the first number is the product between A and the least significant digit (LSD) of B, the second number is the product between A and the second LSD of B, and so on. We reverse all the N numbers, zero-pad them to match the length, and concatenate them in order with '+' tokens in between. Regarding the position ID assignment, observe that it is natural to (i) couple k-th LSD of with k-th LSDs every number of the scratchpad stage and to (ii) couple the k-th LSD in B with every digit in the k-th number of the scratchpad stage 1. This is reflected to PosID1 and PosID2 in Figure 8b, until the second '=' token.
Stage 2: (Modified) Multi Oneri-O Addition. The second stage is basically the same a familiar multi-operand addition, with a slight difference that we shift the operands to the left one by one as add them sequentially. It can be done by viewing the LSD of the k-th number in stage (i.e., the leftmost digit, as it is already reversed) as the k-th LSD when solving stage 2. This is semantically equivalent converting "581+ +70++33 +333" in the stage I into "581 00+04706 +003333" Because of this, we introduce a totally new level of position ID (level-3) that reflects this change of viewpoint on digit significance. Luckily, we can reuse the level-2 position IDs since they only distinguish the numbers. As a result, we expect the model to utilize level-3 and 2 of position IDs to solve the second stage. this spirit, the position ID assignment rule is similar to that introduced in Section 4.1, reflected to Pos PosID2 and Pos ID3 in Figure 8b. Lastly, we fill in unnecessary slots with 0's.
5.2 EXPERIMENTAL SETUP
Data Sampling. We denote the training set by SM (n, m), consisting of multiplication problems where the first operand can have up to n digits and the second operand can have up to m digits. Specifically, the length of the first operand is uniformly selected from [1 n).. The first operand is then randomly generated based on the chosen length. The second operand is chosen in a similar way, but its length is selected from : m). We use SM (10, 10) the baseline training set.

---

Published as a conference paper at ICLR 2025
The test set (n, m) is constructed in similar manner, with the primary difference being that the lengths of both operands are strictly fixed at n and m. For evaluation, we create a 30 x 30 grid heatmap and assess the performance of the model the test set (i, for every entry (1,5) E [1 : 300 x [1 :3]--
Model and Training. We employ the same baseline architecture as in Section 4 (Refer Table 3).
5.3 EXPERIMENTAL RESULTS
We evaluate models trained with 3 different position embedding methods, both with and without the scratchpad, and present the results in Figure For models using Position Coupling without the scratchpad, we adopted similar position ID assignment scheme proposed for solving N- digit x 2-digit multiplication in Cho et al. (2024). When the scratchpad is not applied, which corresponds to the top 3 plots, none of the models manage to generalize, even on in-distribution samples. When NoPE or FIRE is employed limited length generalization, achieving non-tri trivial accuracy on multiplication between two 12-digit integers. However, the combination of Position Coupling and the scratchpad again dominates others.
conjunction with the scratchpad, the models show
Interestingly, Position Coupling without the scratchpad shows weak length generalization when one of the operands is short. This should not come as a surprise, as Cho et al. (2024) already demonstrate that model trained with Position Coupling alone can generalize to N-digit x 2-digit multiplication task: models trained on N < 40 can generalize to N 2 100.
NoPE (-scratchpad) (4 runs) 5 10 15 20 25 30
FIRE (-scratchpad) (4 runs)
Position Coupling
(-scratchpad) (4 runs)
10 15 20 25 30
100
5
10
15
20
25 30
80
80
10.
10
10
60
15.
60
Longth of Fist
60
AC
Lengte
20
25
30
30
30-
of Operand
Length Operand
Length
FIRE (+scratchpad) runs)
Tri-level Position Coupling (+scratchpad) (4 runs)
(+scratchpad) (4 runs)
5 10
25 30
10
29 25 30
5
10
100
100
Lengee Fis Operind
100
80
80
10
10
15-
15:
15-
20.
20
20
40
25
20
25
-20
20
of Second Operand
Len Mengthol Second Operand
Length of Second Ope and
Figure 9: Comparison of methods in the multiplication task. We report exact-match accuracies for the multiplication task, with results taken as the median over 4 seeds for each method. Each axis represents the length of each operand. The red box indicates the training distribution, SM (10, 10).
CONCLUSION
While length generalization in arithmetic Transformers has drawn a lot of attention, especially on operand length for the addition, the ability to generalize on operand counts is considered a difficult challenge and has not been explored yet. To address this challenge, we propose a combination of two techniques: scratchpad and Position Coupling. We show that a Transformer trained on problems involving 1-10 digit integers with 1-10 operands can solve addition tasks with up to 30 operands, each being as long as 30 digits. We also theoretically construct a 1-layer Transformer model capable of adding exponentially many operands with exponentially long integers when our approach is applied. Finally, we demonstrate the effectiveness of our approach for length generalization in the multiplication task, where both operand lengths can vary.
Limitation. One limitation of our approach is that it is only applicable to tasks whose structure is well-defined and can be effectively encoded by scratchpad and Position Coupling. This leaves us with two directions for future work. The first direction is to establish a clear principle for employing scratchpad and Position Coupling when the task structure is known, as the current design choice heavily relies on intuition. The second is to extend our method to the cases where the task structure is implicit even entirely unknown.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors sincerely thank Pranjal Awasthi and Anupam Gupta for their insightful discussions. This work was partly supported by a National Research Foundation of Korea (NRF) grant (No. RS-2024-004212033 and an Institute for Information communications Technology Planning Evaluation (IITP) grant (No. (KAIST)) funded by the Korean government (MSIT). HC, JC, and CY acknowledge support from a Google Gift on the research related to Long Context Transformers. The experiments contained in this work were supported part through a Google Cloud Platform Credit Award.
Artificial Intelligence Graduate School Program
REFERENCES
Cem Amil, Yuhuai Wu. Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models, Advances in Neural Information Processing Systems, 35:38546-385555 2022. 2, 3, 4, 16
Pranjal Awasthi and Anupam Gupta. Improving length-generrlizatton in transformers via task hinting. arXiv preprint arXhv:231000072) 2023. 16
Hanseul Cho, Jaeyoung Cha, Pranjal Awasthi, Srinadh Bhojanapalli, Anupam Gupta, and Chulhee Yun. Position coupling: Improving length generalization of arithmetic transformers using task structure. In Advances in Neural Information Processing Systems, volume 37, 2024. 1, 1, 3, 10, 16, 18, 24, 30, 37, 38
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023. 2
Gregoire Deletang, Anian Ruoss. Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, and Pedro A Ortega. Neural networks and the chomsky hierarchy. In The Eleventh International Conference on Learning Representations, 2023. URL https: httpttss:/ http:///ppeeeeeeeeee/e /ppeeeteep//eeeeeeeeeee... wtt.://eeeouummttooooooooeeeeee.
2, 4, 16
Shaoxiong Duan and Yining Shi. From interpolation to extrapolation: Complete length generalization for arithmetic transformers. arXiv preprint arXiv:2310.11995 2023.
Ying Fan, Yilun Du, Kannan Ramchandran, and Kangwook Lee. Looped transformers for length generalization. arXiv preprini arXiv:2409. 2024. 2, 2, 16
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. Advances in Neural Information Processing Systems, 36, 2023. 17
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann Dauphin. Convolutional sequence sequence learning. In International conference machine learning, pp. 1243-1252 PMLR, 2017. 3
Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, CaantBaptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023. 2
Michael Hahn and Mark Rofin. Why are sensitive functions hard for transformers?
Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 14973-15555 15008, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.180 (0.1863//v12024 acl-long.800. URL httph::t////att www.ccannnoooool..oooooooooooooe wwtoorg////0 2024. acl-long 800. 4
Zhenyu He, Guhao Feng, Shengjie Luo, Kai Yang. Liwei Wang, Jingjing Xu, Zhi Zhang, Hongxia Yang, and Di He. Two stones hit bird: Bilevel positional encoding for better length extrapolation. In Forty-first International Conference on Machine Learning, 2024. URL https://openrre net hot///ormmm//// ferrum?id=lugf11e11...

---

Published as a conference paper at ICLR 2025
Samy Jelassi, Stéphane d' d'Ascoli, Carles Domingo- Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 2023. 3, 16
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy, The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2023. 1, 2, 3, 4, 16, 18
Jeonghwan Kim, Giwon Hong, Kyung-min Kim, Junmo Kang, and Sung-Hyon Myaeng. Have you seen that number? investigating extrapolation in question answering models. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7031-7037, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. 18, 19
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35: 22199-22213 2022. 17
Nayoung Lee, Kartik Sreenivasan, Jason D. Lec, Kangwook Lec, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024. URL WttLhttsp///peee /ttp///wppeeeeeeeeeeeeeeeeeeeeee net. / forum? diddddBdddd--e/ 4bst 9S. 2,
Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang. Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. In The Twelfth International Confer- ence on Learning Representations, 2024. URL attps:///ott /oppr////e//eeeee.///meom net / forum? ?id= rRC3qFesqk. 2, 16
Scan McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, and Tom Goldstein. Transformers can do arithmetic with the right embeddings. In Advances in Neural Information Processing
Systems, volume 37, 2024. 1, 3, 6, 9, 16
Rodrigo Nogucira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2102.13019 2021. 2, 6, 16
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint 1rx1v:2112.00114, 2021. 2,
OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.0777 2023. 2
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. 16, 18
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL https: wpwnreviewwweeee// net/ sateammmdd--sspPppoee
2, 16
Jing Qian, Hong Wang, Zekun Li, Shiyang Li, and Xifeng Yan. Limitations of language mod- els in arithmetic and symbolic induction. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 9285-9298 Toronto, Canada, July 2023. Asso- ciation for Computational Linguistics. doi: 10. 1/033//1/22.aaclllo/o //aclanthee ogy. org 2023 acl- 51
URL https:
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): 2020. 3, 16, 18
12

---

Published as a conference paper at ICLR 2025
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint arXiv:2305.1684. 2023. 3, 6. 16
Mahdi Sabbaghi, George Pappas, Hamed Hassani, and Surbhi Goel. Explicitly encoding structural symmetry is key to length generalization in arithmetic tasks. arXiv preprint arXiv:2406. 0855, 2024. 2
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.052002, 2020. 18, 19
Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lce, Yuanzhi Li, and Yi Zhang. Positional description matters for transformers arithmetic. arXiv preprint 1XXx:2331.1477:5 2023. 3, 16
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024. 3, 16
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc Le, Ed Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them, arXiv preprint arXiv:2210.099661 2022. 17
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Reng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.082 082399 2022. 2
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 16, 18
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022. 17
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-ar natural language processing. arXiv preprint arXiv:1910.03771, 2019. 18
Zhaofeng Wu, Linlu Qiu, Alexis Ross, Ekin Akyürek, Boyuan Chen, Bailin Wang, Najoung Kim. Jacob Andreas, and Yoon Kim. Reasoning or reciting? exploring the capabilities and limitations of language models through counterfactual tasks. arXiv preprint arXiv:2307.02477 2023. 2
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533 PMLR, 2020. 18
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019. 18, 19
Kechi Zhang, Ge Li, Huangzhao Zhang, and Zhi Jin. Hirope: Length extrapolation for code models using hierarchical position. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 13615-13627, 2024. 5
Yi Zhang, Arturs Backurs, Sebastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wag- ner. Unveiling transformers with LEGO: A synthetic reasoning task, 2023. URL https: //openreview nwt./.ee//eeeee wttp:/iwwweeeeetteeeeemeeoem foo.mmiddd1o jor-rmmm-dd--1j--f--fr.r. 2
Hattie Zhou, Arwen Bradley, Etai Littwin Noam Razin, Omid Saremi, Joshua M. Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization. In The Twelfth International Conference on Learning Representations, 2024a. URL https: http:/s////eeeevieeeeee net/ /nrt/foruume forum? Assss Asd.ssssSSSSSH......... THHnHXX. 2, 4, 16
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. ICLR Workshop on Understanding of Foundation Models FoMo), 2024b. 2, 3, 16, 18
13

---

Published as a conference paper at ICLR 2025
CONTENTS
1 INTRODUCTION
2
1.1 SUMMARY OF CONTRIBUTIONS
2 PRELIMINARIES
2.1 RELATED WORKS
3 WARM-UP: LENGTH GENERALIZATION PARITY TASK
3.1 METHOD: SCRATCHPAD & POSITION COUPLING
3.2
EXPERIMENTS & DISCUSSION
4
LENGTH GENERALIZATION ON
ADDITION TASK
4.1
METHOD: SCRATCHPAD & BI-LEVEL POSITION COUPLING
4.2
EXPERIMENTAL SETUP
6
4.3
EXPERIMENTAL RESULTS
4.4
THEORETICAL ANALYSIS ON 1-LAYER TRANSFORMER
LENGTH GENERALIZATION ON INTEGER MULTIPLICATION TASK
5.1
METHOD: TWO-STAGE SCRATCHPAD & TRI-LEVEL POSITION COUPLING
5.2 EXPERIMENTAL SETUP
5.3
EXPERIMENTAL RESULTT
10
6 CONCLUSION
10
A ADDITIONAL RELATED WORKS
16
A.1 LENGTH GENERALIZATION IN ARITEMETII/AL.GORTHMMC TRANSFORMERS
16
A.2 CHAIG-OF-THOUGHTS PROMPTING
17
B EXPERIMENTAL DETAILS
18
C ADDITIONAL EXPERIMENTAL RESULTS
20
D FORMAL CONSTRUCTION OF MULTI-OPERAND ADDITION TRANSFORMER
24
D.I NOTATION
24
D.2 ARCHITECTURE
24
D.3 INPUT SEQUENCE
24
D.4 ENCODING FUNCTION
25
D.4.1 TOKEN EMBEDDING
25
D.4.2 COUPLED POSITION IDS AND POSITION EMBEDDING
25
D.5 TRANSFORMER BLOCK: CAUSAL ATTENTION LAYER
27
D.5.1 ATTENTION HEAD 1: DIGIT-WIST ADDITION WITHOUT CARRIES
28
14

---

Published as a conference paper at ICLR 2025
D.5.2 ATTENTION HEAD 2: CARRY DETECTION
30
D.5.3 ATTENTION HEAD 3: ARROW DETECTION
32
D.5.4 ATTENTION HEAD 4: EOS DETECTION
34
D.5.5 RESIDUAL CONNECTION
36
D.6 TRANSFORMER BLOCK: TOKEN- TOKEN-WISE FEED-FORWARD LAYER
37
D.6.1 SUBNETWORK 1: CONSTRUCTION FOR SUM (DIMENSION 8-17)
37
D.6.2 SUBNETWORK 2: CONSTRUCTION FOR ARROW (DIMENSION 18)
38
D.6.3 SUBNETWORK 3: CONSTRUCTION FOR EOS (DIMENSION 19)
38
D.6.4 RESIDUAL CONNECTION
39
D.7 DECODING FUNCTION
39
E MORE ATTENTION PATTERNS OF TRAINED TRANSFORMERS
41
E.1 ATTENTION PATTERNS WITH SCRATCHPAD
E.2 ATTENTION PATTERNS WITHOUT SCRATCHPAD
43
15

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORKS
A.1 LENGTH GENERALIZATION IN ARITHMETIC/ALGORTHEGG TRANSFORMERS
Several works (Anil et al., 2022; Deletang et al., 2023; Nogueira et al., 2021) have scrutinized the Transformer's lack of ability to length generalize across algorithmic reasoning tasks. Here, we list the studies investigating and enhancing the length extrapolation capability of Transformers for arthhetii/algoorihmii tasks.
Addition Tasks. We first begin by exploring studies focused on encoder-only Transformers. Ruoss et al. (2023) introduce Randomized Position Encodings to address the problem of the appearance of unseen position indices longer sequences. Additionally, Jelassi et al. (2023) demonstrate that the RPE method enables generalization to 15-digit addition problems when the model is trained on problems up to digits.
We next move to the studies considering decoder-only Transformers. Kazemnejad et al. (2023)
investigate the effect of PE methods on length generalization performance, arguing that NoPE outperforms several popular PE methods such as ALiBi (Press et al., 2022), RoPE (Su et al., 2024),
APE (Vaswani et al., 2017), and T5's Relative PE (Raffel et al., 2020). Meanwhile, Shen et al. (2023)
propose the use of a scratchpad and random spacing, which facilitate generalization to 12-digit problems when trained on up to 10-digit problems. Zhou et al. (2024a) introduce the technique of "index-hinting", which inserts appropriate position markers in the sequence. Zhou et al. (2024b) integrate several existing techniques- FFRR (Li et al., 2024), index-binting (Zhou et al., 2024a), and Randomized PE (Ruoss et al., achievin generalization 100-digit problems while training exclusively on samples with fewer than 40 digits. Furthermore, Cho et al. (2024) and McLeish et al. (2024) independently introduce Position Coupling (also called "Abacus"), demonstrating state-of- the-art performance in the literature by generalizing to 100-digit problems after training on samples with up to 20 digits. We lastly remark there a recent attempt to solve addition by utilizing a looped transformer (Fan et al., 2024).
Multiplication Tasks. Most studies on multiplication focus on problems where one operand has fixed digit length. Jelassi et al. (2023) investigate N N-digitx3-dight multiplication but observe length generalization only when train set priming is applied, which involves adding few long samples in the train set. Cho et al. (2024) present the effectiveness of Position Coupling in 7-digitx2-diig multiplication, achieving generalization N > 100 after training on samples with N < 40. Fan et al. (2024) showcase the length generalization capability of looped Transformers to 16-digit problems from training up 11 digits, where the numbers are encoded in binary format and the length of the first operand is up to 2.
McLeish al. (2024) also train their models on multiplication tasks where both operand lengths can vary. They employ the "Abacus" embedding without scratchpud, and report near- sear-perfect in- distribution accuracy but near-zero accuracy on cutoof-distribution multiplication problems. A notable point is the success on in-distribution samples, whereas our results under a similar setting, shown in Figure 9 (upper rightmost plot), exhibit completely different trend. We suspect that this discrepancy arises from differences in their architectural implementation details' compared to ours. One such major difference which we suspect as the main cause of the performance discrepancy is the implementation of attention (sub)layers. They used torch. nn. Mu het.eettttteee tent i oo module predefined inside PyTorch (Paszke et al., 2019), whereas our code includes an explicit implementation of transformer attention layer with bunch of tensor multiplications (specifically based on part of HuggingFace Face's T5 model implementation). We list a few other notable implementational details that exclusively included in their codebase: the recurrent layers (i.e., looped Transformers), input injection technique, noise injection into the input embedding, and additional feed-forwan layer between the last transformer block and the linear read-out layer.
There is also a body of literature focused on length generalization in algorithmic tasks. We highlight few key contributions. Zhou et al. (2024a) introduce the concept of RASP-L, suggesting the conjecture of the Transformer's 's ability to fngth-generaalie may depend on whether the task can be expressed in the RASP-L language. Additionally, Awasthi & Gupta (2023) create an auxiliary task
i wtt.hb.ccmmcleeeet... th
16

---

Published as a conference paper at ICLR 2025
associated with sorting, resulting in substantial length generalization improvements through multitask learning.
A.2 CHAIN-OF-THOUGHT PROMPTING
While the main focus of this work is on training the model with scratchpad, frauin-of-thuughhe (CoT) prompting showing series of intermediate natural language reasoning steps before reaching the answer- is also extensively studied to enhance the reasoning ability of Transformers (Kojima et al., 2022; Suzgun et al., 2022; Wei et al., 2022). Similar to the spirit of scratchpad, CoT allows the model decompose complex problems into several intermediate steps, which has demonstrated its importance in tasks that require arithmetic and reasoning.
In an attempt to understand the success of CoT, Feng et al. (2023) investigate the expressivity of CoT. They prove that the autoregressive Transformers of constant size can solve basic arithmetic/equation tasks, while finite-depth Transformer models cannot directly produce correct answers to these tasks unless their size grows super-polynonee reveal that models trained with CoT-formatted data can generalize to different numbers of operands, but the generalization leap is limited to 3 (from 15 to 18).
with input length. In arithmetic tasks, their experiments
17

---

Published as a conference paper at ICLR 2025
B EXPERIMENTAL DETAILS
We modify and customize the codebase from Kazemnejad et al. (2023) for all our experiments." ? This codebase includes custom implementation of a decoder-only T5 model (Raffel al., 2020) built upon Py Torch (Paszke et al., 2019) and HuggingFace (Wolf et al., 2019), which incorporates several positional encoding methods.
We implemented custom RMSNorm module (Zhang & Sennrich, 2019) and various normalization layer positioning schemes (c.g., PreNorm (Xiong et al., 2020), PostNorm (Vaswani et al., 2017)) to follow the implementation details outlined by Cho et al. (2024); Zhou et al. (2024b).
Below, we provide detailed settings experiments.
Table I: Hyperparameter summary for the parity task (Figure 3).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
6
Number of Attention Heads
8
Embedding Dimension
512
Dimension per Head
64
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU (Shazeer, 2020)
Normalization Layer
RMSNorm (Zhang & Sennrich, 2019)
Normalization Layer Position
PreNorm and PostNorm
Trainable Parameter Count
25M
Training Steps
50,000
Batch Size
20
Optimizer
Adam (Kingma & Ba, 2015)
Learning Rate (LR)
0.00003
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0.1LR)
LR Cool-down
Maximum Position ID (max_pos)
101
Training Dataset Size
10,000
Evaluation Dataset Size
10,000 per query length
Device
NVIDIA RTX A6000 48GB
Training Time
< 3 hours
i wtw.b.cc.mccct....... hwtp./c//wcii11ii///eeeeeeeeeeee htgh:hob..cc///////1//l//g/ggeee iationn
18

---

Published as a conference paper at ICLR 2025
Table 2: Hyperparameter summary for the addition task (Figure 1).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
6
Number of Attention Heads
8
Embedding Dimension
1024
Dimension per Head
128
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU (Shazeer, 2020)
Normalization Layer
RMSNorm (Zhang Sennrich, 2019)
Normalization Layer Position
PreNorm and PostNorm
Trainable Parameter Count
63M
Training Steps
50,000
Batch Size
400
Optimizer
Adam (Kingma & Ba, 2015)
Learning Rate (LR)
0.00003
LR Warm-up
Linear (From 0 to LR), 1% of total steps
LR Cool-down
Cosine Decay (From LR to 0.ILR)
Maximum Position ID 1 (max_ppo___
40
Maximum Position ID 2
40
Training Dataset Size
500,000
Evaluation Dataset Size
1,000 per operand length/count
Device
NVIDIA RTX A6000 48GB
Training Time
< 12 hours
Table Hyperparameter summary for the multiplication task (Figure 9).
Hyperparameter
Value
Architecture
Decoder-only Transformer
Number of Layers
6
Number of Attention Heads
8
Embedding Dimension
1024
Dimension per Head
128
Hidden Width of Feed-forward Layer
2048
Activation Function of Feed-forward Layer
GEGLU (Shazeer, 2020)
Normalization Layer
RMSNorm (Zhang Sennrich, 2019)
Normalization Layer Position
PreNorm and PostNorm
Trainable Parameter Count
63M
Training Steps
50,000
Batch Size
200
Optimizer
Adam (Kingma & Ba, 2015)
Learning Rate (LR)
0.00005
LR Warm-up
Linear (From 0 to LR), 1% of total steps Cosine Decay (From LR to 0.ILR)
LR Cool-down
Maximum Position ID 1 (max Becax_pos_________
64
Maximum Position ID 2 (max pos 2)
32
Maximum Position ID) 3 (max_pos_3)
64
Training Dataset Size
500000
Evaluation Dataset Size
1000 per length combination of first/second operands
Device
NVIDIA RTX A6000 48GB
Training Time
12 hours
19

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENTAL RESULTS
In this section, provide additional experimental results not discussed in the main section.
first present the results experiments investigating how changes in model architecture (the number of layers and attention heads) affect task performance. We explore 12 different configurations. The heatmaps below represent the performance across different operand numbers and lengths, with blue regions indicating higher accuracy and red regions indicating lower accuracy.
1L2H (6 runs)
144H (6 runs)
1LSH (6 runs)
170
500
30
Natoer of Openrdi
Sumer of Oownrcc runs)
Sumber Operrcs
runs)
Corred
4L2H runs)
4L4H
4L8H runs)
6L2H (6 runs)
6L4H (6 runs)
6LSH (6 runs)
100
300
500
28
Narbar of Openral
Sunter of DOMMIIc
Figure 10: Comparison architectures the integer addition task. We report exact-match accuracies for the addition task. with results taken as the median over 6 seeds for each method. The x-axis corresponds to the number of operands, and the y-axis indicates the length of operands. The red box indicates the trained scope S (10, 10).
20

---

Published as a conference paper at ICLR 2025
One might be concerned why IL8H models perform worse than 1L4H models. To address this concern, we conduct experiments by controlling the total number of head dimensions. In Figure 10, since we fix the total number of head dimensions within layer by 1024, each head in IL4H models has 256 dimensions while each head in 1L8H models has only 128 dimensions. To make a fair comparison, decoupled the embedding dimension from the head dimension, fixing the embedding dimension by 1024 for this setup only. The results are presented in Figure 11. We observe that the performance of 1L8H models improves when the dimension per head is increased to 256, while the performance of 1L4H models degrades when the dimension per head is decreased to 128. We conclude that the inferior performance of IL8H models in Figure 10 is due the reduced dimension per head. These findings suggest that larger dimension per head is crucial for achieving strong performance in shallow models.
Exact-Match Acc
1L4H1024D (128d/h) (8 runs)
100
1L4H1024D (256d/h) (8 runs)
100
5
15
25
5
15
25
1-
Length of Each Operand
5
80
Length of Each Operand
5
80
10
10
60
60
15-
15-
20
40
20
40
25
25
20
20
30
30
Number of Operands
Number of Operands
1L8H1024D (128d/h) (8 runs)
100
1L8H1024D (256d/h) (8 runs)
100
5
15
25
5
15
25
1-
Length of Each Operand
5-
80
Length of Each Operand
5-
80
10
10-
60
60
15-
15-
20
40
20
-40
25
25
20
30
30
Number of Operands
Number of Operands
Figure 11: Comparison of head dimensions in the integer addition task. We report exact-match accuracies for the addition task, with results taken as the median over 8 seeds for each configuration. The x-axis corresponds to the number of operands, and the y-axis indicates the length of operands. The red box indicates the trained scope S (10, 10).
Next, we provide ablation results on varying embedding dimensions. We trained 1L4H models, each with embedding dimensions of 64, 128, 256, and 512 (with dimensions per head of 16, 32, 64, and 128, respectively). The results are illustrated in Figure 12. We observe that there a significant performance gap between these configurations. While models with small embedding dimensions have sufficient expressive capacity for solving the task (Theorem 4. 1), we believe larger embedding dimensions are crucial for enabling more effective optimization.
We now present the results of an ablation study on the training set size. Our baseline training data size is 500K, and vary the size to 20K. 100K, 2M, and 10M. We fix the architecture as 6-layer 8-head model with an embedding dimension of 1024. The results are illustrated in Figure 13. We observe that the training set size has significant impact on the model's performance.
21

---

Published as a conference paper at ICLR 2025
Exact-Match Acc
1L4H64D runs)
1L4H128D (32d/h)
runs]
2 5
10
15
20
25
30
2
5
10
15
20
25
30
2
5
10
15
20
25
30
Length of Each Operand
80
Length of Each Operand
80
Length of Each Operand
80
10
10
10-
60
60
60
154
20
20.
20.
25-
25
20
20
20
30
30
30
Number of Operands
Number of Operands
Number of Operands
1L4H512D (128d/h) (8 runs)
100
1L4H1024D (256d/h) runs)
100
5
15
25
5
15
25
Length of Each Operand
5-
80
Length of Each Operand
80
15-
15
20
20
25
20
20
30-
30-
Number of Operands
Number of Operands
Figure 12: Comparison of embedding dimensions in the integer addition task. We report exact-match accuracies for the addition task, with results taken as the median over seeds for each configuration. The x-axis corresponds to the number of operands, and the y-axis indicates the length of operands. The red box indicates the trained scope S (10, 10).
Exact-Match Acc
|=OK (8 runs)
[Pran 100K (8 runs)
[Pyan 5500 (8 runs)
100
100
100
15
Length of Lach Operand
5-
80
Length Each Operand
5-
Langth of Each Operand
5-
15.
15.
15.
20.
40
20.
40
40
25.
25.
25.
30
Number of Operands
Number of Operands
Number of Operands
1 runs)
[Dran] [Dran|=10M runs)
100
100
2
5
10
15
20
25
30
10
15
20
25
30
Length of Each Each Operand
Length of Each Each Operand
10--
10-
66
60
15-
15.
20.
40
20.
40
25-
25
-20
30
Number Operands
Number Operands
Figure 13: Comparison training set size the integer addition task. We report exact-match accuracies for the addition task, with results taken the median over 8 seeds for each configuration The x-axis corresponds to the number of operands, and the y-axis indicates the length of operands. The red box indicates the trained scope S (10, 10).
22

---

Published as a conference paper at ICLR 2025
We lastly investigate the impact of input sequence formatting. We considered three factors: (I) reversed or plain query, (2) reversed or plain response, and (3) zero-pan zero-paadiig pro--paddin or no zero zero-padees aro-padding. We exclude the case where reversed query and plain response are used together, resulting in plain query, reversed response, 6-layer 8-head model with embedding dimension
total of configurations. Note that our baseline format consists and zero-padding. We fix the architecture of 1024. The results are illustrated in Figure 14. To clarify, the top 3 figures correspond to the zero-padding setting, while the bottom 3 figures correspond the no zero-paddinn prro-padding setting, From left to right, the figures represent: plain query with reversed response, reversed query with reversed response, and plain query with plain response.
Exact-Match Acc
ReverseResponse runs
ZeroPad
ReverseAllM runs)
NoReverse runs)
2
5
10
15
20
25
30
2
5
10
15
20
25
30
2
5
10
15
20
25
30
Length of Each Operand
80
Length Each Operand
80
Longh of Exch Ogerand
80
10
10
10
60
60
60
15-
20
20
20.
25-
25
20
20
20
30
30
30
Number of Operands
Number of Operands
Number of Operands
NoPad ReverseResponse (8 runs)
100
NoPad ReverseAll (8 runs)
100
NoPad NoReverse (8
15
25
15
25
15
25
Length of Each Operard
60
Length of Each Operand
80
Longh of Each Operand
5-
10
10
10
15
15
15.
20
20
20
25
20
20
20
30-
30
30-
Number of Operands
Number of Operands
Number of Operands
Figure 14: Comparison of the input formats in the integer addition task. We report exact-match accuracies for the addition task, with results taken as the median over seeds for each configuration. The x-axis corresponds to the number of operands, and the y-axis indicates the length of operands. The red box indicates the trained scope S (10, 10).
23

---

Published as a conference paper at ICLR 2025
D FORMAL CONSTRUCTION OF MULTI-OPERAND ADDITION TRANSFORMER
In this section, we prove Theorem 4.1 1 by formally constructing a 1-layer 4-head Transformer model capable of solving multi-operand addition problems. The proof framework mostly follows the proof idea of Theorem 5.1 in Cho et al. (2024). For the sake of readability, we restate our theorem statement.
Theorem 4.1. With a proper input format, scratchpad, and Position Coupling, there exists I -layyer 4- head decoder-only Transformer that solves the multi-operand integer addition task involving up to m operands each with up to n digits. Here, a sufficient choice of the embedding dimension is d (logg (m + log2(n (n 1)).
D.1 NOTATION
Let ef represent the i-th standard basis vector of Rd Define In as the identity matrix of size m m. The vectors o and 1 are p-dimensional vectors filled with zeros and ones, respectively. Let Omxn denote the m x zero matrix. For N. we use [n] to denote the set [1...., For any matrix A, we use A.o and A.j to refer to the i-th row and j-th column of A, respectively.
Define an ordered vocabulary V (0, 1. 3, 4, 5, 6, 7, 8, +.=, =, -, represents both the BOS and EOS tokens do not have to be identical, we use single symbol for simplicity, Let V represent the k-th element of V.
-, 8). The special token 's' (BOS) token and the end-of-sequence (EOS) token. While
D.2 ARCHITECTURE
We adopt the same architecture as explained in the appendix of Cho al. (2024). We direct the reader to Appendix D of their work for architectural specifications. In summary, we use decoder-only Transformer with softmax operation, but omit the normalization layer for simplicity. Note that, since our construction involves single-layer model, we omit the superscripts (/) in the parameter matrices/vectors and the size of dimensions d@K,n and d() for simplicity.
D.3 INPUT SEQUENCE
We aim to construct a decoder-only Transformer model capable of solving the addition ai + a2 + + am of m operands whose lengths are at most n, where we regard every single digit as single token. We employ the same input format illustrated in Figure 4. Now, we describe how to transform this addition problem into the input sequence I, which will be fed to the Transformer.
We begin by introducing the scratchpad. Let bj Σl= ag (for Vj [m]) and by O, representing the cumulative sum up to the j-th operand; thus, 6m = b. As in the experimental setup, the scratchpad contains every intermediate result (1b:) 1 that arises during the addition process, with each result separated by an arrow (-+).
Next, apply zero-padding to every number in {arm and ( = length, L, which is determined by the maximum (possible) length among them. We set & = 1 1+ [logıo m] since the result of adding numbers each with n digits can have a length at most 1 1+ ((10" 1)m) < + [log m.j.
that they have the equal
Also, we reverse the digits within each number in the response sequence (fb,) o)),, while keeping the numbers in the query (fa;)." ) in their original order, which is consistent with the experimental setup.
To recap, the input sequence I can be formally written as 0102." yN of length N (2m + 1)(E + 1) consisting of the following parts:
1. BOS token d1 '8'
2. i-th operand A, "( "x (((+1) where og E (O.....9) (for i [m], zero-padded a;);
i-th addition symbol 0((++1)+1 = '4' (for i im 1)):
4. equality symbol 0m(4+1)+1 '=:
24

---

Published as a conference paper at ICLR 2025
5. placeholder zeros Bu 0m(++))+2" (m++1](4+1) -0... izero-padder bu);
i-th arrow symbol (m++)((+1)+ " (for [m)):
7. i-th (reversed) intermediate step B; = (m++)(e+1)+ (m(+it))((+)) where oj E (0,....9) (for E [m]. reversed & zero-padded b;).
We call 01 ... ..--m(((+1) by the query sequence, and 0m(f+1)42"' by the response sequence. We note that during the inference process, the response sequence might be incomplete (i.e., N < (2m + +)(E + 1)), as each digit these components will be inferred one by one using the next-token prediction mechanism. However, in this section formal construction, setup, where we infer all the digits of the response sequence simultancously in a single forward pass via next-token prediction. Specifically, aim output sequence O of ..N.. where o' "m(++1)+1 '' is identical to mm(++))+2 . o and 0'x = * (EOS).
focus the training
input sequence I O1''.ON" to produce 1
We summarize this process with example. Given an addition problem 57 + 48 96 201, the transformed input sequence I becomes $057 048 096 000 - 750 + 501 102, with m = 3, n 2, f 3, and = 28. The goal to generate an output sequence that ends with 000 - 750 + 501 + 1028.
D.4 ENCODING FUNCTION
We now define the encoding function, which maps the input sequence V to the initial embedding matrix (0) E RdxN In this representation, each column corresponds to the embedding vector an individual token. The embedding matrix X constructed by concatenating the token embedding matrix and the position embedding (PE) matrix. We note that this construction can also be interpreted as the clement-wisc sum of these two different embedding matrices.
An example of the embedding matrix is presented in Table 4. The first 19 rows correspond to the token embedding matrix, while the subsequent (2P1 + 222) rows represent the PE matrix. will use this example throughout this appendix to visualize construction with tables.
D.4.1 TOKEN EMBEDDING
The token embedding consists of 19 dimensions, which will refer by the following names for clarity:
1=NUM, 2=1S _Bo,, SEFULL_ONES, 1=PRE_SUM, S=PRE_CARRY. CARRY,
6=PRE_ARROW, -=PRE_EOSS (8...,,7)=SUM,, 18=ARROW, and 19=EOS.
To enhance readability, we will refer to each dimension by its corresponding name rather than by its index.
Initially, the last 16 dimensions are set zero, while the first three dimensions are filled with their corresponding values. Below, we explain how the values for NUM, IS BOS, and FUL.L_ONES are determined:
Dimension 1 (NUM). If the token is digit (O...., 9), we fill the dimension NUM by the token's value. Otherwise, for tokens +, =, >,, and $ we put zero.
Dimension 2 {s_BoS). If the token is the BOS token ('S',, we put the dimension IS_BOS. Otherwise, we put zero.
Dimension (FULL_ONES), The dimension FULL_ONES set to for every token.
D.4.2 COUPLED POSITION IDs AND POSITION EMBEDDING
In this section, we explain how the PE vector determined for each token. As in our experiment, we first assign position IDs p1 (k) and p2(k) for each token σg- Then, we map each position IDD to
25

---

Published as a conference paper at ICLR 2025
Table 4: Example initial encoding. We consider $057 + 048 + 096 = 000 - 750 - 501 - 102 as an input sequence and the position ID offsets are chosen by (for the first PE module) and 2 (for the second PE module). We denote dimensions for PE vectors as POS POS_1=(20.....
19), POS POS___________ [__EETEEEPPP + 20.....P. + 19). POS 2os_2=((P) + 20.... .P1 + P3 + 19). and POS NEXT= P2 20... 2P + 2P + 19). The vectors of the form da are defined in Equation (3). The gray cells will be filled in later.
I
$
0
0
4
8
0
9
6
PI (2(-)
2
2
2
3
3
3
3
4
4
4
2
1: NUM
7
0
0
4
8
0
0
9
6
0
2: IS_BOS
O
O
0
0
0
0
0
0
0
0
0
0
3: FULL_ON WULL_ONES
1
1
1
1
1
1
1
1
1
1
1
PRE_SUM
0
0
0
0
0
0
PRE_CARRY
0
0
0
0
0
6: PRE
PRE_ARROW ARROW
0
0
0
0
0
0
0
0
0
0
0
7: PRE_ EOS
0
0
0
0
0
0
0
0
0
8-17: SUM
O10
010
010
010
0 10
110
O1d
01
18: ARROW
19: EOS
POS
POS 1 NEXT
OP
POS POS_2
0
172
P2
NEXT
OP.
0/
mg
5
0
1
1
0
2
P1
7
6
7
4
5
6
7
4
5
6
7
(-(-)
2
2
2
3
3
3
3
4
4
4
4
5
5
5
5
1: NUM
0
0
0
0
7
5
O
0
5
0
1
0
1
0
2
2: IS_ S_S
0
0
0
0
O
O
O
0
0
0
0
0
0
0
0
FULI-OONES
1
1
1
1
1
1
1
1
1
1
PRE_SUM
0
0
0
0
0
5: PRE CARRY
0
0
0
0
0
0
0
0
0
0
0
6: PRE ARROW
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
PRE EOS
0
0
0
0
0
8-17: SUM
10
010
010
010
110
11
010
010
10
10
110
10
110
010
010
18: ARROW
19: EOS
POS
0
P1
P1
P1
P1
P1
POS NEXT
POS_2
0
12
17
POS NEXT
1/
0/
1/
PE vector: from p. we map a 2iideiinnsio vector (i = 1,2). We use two different embedding modules to map position IDs and then concatenate them.
We begin by assigning two position IDs to each token. We introduce two hyperparameters: max [ox__oss(22 + 1) and max _ox__pos((2 m 1), which the maximum possible position ID for the first and the second PE modules, respectively. Then, we select a position offset for each module: E [max posı max__pos_ L) and E [max_poss, pos2 m/. The position IDs, Pı (k) from the first PE module and (2(k) from the second PE module, are assigned to each token dg in the input sequence I as follows:
0,
k= 1, (corresponding to '8' token) k=
p. (k))
81 (i(+1)))))
for [m], k= k= (m+))((+)))))) for [m+1]]
81 11-(mmii-1)((t))-
(I)
26

---

Published as a conference paper at ICLR 2025
and
0,
(corresponding to 'S' token)
k
82 +
k ,,...,m((. 1);
P2(k)
i
(2)
k mll +
82 +
mll + 1) + l...... (2m 1)(e + 1).
E+1
Due to the complexity of the formal expressions, we encourage readers to refer to the dimensions POS_ and POS_2 in Table 4 for concrete examples.
Next, we explain the design of the PE vector for each position ID. We define by (( as the i-th (from left) digit of D-digit binary representation of k: 1. The vector oe (k [2-]) is defined as:
(-1) 10 D
ERD
(3)
This can be interpreted as a vertex of D-dimensional hypercube [-1, Importantly, for +
||0 D, < D
hold. This property will later be utilized in the construction of the attention layer. We set PE vector for the level-1 position ID P1 (i) 0 (indicating BOS token) O2p, For cases where (i) >> 0, the level-1 PE vector is defined as:
v Pi (1)
E
R22
(5)
v (1)++]
but we use of instead of v o (1))+) when (i) 2nh
Similarly, the second PE module is defined such that the PE vector is set to O2r, if P2(i) 0;
otherwise, is defined as:
((())
E RRP.
(6)
Ps
p2(1+1
we use via instead of Po (2(0))1 when p2(1) 2%
Recall that the format of can represent 2 distinct directions. So, can set the hyperparameters for maximum position IDs as max_pos- < 2P: and max_pos2 < 2P, Also, recall that & + < max_ poo__pos1..0 m + max 2o82, and & n + 1 + [log1o m.. Thus, it suffices to choose Pi 2
log, (n + 1 [log m]) and Pa > lo 1). Since d 2(P: + P2) + 19, sufficient choice of embedding dimension is
d = 2 [log2 (n + + log m) m])] + 2 [log,(m + 1)] + 19 + log, 1))
O (logg(n
In the next section, we present the construction of the model and demonstrate its capability solve the multiple operand addition problem when the above conditions are met.
D.5 TRANSFORMER BLOCK: CAUSAL ATTENTION LAYER
We now introduce the construction of the Transformer block, which utilizes 4 attention heads. The output from these heads are recorded to the dimensions PRE_ PRE_SUM, PRE CARRY, PRE_ PRE_ARROW, and PRE EOS. The feed-forward layer will then process these 4 dimensions and store the results in dimensions SUM and EOS. Finally, the linear readout at the final layer uses these outputs to predict the next token in the sequence.
27

---

Published as a conference paper at ICLR 2025
D.5.1 ATTENTION HEAD 1: DIGIT-WISH ADDITION WITHOUT CARRIES
The goal the first attention head is to perform digit-wise addition between two single-digit integers and store the result the dimension PRE SUM. Recall that by. defined as Ei=1 0g can be computed by adding aj to bj-1. This means that predict digit in bj. the model needs attend to two tokens: the first token (placed in az) which is in the query sequence, and the second token (placed in b which is in the response sequence. Thus, we aim design the query weight matrix Q1 and the key weight matrix K1 that generates an attention pattern satisfying this requirement.
Let P Pi + P2 and recall that d 2P + 19. Let the dimension of the first attention head be dok, = P 1. We define the query matrix Q and the key matrix K as follows:
Op x 19
yailp
0 Vai Ip
Q1
0 0p
P
0ppPP
OxxP
E RooK,.xd (7)
01
P
01xP:
01xP:
01xP2
01xPa
X11x19
VolLr
0ppxP
Op
OptPp
K1
0 Op2119
0pxxP
OpxxP
Vailp,
E Rdex
(8)
0n (( P(esspo ps_))))
01xPi
01xP:
01xPa
01xPa
where 0₁ is scaling factor, which we can choose to be arbitrarily large. Also, let dy.1 1 and define
V1 3(c T Rdv.idd
(9)
U1 = CREE SUM
pdxdy.
(10)
For better understanding, explain the structure of Q1 ( E RdoxxIxN The block yan Ir in the first row of Q1 extracts the dimensions POS 1 NEXT from (()), and scales them by vai Similarly, the block Vailp, in the second row of Qı selects the dimensions POS from (0) and scales them by Van. Lastly, the block ((eeP((e dimension FUL.L_ONES from X and scales it by P., X is concatenation of these three matrices. The computation of K X (0) follows similar steps, but it operates on the dimensions POS_I, POS POS_2_NT 2_NETT, and s__oS. For U, V1 x E RdxNN it simply copies the dimension NUM, scales by 3, and shifts it to the dimension PRE_SUM. An example of Q illustrated in Tables 5 to 7.
in the third row of Q1 takes the
(0) (0), and U₁ V
(0)
Table 5: Example of
QiX (0),
continuing from Table 4.
Var
I
5
9
I-P:
OP
V
P
V
P
P
P1
P1
P
P
(P (P(+1)-P:
O
02
02
P2
12
0
P3
0
V
P2
0
P+1:
VP
VP
JP
/P
VP
VP
VP
VP
vP
VP
VP
/P
I
0
+
5
+
5
4
O
I-Pr:
P1
v
F1
P
P1
v
P.
b
(P +11-PPP:
1
0
0.
P
U.
8
P2
V
P1
U
P+1:
VP
PP
VP
PP
VP
V
P
VP
v
P
P
/P
P
VP
28

---

Published as a conference paper at ICLR 2025
Table 6: Example of I
K, X
(0)
continuing from Table 4.
Vai
I
S
0
5
7
+
0
4
8
+
0
9
6
1-P::
Op,
11
17
11
177
1/
17
(R+1))((PH)) (P)+P))
OP
03
Pg
13
P4
13
44
P1
44
P1
U3
+1:
VP
0
0
0
0
0
0
0
0
0
I
0
O
0
7
5
0
>
5
0
1
+
0
2
I-A:
P
10
P1
P1
07
P1
P1
P1
P1
(1+1)((P(PP3)
U3
03
13
05
175
16
Pa
P1+ P2- +1:
0
0
Table 7: Example of U, V, x ( continuing from Table 4.
I
$
0
5
+
0
4
8
+
0
9
6
1: NUM
0
0
0
0
0
0
0
0
0
0
0
0
IS_BOS
0
0
0
0
0
0
0
3: BULI_ONESS
0
0
0
0
0
0
0
0
0
0
0
0
0
4: PRE PRE_SUM
0
0
15
21
0
0
12
24
0
0
27
18
0
5: PRE CARRY
0
0
0
0
0
0
0
0
0
0
0
0
0
PRE ARROW
0
0
7: PRE. EOS
0
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
01
10
0
0
0
0
0
0
0
0
110
010
0
010
18: ARROW
0
0
0
0
0
0
0
0
0
0
19: BOS
0
0
0
0
0
e0-end:
O2P
O2P
O2P
O2F
O2P
02P
02P
02P
02P
02F
02F
02F
02
I
0
0
7
5
0
5
0
1
++
0
2
1: NUM
0
0
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
0
0
0
3: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
PRE_ PRE_SUM SUM
21
15
0
15
0
3
0
3
0
6
5: PRE_ RRE_CARRY
0
0
0
0
0
0
0
0
0
0
6: PRE__RREEO ARROW
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
7: PRE_EOS
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
10
10
0
10
0
10
0
10
0
10
0
10
0
10
0
10
0
10
0
10
0
10
0
10
10
010
18: ARROW
0
19: EOS
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
20 end:
O2p
O2P
O2F
O2P
O2p
02F
02P
02F
02P
O2P
O2P
O2F
O2p
O2P
O2p
and the attention matrix A == sof softmax (Ci) E NNN (including the causal masking operation inside sof tmfx). To understand the structure A,, we first focus on the entry [Cileg. which can be expressed as following:
Now, we consider the attention score matrix C1 == ( Q1 (
P
if i-
[Cily
Pn (P 2)
else if [K = [Q.x otherwise.
In essence, [Ci], equals on P if and only if i 1 or the key vector for o; is identical to the query vector for Tj. Otherwise, [Ci]ay is less than a, P. Let zy denote the number of indices i (<) j) that satisfy [Ci], = 0r P. This allows that, with sufficiently large on, we have
1/2,
i
0
else if i> j, else if [K| = [Q)x
AAl
1/x,
0
otherwise.
29

---

Published as a conference paper at ICLR 2025
The output of the attention layer is computed by multiplying the matrix U₁ V with the matrix A An example of V A₁ is illustrated in Table 8.
Table 8: Example of U, Vx A1, continuing from Table 4. For simplicity, we omit the columns corresponding the tokens preceding the equal sign '=', they do not influence the next-token prediction in the response.
I
0
-
5
0
5
+
0
I: NUM
0
0
0
0
0
0
0
2: IS_BOS
0
0
0
0
0
0
0
0
3: FULL ONES
0
0
0
0
0
0
0
0
0
0
0
0
4: PRE SUM
PRE_SUM
0
0
0
0
7
0
0
15
9
0
11
9
1
0
5: PRE ERERA
0
0
0
0
0
0
0
0
PRE TRE_ARROW
0
0
7: PRE_ EOS
0
0
0
0
0
0
0
0
8-17: SUM
01
01
0
10
0
0
10
0
0
10
0
10
0
10
10
0
11
0
10
10
0
010
18: ARROW
0
0
0
0
0
0
0
0
0
0
0
0
19: FOS
0
0
0
0
0
20-end:
O2p O2p
O2P
O2p
O2p
02p
O2p
O2P
O2e
02p
02p
02p
02p
02p
02p
O2p
To clarify our example, consider the input sequence $057 + 048 + 096 = 000 + 750 501 -. At this stage, the goal of the model is to predict '1' as the next token, which corresponds significant digit (LSD) of the sum of (the third token of 096) and 5 (the first token of 50T). It is important to note that the query vector for the last arrow token is identical to the key vectors for 6 and 5. This can be easily verified by comparing Table 5 and Table 6. As a result, the dimension PRE SUM of UG V (0) A1 for the last arrow token can be obtained by computing 3 (0 + 18 + 15) = 11. Here, the "3" in å originates from the number of tokens that the last arrow token attends to: BOS, 6, and 5, and this is the reason why we design the matrix V1 with the scalar 3. The operation of extracting the LSD from 11 will be handled in the subsequent feed-forwal layer.
the least
D.5.2 ATTENTION HEAD 2: CARRY DETECTION
The goal of the second attention head is to fill the dimension CARRY with the appropriate values. To illustrate this, consider the partial input sequence $057 + 048 + 096 = 000 + 750 - 501 - I. the next token, the model needs to (1) compute the sum of 9 (the second token of 096) and 0 (the second token of 501), and (2) detect the carry from the previous digit's sum, which was 6 6-+55 5. The first attention head handles the first part, by making (the last token of our example input sequence) attend to 9 and O. We aim to design the second attention head to handle the second part. A key observation is that detecting the carry becomes possible when the model attends to 6 and 5. by verifying that 6 + 5 - E (9, 10) (see Section E.4.2. of Cho et al. (2024) for more details). Thus, the second attention head's 's goal is to compute + 5 and assign a value of 11 to the dimension PRE_ CARRY of the token 1.
To predict
Recall that P = P1 + P2 and d = 2P + 19. Let the dimension of the second attention head be dok,2 P We define the query matrix Q2 and the key matrix K₂ as follows:
Opix19
/azIp
0PtP:
O Vadlp,
Q2
1 Vaz VazP(ee (e ONES)
O.pPP
OppPP
0.
E R Roxxxx
(11)
01x
01xP7
01xP3
01x1s
Op 19 0 LapP(e) (e BOS)
VaIlP P. OptxPs
0 x
Op
Op Ps
K₂
0ppPP
OppxP
Va2lps
E RRdoxx2x
(12)
01xP,
01xP:
01xP2
01xP:
where a2 is a scaling factor, which we can choose to be arbitrarily large. Also, let dy.2 = and define
V22 = 3 (( T E Rdv.axd
(13)
U, = e CARRY Rdddvaa
(14)
30

---

Published as a conference paper at ICLR 2025
An example of Q2X(0, K2X(0), Similar to the first attention head, with sufficienlly large 02 the j-th token o, will only attend to the BOS token and tokens whose key vectors match the query vector of o.. However, the tokens that Jj attends to differ from those in the first attention as the design of the query matrix Q2 and K₂ is slightly modified. By adjusting the placement of VazIr, and vazIrz rg in the Q2 and K2, we control which tokens are attended to.
U2V2X(0) and U2V2X(00 (0) (0/A2 is illustrated in Tables 12.
Table 9: Example of
Q2 X ( continuing from Table 4.
vag
I
$
5
7
0
4
0
6
I-P:
OP
U
P1
16
0
0
P1
1
01
0
P1
(P(+1)-P:
Op4
U2
P2
V2
P
12
02
0/
P2
1/
04
P2
P+1:
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
I
0
0
7
5
0
5
0
1
0
1-P;
U
1
1
21
P1
P1
11
17
11
P1
P1
16
11
11
F1
P1
1
1
V
P
(P++1)-P:
12
P2
12
17
P3
P2
1/3
44
04
3
11
P2
P2
p+1:
VP
VP
P
V
VP
VP
V
P
VP
P
V
P
PP
V
P
P
VP
Table 10: Example of 1
K2X (0), continuing from Table 4.
Vag
I
$
5
7
4
8
6
I-P::
Op:
(P1 11-PP:
P
R2
0
P2
P2
1
P3
P2
Pg
P+1:
VP
0
0
0
0
0
0
I
0
0
7
O
5
2
I-Pr:
10
10
10
(P) P)-P:
P2
P2
Pg
P+1:
31

---

Published as a conference paper at ICLR 2025
(0), Table 11: Example of U2 V2 x
continuing from Table 4.
I
$
0
5
7
+
0
4
8
+
0
9
6
NUM
0
0
0
0
0
0
0
0
1S_BOS
0
0
0
0
0
0
0
3: FULL_ONES
0
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
0
PRE CARRY
0
15
21
0
0
12
24
0
0
27
18
PRE ARROW
0
0
0
0
7: PRE EOS
0
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
20
0
0
0
0
0
0
0
0
0
110
0
0
110
18: ARROW
0
0
0
0
0
0
0
0
0
0
0
19: EOS
0
0
0
0
0
0
20-end:
O2P
O2P
O2P
O2P
O2F P
O2F
02P
02P
02P
02P
02F
02p
02p
I
0
0
7
5
0
->
5
0
1
0
2
I: NUM
0
0
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
0
0
3: FULL_ONEE,
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
PRE_SUM
0
0
0
0
0
0
0
0
0
5: PRE_ EE__RARRY
21
15
0
0
15
0
3
0
3
0
6
6: PRE_ PRE_ARROW
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
7: PRE EEE
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
00
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
010
18: ARROW
0
19: EOS
0
0
0
0
0
0
0
0
0
0
0
0
0
0
20 end:
O2p O2P
O2P
O2P
O2p
O2F
O2P
O2F
O2p
OZF
O2p
Oup
O2p
O2p
O2P
Table 12: Example of U2
V22 (0) A2. continuing from Table 4. For simplicity, we omit the columns corresponding to the tokens preceding the equal sign '=', as they do not influence the next-token prediction in the response.
I
0
7
5
0
5
1
0
2
I: NUM
0
0
0
0
0
0
0
0
0
0
S__OSS
0
0
0
0
0
3: FULL_ONES
0
0
0
0
0
0
0
0
0
4: PRE_SUM
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
5: PRE_ _AE_R
0
0
0
0
5
0
15
0
0
11
9
6: PRE_ARE PRE_ARROW
0
0
0
0
0
0
0
0
0
0
7: PRE_EOS _ES
0
0
0
0
0
0
8-17: SUM
0
10
00
0
10
10
0
10
0
10
0
10
0
10
10
0
10
10
0
10
0
10
10
0
10
010
18: ARROW
0
0
0
0
0
0
0
0
0
0
0
19: EOS
0
0
0
0
0
0
0
0
0
20-end:
O2p
O2P
OZP
O2p
O2p
O2P
O2P
O2P
O2p
O2p
O2p
O2p
O2p
O2p
O2p
O2p
D.5.3 ATTENTION HEAD 3: ARROW DETECTION
The goal of the third attention head is to fill the dimension PRE_ ARROW. We aim to put 1 if the next token the model has to predict is the arrow (-), and otherwise, we will put strictly smaller values (below 1/2).
Recall that d = 2P Let the dimension of the first attention head be dok,3 = P1 + We define the query matrix Q3 and the key matrix K3 as follows:
Q3-
Op (e
Opt
Vaglp
0pxP2
Optxps
Rdoxaad
(15)
O
O
O
01xP,
K₃-
O Xp1x19
yaIp. 0
O 01xP,
Opixp
OptxP=
E Rddokadd
(16)
(eel
T
01xP3
01x13
32

---

Published as a conference paper at ICLR 2025
where ar is scaling factor, which can choose be arbitrarily large. Also, let dy,3 1 and define
V (e (eis_no)) T Rddaaxd fss_
(17)
U = PRE PRE_W E IR Iddxdv
(18)
An example of Qx) (3) V3 and 33V (0) A3 illustrated in Tables 13 16. Similar to the first and the second attention head, with sufficiently large 03, the j-token o, will only attend the BOS token and tokens whose key vectors match the query vector of Ø,
To enhance understanding, consider two input sequences I1 = $057 048 + 096 000 + 750 and I2 = $057 + 048 + 096 000 - 75. We first analyze I1. By comparing Table and Table 14, observe that the token 0 (the last token in 750) only attends to the BOS token. Therefore, the dimension PRE_ ARROW of the token 0 will be set to in the matrix U3 V X (0) A3. Next, for I2, can see that the token (the second token in 750) attends to the BOS token and four other tokens (0 from 057, from 048, 0 from 096, 0 from 000). As result, the dimension PRE ARROW of the token will be filled by 1/5, where the denominator comes from the sof tmax operation.
Consequently, the model can decide to predict the arrow (+) the next token if the dimension TRE_ARROW in the matrix U3 V3.X (0) A3 of the last token of the given input sequence is equal to 1. Otherwise, for the case where the model should not predict the arrow as the next token, PRE_ARROW is set to the value less than 1/2.
Table 13: Example of
Qa X
(0), continuing from Table 4.
yas
I
8
0
4
8
0
I-P::
OP
P1
U
U
17
P
V
P1
0
U
P 11
VPI
VPI
VP
VPI
PP
VI
P
VP
VP
VPI
VPI
VP
VPI
I
+
0
+
1
I-P::
U
0
P
1
P1
v P1
P
P
v
P1
0
P
P1
0
P
P
P1 ++
VP.
/
P1
VT
VT
v
P
V
P
P
v
P
VI
RR
VT
Table 14: Example of
1
K X
(0), continuing from Table 4.
Vag
I
$
0
7
+
4
+
0
I-P;:
0p,
P1
P:
1/
P1
11
P1
17
P1
1/
P1
P:
1/
P1
87
P1
1/
P1
P 11:
VPi
0
0
0
I
0
0
+
7
0
0
I-P:
0.
06
P1
P1
P1
07
P1
P1
Pi
P +1:
0
33

---

Published as a conference paper at ICLR 2025
(0), Table 15: Example of U3 V2X x
continuing from Table 4.
I
$
0
5
+
0
4
8
+
0
9
6
NUM
0
0
0
0
0
0
0
0
1S_BOS
0
0
0
0
0
3: FULL__ONE
0
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
0
5: PRE EEERY
0
0
0
0
0
0
0
0
0
0
0
0
PRE ARROW
7: PRE EOS
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
10
0
0
0
0
0
10
0
0
0
0
0
0
10
0
110
18: ARROW
0
0
0
0
0
0
0
0
19: EOS
0
0
e0-end:
02P
O2P
O2P
O2P
O2F
O2F P
02P
02P
02p
02F
02F
02P
02p
I
0
0
7
5
0
->
5
0
1
0
2
I: NUM
0
0
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
0
0
3: FULL_ONES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
PRE_SUM
0
0
0
0
0
0
0
0
PRE_ PRE_CARRY
0
0
0
0
0
6: PRE_ARROW
0
0
0
0
0
0
0
0
0
0
0
0
0
0
7: PRE EES
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
10
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
010
18: ARROW
19: EOS
0
0
0
0
0
0
0
20 end:
O2P O2P
O2F
O2P
O2p
02F
O2P
O2F
O2P
O2P
O2P
Oup
O2p
O2P
O2p
Table 16: Example of U3V (0) As. continuing from Table 4. For simplicity, we omit the columns corresponding to the tokens preceding the equal sign '=', as they do not influence the next-token prediction in the response.
I
0
7
5
0
5
0
I: NUM
0
0
0
0
0
0
0
0
0
0
0
IS BOS
0
0
0
0
0
3: FULL_
0
0
0
0
0
0
0
0
0
0
4: PRE SUM
PRE_SUM
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
5: PRE_CARRY
0
0
0
O
0
0
0
0
0
0
0
0
PRE_ ARR_WRROW
1/4
1/4
1/4
1/5
1/5
1/5
1/6
1/6
1/6
1/7
1/7
1/7
7: PRE_EOS
0
0
0
0
0
0
0
8-17: SUM
10
00
0
10
10
0
10
0
10
0
10
0
10
10
0
10
0
0
10
0
10
10
0
10
00
18: ARROW
0
0
0
0
0
0
0
0
0
0
0
19: EOS
0
0
0
0
0
0
end:
O2p
O2P
O2P
O2p
O2p
O2P
O2P
OZP
O2p
O2P
O2P
02p
O2P
O2P
O2p
O2p
D.5.4 ATTENTION HEAD 4: EOS DETECTION
The goal of the fourth attention head is to fill the dimension PRE_EOS. We aim to put 1/2 if the next token the model has to predict is the EOS token.
Recall that P = Pt + P2 and d = 2P + 19. Let the dimension of the second attention head be dQK,4 P We define the query matrix Q4 and the key matrix K₁ as follows:
Op,x19
/a
0ppP:
Opix2 vailp,
0ptxP
Q1
Op
OppPP
OptxPs
0
E BRdoxaxd
(19)
Pee
01x
01xP3
01xP3
01xP;
Op x19 Op 19
VasIp p, OptPP
Op xP
Xp PP Voulry
Optxp
K,-
OppPP
Op
RRdaxa
(20)
(e
01xP.
O1xP,
01xPa
01xP,
34

---

Published as a conference paper at ICLR 2025
where a4 is scaling factor, which can choose be arbitrarily large. Also, let dy, 1 and define
V = (ess (Eis_oo) ) E Rdvad
(21)
U = ERE_ EOS E Rdxdv,
(22)
An example Q4 K2 UVV,X(0), and VVV4X A, illustrated in Tables 17 20. Like the previous attention heads, with sufficiently large a4, the j-token oj will only attend to the BOS token and tokens whose key vectors match the query vector of g.
As mentioned earlier, the fourth head aims to fill 1/2 in the dimension PRE_EOS if the model has to predict the EOS token as the next token. However, it might not be very clear that such a token is not uniquely defined, as multiple tokens in the matrix U. V (0) A. can have their PRE_ [RE_ROS entry filled with 1/2, as illustrated in Table 20. To clarify. we note that the final decision regarding the EOS token made by combining the outputs from both the third and the fourth attention head (readers can check this in the subsequent feed-forware layer construction). This approach enables the model to correctly determine whether the next token should be the EOS token or not, as the token with PRE ARROW and PRE EOS set to and 1/2 is uniquely identified. We also note that PRE_FOS can only be either 1/2 or 1/3 which will be utilized the feed-forward layer construction.
Table 17: Example
Qx
continuing from Table
Vou
I
$
5
7
4
8
6
I-Pr:
OP:
V
P.
1
P
10
11
w
P.
P
15
10
P
P
(1++))-P:
On
U2
P2
P2
22
12
P2
P2
P2
03
Pg
P2
P2
P+1:
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
VP
/P
I
0
0
7
7
0
7
5
2
I-Pr:
U
U
P
P1
P
05
10
P
v
P
05
V
(Pi +1)-P:
14
12
12
12
12
15
P+1:
VP
P
P
P
P
P
P
P
V
VP
Table 18: Example of 1
K K (0), continuing from Table 4.
yau
I
$
5
7
4
8
6
I-P::
Op:
V.
(P1 1)-PP:
P
P2
02
P2
02
P2
Pg
03
P2
Pg
P+1:
VP
0
0
0
0
0
0
0
0
I
0
0
7
O
5
2
I-Pr:
10
10
10
10
(P) P)-P:
0
Pg
P+1:
35

---

Published as a conference paper at ICLR 2025
Table 19: Example of U2 (O), continuing from Table 4. For simplicity, we omit the columns corresponding the tokens preceding the equal sign '=', they do not influence the next-token prediction in the response.
(0),
+
4
8
+
0
9
6
NUM
0
0
0
0
0
0
0
0
0
0
2: IS_BOS
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
0
4:
0
0
0
0
0
0
0
0
0
0
0
0
0
5: PRE_CARRY
0
0
0
0
0
0
0
0
0
0
0
PRE PRE_ARROW
0
0
0
0
0
7: PRE EOS
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
1D
0
0
0
0
0
0
0
0
0
0
0
0
18: ARROW
0
0
0
0
0
0
0
19: EOS
0
0
0
0
20-end:
O2F
O2F
O2F
O2F
O2F
O2F
02P
O2F
O2f
O2F
O2f
02f
02F
7
5
0
5
0
1
0
2
I: NUM
0
0
0
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
0
0
0
3:
0
0
0
0
0
0
0
0
0
0
0
0
0
0
PRE_
0
0
0
0
0
0
5: PRE CARRY
0
0
0
0
0
0
0
0
0
0
0
0
0
0
6: PRE_ARROW
0
0
0
0
0
0
0
0
0
0
0
0
0
0
7: PRE_
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
8-17: SUM
11
10
0
10
10
10
10
110
110
01
110
010
10
10
10
010
18: ARROW
19: EOS
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
20-end:
O2P
02P
02P
02P
O2P
O2P
O2P
O2P
02P
02F
O2F
02F
02P
02p
02p
Table 20: Example of U, U,V.XX A4, continuing from Table 4. For simplicity, we omit the columns corresponding the tokens preceding the equal sign "=", as they do not influence the next-token prediction in the response.
I
-
5
-
5
+
1: NUM
0
0
0
0
0
0
0
0
2: IS_BOS
0
0
0
0
0
0
0
0
0
0
0
0
0
3: FULL_ OOLESSS
0
0
0
0
0
0
0
0
0
0
0
0
4: PRE_ RRM_
0
0
0
0
0
0
PRE PRE_CARRY
0
0
0
0
0
0
6: PRE ARROW
0
0
0
0
0
0
0
0
0
7: PRE_EOS
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/2
1/3
1/3
1/3
1/2
1/2
1/2
1/2
8-17: SUM
010
010
100
110
110
110
110
010
11
110
010
110
0
110
010
010
18: ARROW
0
0
0
0
0
0
19: EOS
0
0
0
0
cn-
O2p
O2P
02P
OZP
O2p
02P
O2P
O2P
O2P
O2P
O2P
O2p
O2p
O2P
02P
O2p
D.5.5 RESIDUAL CONNECTION
We now consider the residual connection. The output of the attention layer can be expressed as follows:
YO)= ())
Σ
V.V.X (0) An
(23)
hE(1.23344
One can observe that the dimensions PRE_SUM, PRE CARRY, PRE ARROW, and PRE_EOS in X (0) are empty, whereas these same dimensions in Ene (123334) U, V, (0) A, contain non-empty values. Thus, the residual connection effectively "fills in the blanks" " in the input embedding matrix. An example of the output of residual connection is presented in Table 21. Again, we note that the
36

---

Published as a conference paper at ICLR 2025
error from the softmax operation be made negligible by setting the scalars 0r,, ag, Org, and Crq sufficienlly large.
Table 21: Example output of residual connection, Y continuing from Tables 4, 8, 12. 16 and 20. The orange rows correspond the values computed through the attention layer. We omit the columns corresponding to the tokens preceding the equal sign '=', as they do not influence the next-token prediction the response. The gray rows will be filled during the subsequent nerdfforwrrd layer.
I
0
5
0
5
0
2
I: NUM
0
0
0
0
7
5
0
0
5
1
0
0
IS_BOS
0
0
0
0
0
0
0
0
FULL ONES
4: PRE SUM
0
0
0
7
0
0
15
9
0
11
9
1
5: PRE CARRY
PRE_CARRY
0
0
0
O
7
5
0
0
15
0
0
11
9
6: PRE_AN PRE_ARROW
1/4
1/4
1/4
1/5
1/5
1/5
1
1/6
1/6
1/6
1/7
1/7
1/7
PRE_EOS
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/2
1/2
8-17: SUM
Ozo
O10
O10
O10
010
110
100
0
11
1
O10
O10
1uu
00
10
18: ARROW
0
0
0
0
0
0
0
0
0
0
0
0
0
19: EOS
0
0
0
0
0
0
0
POS 1
O
0
1
0
0
17
10
0
0
P1
POS DOS_I_NEE NEXET
V4
18
06
POS 2
02
01
POS NEXT
03
03
03
Us
D.6 TRANSFORMER BLOCK: TOKEN-WISE FEED-FORWARD LAYER
The goal of the feed-forward layer is fill the dimensions SUM, ARROW, and EOS. While the attention layer enables interactions between different tokens, the leed-forwart layer operates solely on the dimensions within each token. After processing, predict as the next token, while ARROW and EOS serve as flags of whether the next token should be an arrow (-) or the EOS token ($), respectively. Based on the output of the feed-forward layer, the next-token prediction is carried out in subsequent linear readout process.
specifies the number the model will
We will construct one-hidden-lug 1 to 7. Each network will handle specific output: SUM (8-17), ARROW (18), and BOS (19), respectively. The goals of each network are follows:
ReLU networks (FF). FF?. FF that take inputs from dimensions
FF If the model has to predict a digit (let's say & E {0, 10.1...., 9)), it assigns the value of the dimension 8 + ki which is the (A + 1)-th dimension of SUM, while setting all other dimensions in SUM to
FFî: If the model has to predict the arrow (-+ as the next token, set ARROW by 1.
FFî: If the model has to predict the EOS token the next token, set EOS by 1.
construct single one-hidden-layer ReLU network FF1, that takes inputs from dimensions to 7 and outputs the proper values at dimensions 8 to 19. We provide our example in Table 22.
By combining these three net-nes
D.6.1 SUBNETWORK 1: CONSTRUCTION FOR SUM (DIMENSION 8-17)
For FF], we can apply the construction provided in Section E.5.1 of Cho et al. (2024). Below, we provide an overview of their approach, and refer readers to Cho et al. (2024) for the underlying principles.
Since the feed-forward layer processes tokens individually, we denote each column vector ( as IRd in unified manner. We first define linear function g Rd IR as:
g(x)= = 3PRF_ SUM + IPRE_CARRY CARRY
@NUM
+ 0.21.
10
37

---

Published as a conference paper at ICLR 2025
,.. Using 9, construct one one-hidden-layer ReLU network fs :R R (k 0, ........ 0.......9 defined
(z) = 20 -0.)) o(x k) o(x (k: + 0.5)) + 6(x (k+ 1))
:cx (I: + 9.5)) <(x - +k+ 10)) o(x (A: 10.5)) o(x (I +11)]
Now, we construct FF} as
(x:) SUM = [fo (g (x))
fo (g
We note that the construction of g slightly differs from the original formulation in Cho et al. (2024), due to the difference in how value stored in the PRE_ CARRY. Specifically, we store the sum of digits at the (previous) same significance level, whereas Cho et al. (2024) stores the result of additional summation of the value in NUM to the sum of digits.
Intuitively, FF} is designed to output a one-hot vector that acts flag indicating the digit which the model has to predict as the next token. The difference between the value in the PRE_CARRY CARRY and NUM takes role of detecting whether a carry occurs from the previous digit, while fú. identifies whether x corresponds to l or k 10. Certain constants, such as 0.21 and 0.5, are chosen to manage numerical errors.
D.6.2 SUBNETWORK 2: CONSTRUCTION FOR ARROW (DIMENSION 18)
We will construct a subnetwork FF Rd + Rd that outputs the dimension ARROW the dimension PRE ARROW is set to 1; otherwise, outputs O. As PRE ARROW can have value less than 1/2 if is not 1, this can be easily achieved by constructing FF as
26 ARROW 1/2),
where æ represents each column vector of Y
D.6.3 SUBNETWORK 3: CONSTRUCTION FOR EOS (DIMENSION 19)
We will construct subnetwork FF? Rd I that outputs in the dimension EOS if the dimension PRE_ARROW is set to 1 and PRE_EOS BOS is set to 1/2. We construct FF as
|FF [])s 2 ARRO 1/2) + 66( (DREEEES 1/3) - 1,
where represents each column vector of (1), Note that PRE_ARROW can take values of either 1 or less than 1/2, and PRE_EOS can be either 1/2 or 1/3. Therefore, the output FF equals I only when the dimension uniquely identifies the token as the most significant digit of bm-
is set to and PRE_BOS is set to 1/2. Importantly, this condition
Table 22: Example of the output of the feed-forward layer. The yellow rows represent the values that are generated during the feed-forward operation. We omit the columns corresponding to the tokens preceding the equal token '=', as they do influence the next-token prediction.
I
0
0
7
5
0
-
5
1
0
2
1: NUM IS_BOS
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
FULL_ONES
0
4: PRE_ SUM
0
0
0
0
0
0
0
0
5: PRE PRE_CARRY
0
0
0
0
0
0
0
0
0
0
0
6: PRE_ARROW
0
0
0
0
0
0
0
0
0
0
0
PRE_FOS
0
0
0
8-17: SUM
e e1
10
10
C1
10
10
e
6
10
10
e1
10
C1
18: ARROW
0
0
0
0
0
0
0
0
0
0
19: EOS
0
-1
-1
-1
0
-1
-1
0
20-end:
O2p
O2p
02p
O2p
O2P
02p
02p
O2p
02p
O2p
O2p
O2p
02F
O2p
Ogp
O2p
38

---

Published as a conference paper at ICLR 2025
D.6.4 RESIDUAL CONNECTION
Similar to the residual connection applied after the attention layer, the residual connection following the feed-forward layer "fills in the blanks" of the matrix Y with the output of each subnetwork as
X (1) Y + + (Y
The example of (1) is illustrated in Table 23.
Table 23: Example of after applying the residual connection, X (1), continuing from Table 22.
+
+
5
NUM
0
5
IS_BOS
0
0
0
0
0
0
0
3: FULL_ONE LONES
1
1
1
1
1
1
1
1
1
1
1
4: PRE_ SUM
0
0
0
7
0
0
15
9
0
11
9
1
PRE_ PRE_CARRY
0
0
0
5
0
15
0
11
PRE PRE_ARROW
1/4
1/4
1/5
1/5
1/6
1/6
1/7
1/7
7: PRE PRE_EOS
1/2
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/3
1/2
1/2
1/2
1/2
8-17: SUM
61
eg
G6
e
CG
C2
C3
18: ARROW
0
19: FOS
-1
-1
POS_I
POS POS_I __NETT
U
U
1
1
0
v.
U
POS_2
V
12
17
14
Ve
P3
P2
POS NEXT
0/
03
03
08
D.7 DECODING FUNCTION
The final step is decoding: the model decides which token to predict as the next token based on embedding matrix. Specifically, with weight matrix W... ou E RIV|xd, the model first compute the multiplication between Won. gd and X (1), Then, the model takes (token-wise) arg-max
operation for greedy decoding. Mathematically, the next-prediction at i-th token can be written as follows:
h; arg max ok W Wout. X x [[or
oie }
(24)
kE[lv]]
The design of the weight matrix W... R illustrated in Table 24, and the example of the matrix X (1) and the output sequence is presented in Tables 25 and 26, respectively.
Table 24: The transposed weight matrix W
out of the linear readout in decoding function.
0
7
$
1-7: NUM-PRE_EOX
07
07
07
07
07
07
07
07
07
07
07
07
07
07
8: SUM)
0
0
0
0
0
0
0
9: SUM2
0
0
0
0
0
0
0
0
0
0
0
10: SUM3
0
0
0
0
0
0
0
SUM,
0
0
12: SUM;
0
0
0
0
0
0
0
13: SUM6
0
0
0
0
0
0
0
0
0
0
14: SUM7
0
0
0
0
0
0
0
15: SUM$
0
0
0
0
0
16: SUM9
0
0
17: SUM10
0
0
0
0
0
0
18: ARROW
0
0
0
0
0
0
0
10
0
19: BOS
0
0
0
0
0
0
100
20-end: POSITIONS
OZP
O2F
OZP
O2P
O2F
O2P
O2F
O2P
02F
O2F
Ozp
02P
02F
O2F
39

---

Published as a conference paper at ICLR 2025
Table 23: Example output linear readout, W X continuing from Tables 23 and 24. The yellow cells represent the maximum value of each column.
0
7
0
0
0
2
0
1
0
0
0
0
0
1
0
0
0
0
0
0
0
0
2
0
0
0
0
0
0
0
4
0
0
0
0
0
0
0
0
5
0
0
0
0
0
0
0
6
0
0
0
0
0
0
0
0
0
9
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
10
-100
-100
-100
-100
-100
0
-100
-100
-100
100
Table 26: Example output sequence O continuing from Table 25.
0
0
-
5
0
4
0
-
0
0
0
7
0
5
0
1
0
2
40

---

Published as a conference paper at ICLR 2025
E
MORE ATTENTION PATTERNS OF TRAINED TRANSFORMERS
Continuing from the discussion in Section 4.4 the attention patterns due the (non-Jexistence of scratchpad, we showcase more examples of the attention matrices sof softmax( QK + A) of actually trained Transformers (where A is causal mask).
E.1 ATTENTION PATTERNS WITH SCRATCHPAD
Attention Puttem of layer head
Attention Pattern of layer head
IReft Scratchpud!
(toFE: Scrutchpedi
Figure 15: Some attention patterns of 6-layer 8-head decoder-only Transformers trained with NoPE and using the scratchpad.
41

---

Published as a conference paper at ICLR 2025
(PILE bcratcppe)
(Pine scratchpal)
Attention Puttem llaeer head
Attention Patten of layer head (MAE
Figure 16: Some attention patterns of 6-layer 8-head decoder-only Transformers trained with FIRE and using the scratchpad.
(Positinn Coupling
i Scrutchpad!
(Pestion Coupling Scr Scralchoud)
Figure 17: Some attention patterns 6-layer 8-head decoder-only Transformers trained with bi-level Position Coupling and using the scratchpad.
42

---

Published as a conference paper at ICLR 2025
E.2 ATTENTION PATTERNS WITHOUT SCRATCHPAD
(NorE, atchppd)
INoPe, Scrutchpall
Figure 18: Some attention patterns of 6-layer 8-head decoder-only Transformers trained with NoPE but not using the scratchpad.
Attevttion Puttem layar 6 head
Attention Pathern of layer head
Figure 19: Some attention patterns of 6-layer 8-head decoder-only Transformers trained with FIRE but not using the scratchpad.
(Position Coupling i Scrutchpeei
(Pesttion Couplng Scralccppe))
Figure 20: Some attention patterns of 6-layer 8-head decoder-only Transformers trained with (single-
level) Position Coupling but not using the scratchpad.
43


Reviews:
[{"review_id": "ddI8foPPJ0", "reviewer": "Reviewer_joP4", "summary": "This paper studies the problem of length generalization in transformer models over arithmetic operations, particularly over both number of operands and length of operands. They achieve 2-3$\\times$ length generalization on both addition and multiplication using task-specific scratchpads and multi-level versions of Position Coupling (Cho et al., McLeish et al.). They also provide a theoretical result for a 1-layer transformer showing that it can solve multi-operand addition up to lengths and counts exponential in the embedding dimension.", "questions": "1. Is the main contribution of the paper over Cho et al., proving length generalization over number of operands and the introduction of multi-level position coupling? I would appreciate this being made more clear.\n2. For the comparison in Figure 3, why use NoPE? What about FIRE or Rotary position embeddings? With NoPE, is the model trained for longer durations?\n3. Is there a comparison with He et al., and Zhang et al. on multi-level relative PEs?\n4. In the data sampling, what is the purpose of the second chunk? Is it so that you ensure that you see every position?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well written with well described experiments and intuition.\n2. Length generalization is known to be challenging on arithmetic problems with Transformers. The results in this paper are extremely impressive in that they not only achieve length generalization in operand lenght, but also operand count.\n3. The authors perform extensive ablations on th effect of zero-padding, trained-lengths and also connect some of their experimental results with their theoretical construction.", "weaknesses": "1. While the problem of length generalization is important, it is mainly relevant in how it applies to more general tasks. However, position  coupling seems to be highly task specific and it is not clear to me how to generalize this idea to non-arithmetic tasks other than very artificial problems.\n2. The size of the test set seems to be too small for 30 operands of length up to 30. I would like to see if the test accuracy remains roughly the same even when scaling up the test set sizes.\n3. The model assumes a singly digit tokenizer. It would be interesting to see if the results extend to more general tokenizers at all.\n4. The zero-padding to match lengths seems to make the input size explode. If you have n inputs with max length n, then the input length will be $O(n^2)$ even if only one operand is of max length.\n5. Do you have any intuition on why 1L8H has particularly bad performance? Shouldn't it be strictly more expressive than the 1L4H construction?\n6. In Theorem 4.1, does \"solve\" mean solve perfectly? Are there any constants in max digits or operands? $2^d$ is a very large number, so this is an extremely impressive result!", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GrscOqjBKk", "reviewer": "Reviewer_ois9", "summary": "The paper extends the idea of position coupling to get better generalization for the tasks of addition and multiplication. Position coupling is one of the several methods for leveraging the symmetries of a task, and in this paper, it is used in tandem with scratch-padding, padding the integers, and reversing. The paper instantiates generalization for up to 30 digits for addition with the mentioned techniques, and around 20 digits for multiplication, which is impressive.", "questions": "Please address the concerns expressed in the weaknesses.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1- The paper is well written, with their ideas well articulated in all the sections. By comparing their methods to some of the previous work they have made the reader convinced that their methods are necessary in the given framework (leveraging the structure grounded on APE) to achieve length generalization in arithmetic. \n\n2- The idea of position coupling is applied to parity to achieve perfect length generalization. \n\n3- Tri-level position coupling facilitates length generalization on both the first and second operands in multiplication.", "weaknesses": "1- In line 18, the paper claims to offer the first length generalization results for arithmetic tasks. To my knowledge, there are several other works that have argued the same. For instance, [1] has leveraged the same symmetries of addition (In a different framework, and using RPE)  to offer generalization up to 50 digits (for both operands) when trained on samples up to 5 digits. Can you explain the advantages of your work? Authors should have done a better comparison with previous work and stayed truthful about the results. \n\n2- Scratch-padding has been explored before for the problem of parity as also mentioned in line 179. I'm not convinced that position coupling is the core solution for this problem, since it's been only compared with a NoPE  + scratchpad approach. To me it is clear why such an approach cannot work in a next-word-prediction setting as the generated string in the output relies on the order of the the bits in the input, and  scratch-padding makes the final answer reliant on the previous bits in the output. Therefore, even though the final answer of parity is intrinsically independent of the string order, with this solution it is in fact regarding the order. Position coupling had to be compared with other methods that utilizes positional encoding as well. \n\n3- The contribution of the paper is pretty limited considering that position coupling is introduced in previous work, and bi-level and tri-level versions are two extensions of that method. Besides, the effectiveness of padding and reversing is already well-known and widely studied. \n\n4- The scope of the experiments are limited to solely having numbers and in the absence of any text. How does this apply to a broader range when numbers are amongst text? For instance, [1] has tried to address this with training some additional attention heads tailored for arithmetic along with rest of the architecture. Besides, is it true that with your method multiplication and addition require different sets of positional encoding? Please correct me if I'm wrong. \n\n\n[1] Sabbaghi et al., Explicitly Encoding Structural Symmetry is Key to Length Generalization in Arithmetic Tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nHfhywfYJ3", "reviewer": "Reviewer_uejG", "summary": "The paper describes methods allowing decoder-only transformers to generalize to longer operands when performing integer addition or multiplication, and to generalize to a larger number of operands when performing addition.\n\nThe authors mix two existing techniques, scratchpad, which forces the model to output a sequence of intermediary steps before producing the result, and position coupling, problem-specific hierarchical positional encodings. They experiment with three tasks. \n\nFor parity calculation over bit strings, they use a scratchpad which computes parity integration, ie force the transformer to compute the parities of successive prefixes, and use the same (absolute) positional encoding in the input and output sequences. \n\nFor addition of $n$ positive integers, they use a scratch pad of partial sums (0, first operand, sum of the first two...) and two positional encodings: one for the position of digits in each operand and output, and one for the position of each operand and output in the operation and scratchpad.\n\nFor multiplication of 2 position integers, they use a scratchpad of single digit products and partial sums, and three positional encodings, which reflect the different subtasks in multi-digit multiplication.\n\nExperimental results indicate that such techniques allow for generalization to longer lengths on all three tasks, and a larger number of operands for addition.", "questions": "* At test time, for a model prediction to be correct, you insist on both the output and the scratchpad to be correctly predicted. Prior work suggests that models sometimes learn despite incorrect scratchpads. Have you tested models by judging them on the final result only? (this could have an impact ablation results, and comparisons with other approaches)\n* Is reversing the order of digits in the scratchpad and output sequence really necessary? Can you provide ablations for this? \n* What would be the impact of reversing the digits in the input sequence? \n* Your architecture ablations suggest that shallow models can learn and length generalize. Unfortunately, they all use a large embedding dimension (512), which translates into a large number of parameters. Have you tried models with smaller dimensions (128, 64 even)?\n* In figure 10, you report better performance of 2 layer models than 4 or 6 layer models with 2 and 4 heads, and better performance of 4 layer models than 6 layer models with 8 heads. Could this be due to the size of your training sample (500k may be on the low side for the larger architectures), or the number of steps you allow? \n* An ablation on the training set size would be useful.\n* For addition, what is the point of adding 0 and a copy of the first operand in the scratchpad? \n* The main text claims the model has 63M parameters, the appendix claims 25M, which is right?\n* Would this approach benefit from an encoder-decoder model? (intuitively, bidirectional attention over the input sequence should help, all the more as the digits in input sequences are not reversed)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper studies an important limitation of arithmetic transformers, their inability to scale to longer sequences that those seen at training. The authors propose a solution to this problem.\n\nThe paper is clearly written, the methods is clearly described and the experimental results are compelling. Some ablation results are provided.", "weaknesses": "The authors seem to only focus on positive integers. This greatly reduces the practical interest of the findings, since integer arithmetic usually involve positive and negative numbers. Could the results be extended to relative numbers, either by adding a sign in the number representation, or changing the tokenization scheme to a number system that encode negative numbers as well (e.g. encoding them in base -10 instead of base 10, or using balanced base 10, with digits -4, -3, -2, -1, 0, 1, 2, 3, 4, 5)? Since the only change is the tokenizer, this could work out of the box. (besides, arithmetic might be easier to learn in balanced base, because the addition and multiplication tables feature less carries).\n\nThe paper makes heavy use of techniques introduced in prior work, this is not a problem, but it would be fitting to have a proper \"related work\" section in the main, describing these methods and acknowledging anteriority. At present, these descriptions are scattered in sections 1 and 2. This could be done at the expense of the preliminaries section (sec. 2, which serves little purpose), and perhaps some of the parity results, which are underwhelming, compared to the rest of the experiments.\n\nBoth the scratchpad structure and the position encodings proposed in this paper are strongly dependent on the task. This suggests that the model can be trained to perform one operation (addition OR multiplication), but not calculations featuring both. This is a strong limitation. Have you experimented with such tasks (add and multiply)? Can they be learned zero or few shots? A discussion of these questions would be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uMsg5l5o5k", "reviewer": "Reviewer_3hhr", "summary": "In this paper, arithmetic transformers refer to transformers trained on arithmetic tasks.  The core idea of the paper is to combine position coupling and scratchpad, and the authors demonstrate that this allows for length extrapolation in arithmetic tasks.", "questions": "- Is there a way for the position coupling to be done “automatically” that you know of? I am unfamiliar with this specific line of work. \n- For PARITY, at least, I have seen a prompt construction that asks the model to number the bit that the model has generated as it generates the internal state (odd or even) given the prefix. \n  For example, given the sequence 1101. \n  ```\n    1. 1 odd\n    2. 1 even\n    3. 0 even\n    4. 1 odd\n\n    Final answer: odd.\n  ````\n  The intuition here being that the numbering allows the attention mechanism to more accurately pin-point the required informaton, which has a similar flavour to your approach.\n  I was wondering if you have attempted to prompt a model so that it generates the same positioning scheme (just in human readable form), and I wonder if that will improve multiplication generalisation.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "They show that by using a specific position embedding encoding (tri-level position coupling), they are able to also perform well on multiplication. This is an interesting observation for multiplication that when baked into the ‘structure’ of the transformer can give generalisation on arithmetic.", "weaknesses": "My first slightly minor nitpick is that the scratchpad requires scratchpad data generated in that particular format, for the multiplication specific task. In a general pretraining or fine-tuning setting, this could probably be added in as additional data to shape the model’s behaviour.\nHowever, the specificity of position coupling (for all variants of position coupling) means that it can’t be a general modification to standard Transformers in a pretraining setting, and have it improve arithmetic capabilities. To their credit, the authors have already addressed this in the Limitations section of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)