Transformers Can Do Arithmetic with the Right Embeddings
Authors: Sean Michael McLeish
Arpit Bansal
Alex Stein
Neel Jain
John Kirchenbauer
Brian R. Bartoldson
Bhavya Kailkhura
Abhinav Bhatele
Jonas Geiping
Avi Schwarzschild
Tom Goldstein
Venue: NeurIPS 2024 poster (2024)
Keywords: Algorithmic Generalization
Arithmetic
Transformer
Recurrent
Primary area: generative_models

Abstract: The poor performance of transformers on arithmetic tasks seems to stem in large part from their inability to keep track of the exact position of each digit inside of a large span of digits. We mend this problem by adding an embedding to each digit that encodes its position relative to the start of the number. In addition to the boost these embeddings provide on their own, we show that this fix enables architectural modifications such as input injection and recurrent layers to improve performance even further.

With positions resolved, we can study the logical extrapolation ability of transformers. Can they solve arithmetic problems that are larger and more complex than those in their training data? We find that training on only 20 digit numbers with a single GPU for one day, we can reach state-of-the-art performance, achieving up to 99% accuracy on 100 digit addition problems. Finally, we show that these gains in numeracy also unlock improvements on other multi-step reasoning tasks including sorting and multiplication.

Transformers Can Do Arithmetic with the Right Embeddings
Sean McLeish Arpit Bansal" Alex Stein', Neel Jain', John Kirchenbauer', Brian R. Bartoldson", Bhavya Kailkhura", Abhinav Bhatele', Jonas Geiping",
Avi Schwarzschild", Tom Goldstein University of Maryland, Lawrence Livermore National Laboratory, ELLIS Institute Tübingen, Max Planck Institute for Intelligent Systems, Tübingen AI Center, 4 Carnegie Mellon University
Abstract
The poor performance of transformers on arithmetic tasks seems to stem in large part from their inability to keep track the exact position of each digit inside of a large span of digits. We mend this problem by adding an embedding to each digit that encodes its position relative to the start of the number. In addition the boost these embeddings provide on their own, we show that this fix enables architectural modifications such as input injection and recurrent layers to improve performance even further.
With positions resolved, transformers. Can they solve arithmetic problems that are larger and more complex than those in their training data? We find that training only 20 digit numbers with a single GPU for one day, we can reach state-of-the-art performance, achieving up to 99% accuracy 100 digit addition problems. Finally, show that these gains in numeracy also unlock improvements on other multi-step reasoning tasks including sorting and multiplication.
study the logical extrapolation ability of
Positional Embedding:
Positional Embedding: Abacus
FIRE
Length of Operand One
100
75
Accuracy
50
50
25
100
50
100
50
100
Lenghh of Operand Two
Figure 1: Zero shot exact match accuracy on addition using depth sixteen transformer (decoder only) models trained on operands of up to 20 digits. Compared to state-of-the of-the-art embeddings (left). our new Abacus Embeddings (right) dramatically improve generalization to unseen digit lengths. The interior of the red square denotes the training distribution. Accuracies are averaged over three trials.
"Equal Contribution, correspondence to: smcleishOumd, edu, bansal01 @umd. edu.
"Code available on GitlHub: htttuu..com/meeeii//iiittteeee
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Introduction
Much of the recent work on Large Language Models (LLMs) focuses on their ability to solve problems in natural language and code generation. Despite progress in these domains, transformers still struggle to perform complex multi-step and algorithmic reasoning tasks in zero shot setting without resorting to tool use. To study algorithmic reasoning in a sterile laboratory setting, the academic community focuses on simple arithmetic test problems like addition. Addition is simple enough that modest-sized LLMs can (in principle) be trained from scratch to do it without running into capacity and training budget limitations, yet complex enough that even large industrial models fail on large numbers without a code interpreter [Loeber, 2024].
Training transformers for arithmetic enables to study several important questions. First, ask what architectural design choices, dataset characteristics, and training pipeline variants are required to learn many-step reasoning process like multi-digit addition? Going deeper, then investigate whether these models are capable of logical extrapolation can they solve problems of greater size and difficulty than those that appear in their training set?
Prior studies indicate that addition is hard for transformers [Lee et al., 2023, Shen et al., 2023, Zhou et al., 2023, 2024]. Our experiments indicate that this difficulty stems from their inability to clearly represent the exact position of a digit within a long sequence of digits. To address this problem, we propose simple modification the data representation that directly addresses this shortcoming. Our Abacus Embeddings are simple learned positional embeddings that are used to encode positions within each span of numerical tokens. Combining Abacus Embeddings and standard positional embeddings, we observe dramatic improvements in accuracy such that models trained with at most 20 digit operands can generalize to problems with 120 digit operands. This represents a state-of-the-an generalization factor of 6x, with the previous state of the art being only 2.5x. To the best of our knowledge, these the longest sequences which learned addition has ever been demonstrated.
We also study several other methods of improving arithmetic and generalization in transformers. We find that incorporating input injection-ske skip connections inserted between the input layer and each decoder layer can reduce generalization errors by 50% over the Abacus Embedding bascline. We also find that together with our embeddings looped transformer architectures, which contain recurrent layers in which the same parameters are re-used multiple times, can achieve pear-perfect generalization on addition problems consider.
Since our proposed methods solve large addition problems successfully, we evaluate whether the same approaches can be used to improve other kinds of algorithmic learning. We explore multiplication problems of up to 15 digit numbers and sorting over arrays of up to 10 numbers, making this the first study of extreme length generalization techniques for addition that transfer to other algorithmic tasks. Our contributions can be summarized as follows.
We propose a new positional embedding called Abacus Embeddings to better capture the significance of each digit, which leads to near- near-perfect in-distribution generalization.
We show that when combine Abacus Embeddings with input injection and looped transformers performance further improves, increasing from 92.9% to 99.1% in out of distribution accuracy, 87% reduction in error compared to using the embeddings with standard architectures alone.
We push length generalization beyond existing work and show that our models can solve problems with six times as many digits as the largest samples in the training set, whereas the previous state of the art is only two and a half times
We extend our findings to more complex problems including multiplication and sorting where we show length generalization in these domains.
2 Related Work
Arithmetic and Algorithmic Reasoning. Solving arithmetic with next token prediction is a difficult problem that attracts lot of attention [e.g. Saxton et al., 2019]. However, in zero-shot settings, even incredibly strong commercial API models struggle with very large addition problems (e.g. up to 100 digits) without access to tools. Among attempts to improve arithmetic performance of transformer-bau based models, reversing the digits so the arguments are written with the least significant
2

---

digit first is popular [Lee et al., 2023, Shen et al., 2023, Zhou et al., 2023, 2024]. Furthermore, changing the data format by adding explicit index characters improves model capability for addition [Zhou et al., 2023, 2024, Olsson et al., 2022]. Other work approaches arithmetic by embedding real numbers by scaling a single fixed aaken-embedding for numbers [Golkar et al., 2023]. Moreover, Dziri et al. [2023] show multiplication is hard problem for GPT-3 [Brown et al., 2020] even when finetuned on this task. Dziri et al. [2023] further show that GPT-4 [OpenAI, 2023] struggles to obtain high in-distribution accuracy on multiplication, even with scratchpad. However, Lee et al. [2023] find that with a detailed scratchpad, small transformers can perform multiplication in-distribution.
Arithmetic is a subset of the larger class of algorithmic reasoning problems that focus on the ability to learn and execute algorithms and generalize to longer problems [Anil et al., 2022b, Jelassi et al., 2023, Yang et al., 2023b, Veličković et al., 2022, Rodionov and Prokhorenkova, 2024, Testolin, 2024]. The more general algorithmic reasoning field includes work on various architectures and data modalities aimed at learning algorithms from data. Veličković et al. [2022] and Rodionov and Prokhorenkova [2024], for example, train neural networks to execute specific algorithmic tasks by training on input-output pairs well as intermediate steps and hints. In a similar vein and although initially appreciated for efficiency, weight sharing and recurrence can be used to make models adaptive and help generalize to harder problems [Dehghani et al., 2018, Sukhbaatar et al., 2019, Lan et al., 2020, Ibarz et al., 2022]. Schwarzschild et al. [2021] and Bansal et al. [2022] explore an end-to-end learning approach using recurrent convolutional neural networks to learn algorithms from input-output pairs, tackling algorithmic tasks like prefix sums, mazes, and chess. Weight sharing for algorithmic reasoning is also helpful with transformers and we use the looped transformer in some of our experiments below. A looped transformer has a transformer block called recurrently on its own output lending itself to executing iterative algorithms [Giannou et al., 2023, Yang al., 2023a, de Luca and Fountoulakis, 2024]. Additionally, recent work aims to improve reasoning in LLMs [Zhou al., 2023], but McLeish et al. [2024] demonstrate that LLMs, even with code interpreters, are less than perfect at algorithmic reasoning tasks, indicating a crucial need for advancements in our methodologies. This paper takes a step towards improving LLM arithmetic and algorithmic capabilities without tool use.
Positional Embeddings. Indicating the position of tokens in a sequence to transformer models is critical for language modeling [Vaswani et al., 2017]. Absolute positional embeddings (APE) are learned embeddings that are added to token embeddings before the first layer of the transformer [Vaswani et al., 2017]. However, these absolute embeddings inhibit length generalization [Press et al., 2022]. To address this issue. Shaw et al. [2018] propose relative embeddings (RPE) which are embed- ded during the attention computation, a mechanism further simplified by Raffel et al. [2020]. Others build on these works to improve length generalization including Sandwich [Chi et al., 2023]. Kerple [Chi al., 2022], and Alibi [Press al., 2022] positional embeddings. Additionally, Kazemnejad et al. [2023] show that decoder layers can still learn positional information with no explicit positional embeddings. No positional embeddings (NoPE) can achieve good length generalization performance for small algorithmic tasks and even outperform some specialized embeddings. Rotary Positional Embeddings (RoPE) [Su et al., 2024] are commonly used in state-of-the-art open source transformers le.g. Touvron et al., 2023]. However, RoPE does limit the length generalization as models are trained only using rotations based on training data length [Kazemnejad al., 2023, Press et al., 2022]. For improved length generalization, one can add post-training extensions [Peng et al., 2024]. The latest and most useful for arithmetic is Functional Interpolation for Relative Position Embeddings (FIRE) [Li et al., 2023]. FIRE shows the strongest length generalization to date, which leads to length generalization by 2.5x on addition [Zhou et al., 2024] when combined with randomized embeddings [Ruoss et al., 2023]. We go into more detail on some of these positional embeddings in Appendix 1. In this work, we focus on NoPE and FIRE embeddings since these are the best performers for addition in reversed format among existing embeddings [Zhou et al., 2024].
3 Achieving Length Generalization for Addition
We study a range of methods for improving the arithmetic capabilities of language models trained from scratch centering on two main hypotheses: (1) the positional information for individual digits within numbers is being lost and (2) recurrence can improve the reasoning abilities of transformer architectures on multi-step arithmetic reasoning problems. We briefly discuss the training and evaluation setup before describing each of our improvements in detail.
3

---

Least Significant Digit First: 1
2
3
4
+
2
3
4
4
6
8
Most Significant Digit First: 4
3
2
1
2
1
4
2
Abacus Embeddings: 1
2
8
4
0
1
2
3
4
0
1
2
3
4
Absolute Embeddings:
1 2
3
4
5
6
7
8
9
10
11
1R
15
14
Visualization of data formats and positional embeddings. Abacus Embeddings give the same positional embeddings to all digits of the same significance.
Figure 2:
Experimental Setup. We train decoder-only causal language models to solve addition problems. Following prior work [Zhou et al., 2023, 2024, Shen et al., 2023, Kazemnejad al., 2023, Lee et al., 2023], imputs are formatted least significant digit first, e.g. 98282 + 3859172 = 2787472. Unlike prior work, we do not add any padding between digits [Shen et al., 2023] and do not pad any numbers with zeros, neither in the case of carry digits [Zhou et al., 2024], nor to make all operands the same length [Shen al., 2023]. We train on all combinations of operand lengths less than or equal to and j where i and, j are the maximum lengths of the first and second operands, respectively. For this study all training sets have 20 million samples and = hence we can use one number to define the dataset i, where is the maximum length of either operand. We sample data with replacement and we stratify the data, so that all length pairs j) are equally sampled during training. To facilitate training of many models from scratch, we use language model cramming setup [Geiping and Goldstein, 2023] and limit each training run to 8 exaFLOP of compute (a single Nvidia RTXA0000 GPU for 24 hours); for multiplication results
allow 64 exaFLOP (eight Nvidia RTXA4000 GPUs for 24 hours). During training, we mask the input question and only compute loss on the answer digits. For further details on data construction and training we refer to Appendix A.2.
We report model accuracy for each (i.j) length pair and unlike most existing work, we also include accuracy for pairs where i * to highlight all instances of extrapolation. This extensive tabulation is costly and makes inference the main computational burden of this study, Since our training pipeline produces fairly consistent results, we report the mean over three runs (rather than using a best-of-ten reporting scheme [Zhou et al., 2024J). We measure accuracy in the strict sense where only exact matches of all output digits are counted as correct, i.e. if a single digit is incorrect then the example is marked as wrong and we refer to this as exact match accuracy. We have the following three evaluation (ID) where the models are tested problems up to the maximum size seen during training; (ii) out of distribution (OOD) where the models are tested on problems greater than the maximum size seen during training but both operands
categories: (i) in distribution
at most 100 digits; (iii) and extreme out of distribution (100+ digit OOD) where the models are tested on problems where both operands are of the same length and are both more than 100 digits and less than 160 digits. In the 100+ OOD setting, we only analyze problems where the operands are the same length (i j) due to inference costs this scale.
We consider two standard transformer architectures. First, we use a standard autoregressive trans- former model where multiple decoder layers stacked in feedforward manner. Second, we enhance this standard transformer model by incorporating input injection, where the embedded inputs added to the input of each decoder layer [Ma et al., 2022, Bansal et al., 2022, Anil et al., 2022a]. We visually describe the architectures in the Appendix Figure 22.
3.1 Abauus Embeddings Help Align Digits
From prior work and our own initial experiments, we observe that even when input numbers are presented least-significant digit first and training data is stratified and abundant (several million examples), standard transformers struggle to learn multi-digit addition. We also observe that humans do long addition by first aligning the digits of the same significance into columns. Thus, our first hypothesis is that the significance of each digit (i.c. each digit's 's position relative to the beginning of the number) is not easy for transformers to represent, and that this sub-problen presents more of a hurdle than the actual addition itself.
Prior work addresses this by proposing explicit index hints in the inputs and outputs of the addition, for example a6b7c5 + a1b6c3 = a7b3e9, finding that transformers perform much better on addition with the information provided by such hints [Zhou et al., 2023, 2024]. However, index hints of this form increase the input context length required and double the output length and inference cost of solving a given addition problem. Furthermore, Zhou al. [2024] find that the ability of models trained with index hints to generalize is sensitive to the particular random initialization. Zhou et al.

---

100
97.5
Exad Match Accurazy
30
24-1
40-
201
Bexe Maich
151
30
20-
wT
Archinecure Type FIRE, OUD FIRE, OOD
I1
Architeoture Type NOPE, OOD
Abaras, UOD
NoPk, OOD NoPE, 100H
Abacus, 100-
Figure 3: Left: Mean exact match accuracy of three models of depth sixteen on size 20 data, varying the architecture and embeddings. Abacus Embeddings improve accuracy for addition over FIRE and NoPE Embeddings. Right: Mcan exact match accuracy of three models of effective depth sixteen size 40 data, varying over NoPE FIRE embeddings and architectures. Recurrent looped transformer models improve accuracy for addition for both the FIRE and NoPE embeddings.
Looped transformer (LT): Weight tied decoder layers, with input injection and progressive loss. Standard Transformer (ST): Stacked decoder only layers. Standard Transformer with Input Injection (ST w/ II): Standard Transformer with input features added the hidden representation between each decoder layer.
[2024] highlight this by training models with different random seeds, varying weight initialization and data input order seeds, showing the variance in the performance of these models can vary from near perfect on 100 digit addition 0% accuracy at digit addition.
To address the limitations transformers at representing positional information, we design a specially built positional embedding that encodes the location of each digit relative to the start of the current number. We call this Abacus Embeddings. We apply the same positional embedding to all digits of the significance, providing an explicit signal that the model can align digits. We visually describe these embeddings in Figure 2.3
We take inspiration from Randomized Embeddings [Ruoss al., 2023] but instead of using random ascending indices to represent positions in a sample, we use consecutive ascending indices with a random starting position to allow for length generalization. Specifically, during training we give consecutive positional embeddings to each digit a number, starting from a randomly chosen offset value from [11, kt), where is hyperparameter. Unless otherwise stated the default value for k in this study is 100 and show this can be varied in Appendix A.5. For example, if the input is 123, the positional encodings learned embedding matrix. The value sampled from U[1, ki] is the same for all numbers a batch, significance obtain the same positional embedding. This training
3, 3 1, where 3 ~ [1, 100], which then passed through
meaning all digits of the scheme allows the model to see wide range of positional embeddings, even when training sequences are short. At test time, each positional embedding begins from one, i.c. B
Abacus Embeddings Solve Addition. Abacus Embeddings improve generalization performance up to 100 digits and beyond for standard transformer architectures. In Figure (left), we highlight the comparative boost Abacus Embeddings have standard transformer architectures and embeddings for performing addition, taking the mean accuracy of three models all cases. The accuracy results for the standard transformer models trained with FIRE and Abacus, tested both in-domain (ID) and out-of-domain (OOD), are also shown in Figure 1. Additionally, in Appendix A.6, we present similar 2D grid plots for several other experiments that are depicted bar charts in the main text. Zhou et al. [2024] find that operand lengths of up to forty digits are required during training for good generalization 100 digit addition during testing (albeit not robustly). We find that with our Abacus Embeddings, we can achieve similar accuracy and larger extrapolation using a standard model with input injection trained on maximum operand sizes of 20 digits.
As Abacus Embeddings are a variant of absolute positional embeddings, technically they cannot generalize beyond the relative positions seen during training. However the hyperparameter k that
'In Appendix A.3, we motivate these embeddings further with experiments demonstrating their utility in solving bitwise OR task.
5

---

randomizes the starting offset used for each individual addition example can be increased to enable generalization by training a larger range of embeddings for given computational budget. Relatedly, Appendix Figure 9 shows that training on larger datasets improves performance, even for operands with fewer than 100 digits.
3.2 Recurrence In Transformers Boosts Performance
With positional embeddings addressed, next we explore whether recurrent architectures can further improve the ability of transformers to perform multi-digit addition. We use the term recurrent block to refer to a set of decoder layers with distinct weights and recurrences to refer to the number of times the recurrent block is repeated. We use the term effective depth to mean the number of layers used in a transformer, whether their weights are unique or not. Unless otherwise stated, we use a maximally recurrent architecture, i.e. only one unique layer recurred achieve the effective depth. We also employ input injection, skip-connections that propagate copy of the input to each layer in the network.
The Benefits of Recurrence. In Figure (right), we compare all architecture variants using both FIRE and NoPE embeddings trained on addition over operands with up to 40 digits. Despite having approximately 10x fewer parameters than the other models, we see that the looped transformer (recurrent, with input injection and progressive loss), achieves the best out of distribution performance using either position embedding. In Figure 9 in the Appendix, we show this result is robust across multiple training data sizes.
With recurrent models, we can choose to vary the number of recurrences for each forward pass while training. This tends to improve generalization to harder tasks at test time and is also refered to as progressive loss computation [Bansal et al., 2022]. This loss function a convex combination the loss values from two forward passes, one with the nominal number of recurrences (so 16 for x 16 model) and one with a random smaller number of recurrences.
Next, we explore the effect of varying the size of the recurrent block while keeping the effective depth fixed. We perform this ablation by halving the number of layers in the recurrent block and doubling the number of recurrences, sweeping from model with sixteen layers in the block and single recurrence (16 1, i.c. a standard transformer), through to one layer the block but with sixteen recurrences (1 x 16). Analyzing these results in Figure 4, we show further performance improvements are possible in some cases with the combination of both recurrence and Abacus Embeddings. In particular, model with two recurrences (8 x 2) incurs half the error of the purely non-recurrent model (16 x 1) for OOD problems and enjoys increased accuracy on 100- OOD problems.
Finally, in Appendix A.7.3, we vary the effective depth of the models to analyze the impact of parameter count on this task, across Abacus, FIRE and NoPE embeddings. Although the experiments presented in Figure 4 are a fair comparison across depth, the purely standard transformer models have many more parameters than their recurrent counterparts. In Table the appendix, we record the parameter counts to the nearest million.
4 Pushing the Limits of Algorithmie Reasoning for Transformers
While there is an emphasis on addition as a difficult problem in existing work, our method's strong performance allows to extend even more difficult problems, including multiplication and sorting and even multiple operations at once.
4.1 Addition and Subtraction
We train models on dataset made up of an even mix of addition and subtraction samples. In Figure 5, we show results from models with 8 layers the recurrent block and recurrences trained with exactly the same hyperparameters used to train the addition models above. We see that these small transformer models can simultaneously learn to extrapolate for both the symmetric operation of addition and the anti-symmetric operation of subtraction using Abacus Embeddings.

---

11 960
97.
98.5
Exact Match Accuracy
100-
79. &
80
60
40
30.6
31.3
30.
2911
13.7
20
0
16x1
4x4 Layers in Recurrent Block X Number of Recurrences Abacus, OOD
8x2
2x8
1x16
Abacus, 100+ OOD
Figure 4: Varying the size of the recurrent block, while maintaining an effective depth of 16 and training on size 20 data. We see that a recurrent model with eight layers in the recurrent block and two recurrences the most accurate of all effective depth 16 models, halving the error rate of standard model with input injection in the OOD evaluation. (See Figure 17 for results with FIRE and NoPE.)
4.2 Integer Multiplication
We now study a harder task, multiplication natural numbers, where the length of the output may be the sum of the lengths of the operands. Compared to addition, where the output is at most digit more than the longest operand, multiplication has Longer-distance dependency and the output length scales much faster problem size increases.
To adapt from addition multiplication, we make some small changes to our set-up. First, we remove the input injection from inside the recurrent block and second, we divide the gradients in the recurrent block by the number of recurrences, boighwigg the gradient update from batches with many recurrences [Bansal et al., 2022]. (We analyze the impact of these design decisions for addition models in Appendix Figure 19.) We only examine looped transformers as the compute required for training and hyperparameter search for multiplication is far greater than for addition, limiting to much smaller scale analysis.
Abacus Embeddings help looped transformers reach near- pear-perfect accuracy in-distribution for mul- tiplication. In Figure 6, we show how the training distribution, surrounded by the red square fully saturates with Abacus Embeddings. In fact, models with our Abacus Embeddings achieve higher in
Exact Match Accuracy
100
90
Addition
80
Subtraction
In Distribution
60
50
40
30
20
100
110
120
Operand Length
Figure 5: Models which have layers in recurrent block and recurrences, trained on size 20 addition and subtraction data, each line is the average of 3 models. We see that it is possible to have extreme generalization whilst learning multiple tasks.

---

distribution accuracy on 15 digit multiplication than prior work [Shen et al., 2023] and do not require padding each operand the same length with zeros. In particular, we highlight that the specific problems that models trained with FIRE embeddings struggle to solve are the hardest problems in the training set and Abacus Embeddings outperform them in this key area (see the lower right corner of the red boxes in Figure 6)
4.3 Array Sorting
While both addition and multiplication accept only two operands, we now analyze the task of sorting arrays of multiple variable length num- bers, a more challenging testbed for evaluat- ing the generalization abilities of our Abacus Embeddings. We present each sorting problem using alphabetical indices for each (reversed) number in an input array where the expected output is the alphabetical indices in ascending order. For example,
Table 1: Exact match accuracy for sorting with various positional embeddings. All results are per- centages of the test set and all models here are standard transformers with eight layers.
FIRE
Abacus
Abacus + FIRE
OOD (number length 30)
55.32
68.63
67.28
OOD (urray length 30)
21.35
9.67
21.11
All OOD (30 30) All OOD
3.73
2.65
4,48
64957, 6 99963, 10218, d 7141, e 05781 d, e, b, a, c. We train with arrays of up to 10 numbers each hav- ing up 10 digits and then evaluate with arrays of up to 30 numbers each having up to 30 digits. We give more detail on the sorting data construc- tion process in Appendix A.2.
14.65
9.78
16.91
Table 2: Accuracy for sorting with various architec- tures for sorting. ST denotes standard transformer, ST w/ II denotes standard transformer with input injection, and I.T denotes looped transformer mod- els. The standard transformer has the best exact match accuracy. When measuring the accuracy on identifying only the minimum element of the array, looped transformers outperform All results are percentages of the test set.
In this setting, explore ization. First, we increase the maximum pos- sible length of the input numbers 30 digits while maintaining the maximum array length to 10 and refer to this scenario as "OOD (number length- - 30)." Second, we increase the number of inputs in the array to be sorted 30 while keeping the maximum digit length of each num- ber at 10 and term this scenario "OOD (array length 30)." Finally, we consider scenario where both axes are increased simultaneously, referred to as "all OOD."
axes of general-
others.
ST
ST w/ II
LT
All OOD (exact string match)
4.48
3.84
2.60
All OOD (min. clem. only)
49.73
60.09
68.51
In Table 1, we illustrate the performance of a standard transformer (eight layers) trained with different embeddings- Abacus, and their combination. Again, our results demonstrate that the combined embedding approach enhances the model's 's ability to generalize, surpassing the performance of either
Abacus
Abacus + FIRE
FIRE
100
Length of Operand One
5
Accuracy
10
50
15
20.
0
10
15
20
0
5
10
15 20
5
10
15
20
Length of Operand Two
Figure 6: Exact match accuracy of looped transformer models trained multiplication, with four layers in the recurrent block and four recurrences. The red square denotes in distribution testing on up 15 digit operands. We the models with Abacus Embeddings achieve near perfect in distribution accuracy. Combining Abacus Embeddings with FIRE also improves in distribution accuracy on the hardest in distribution problems (bottom right), comparing to the FIRE-only baseline.

---

embedding alone in the "all OOD" setting. However, in Table 2, we observe mixed results when pairing the Abacus+FIRE Embeddings combination with different model architectures with effective depth eight. For sorting, different architectures appear to be better suited to different types of extrapolation, for example the looped transformer is best at extrapolating for finding the minimum element but not for sorting the whole array.
Overall, the superior sorting performance of the Abacus Embeddings underscores their potential udiltty across broader spectrum of algorithmic tasks beyond basic arithmetic. Abacus Embeddings may be instrumental in use cases requiring transformer models to perform variety of complex positional, numerical, and/or relational reasoning tasks.
4.4 Abacus and Relative Embeddings
As Abacus Embeddings are only applied to numbers, to incorporate Abacus Embeddings into a general purpose model, they must be compatible with other relative embeddings maintain good downstream performance on non-arinn -rnthhitheetic tasks. We examine these types of combinations here and conclude that Abacus Embeddings complement techniques that are good for natural language well, suggesting that these combinations could be powerful for large-scale general models.
Although Abacus Embeddings implicitly combined with NoPE (no positional embeddings) embeddings for all experiments seen so far, most state-of-the-art open source models use Rotary Embeddings. Rotary Embeddings are weak for length generalization. We show that combining Abacus Embeddings with RoPE does, in fact, yield improvement in operand length generalization. However, in Figure 7, demonstrate the potential for integrating Abacus Embeddings into more general system, showing that the combination of Abacus Embeddings with FIRE unlocks generalization well beyond the problems that FIRE embeddings can solve on their own.
ST w/ II
ST w/ II FIRE
ST w/ II Abacus RoPE
ST w/ II RoPE
Length O'eerndd One
Abacus FIRE
100
Accuracy
50-
50
100
50
100
100
$0
100 0
50
100
Length of Operand Two
Figure 7: Exact match accuracy of standard transformer of depth 16 with input injection, trained on up to size 20 data. The red square denotes in distribution testing. Combining Abacus Embeddings with FIRE or RoPE embeddings improves out of distribution accuracy for addition, over the baseline models without Abacus Embeddings.
5 Discussion & Limitations
While the capabilities of LLMs have advanced far enough to encompass complex tasks including code generation and mathematical reasoning, stress testing the limits of these models remains challenge. In this paper, we study mathematical reasoning tasks including addition, multiplication, and sorting to evaluate these capabilities in controlled setting. We analyze the ability of specialized language models to learn algorithmic tasks in zero shot setting, without access outside tools like code interpreters, etc., exploring the benefits of various architectural improvements like improved embeddings and recurrent layers.
Across our experiments, we find that our novel Abacus Embeddings improve performance dra- matically both when applied to standard transformers well as recurrent variants. We repeatedly achieve length generalizations at least 6x (capped by the context length) more than doubling the extrapolation demonstrations in prior work, achieving near perfect results on addition of up 100 digits, with repeatable results across multiple training runs. We demonstrate the the complementary properties of our Abacus Embeddings with other relative embeddings like FIRE, achieving dramatic

---

improvements in in-distribution multiplication performance, and making headway on the challenging problem of variable length array sorting.
Contrasting with prior work, our experiments explore types of extrapolation well beyond just length generalization for addition, presenting an architecture modification that improves performance on multiple algorithmic reasoning tasks simultancously. We hope that our work deepens the community's s understanding of these problems and paves the way for further advancements in the algorithmic reasoning capabilities of large language models.
Limitations There are some intrinsic limitations that accompany any study involving language model training from scratch under compute constraints. However, the primary point of relevance for this study is that although we show the compatibility of Abacus Embeddings with FIRE and RoPE embeddings, we do not actually explore any natural language tasks. In the future, a larger scale study including natural language would be needed to understand further how Abacus Embeddings would perform on heterogeneous tasks comprising both numerical and natural language inputs.
Acknowledgments and Disclosure of Funding
This work was made possible by the ONR MURI program and the AFOSR MURI program. Com- mercial support was provided by Capital One Bank, the Amazon Research Award program, and Open Philanthropy. Further support was provided by the National Science Foundation (1S-22222182), and by the NSF TRAILS Institute (2229885). Computing resources were furnished by the Department of Energy INCITE Allocation Program, and Lawrence Livermore National Labs.
Furthermore, this work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA227221 and was supported by the LLNL-LDRD Program under Project No. 24-ERD-010 (LLNL-CONF-2000155)
References
Cem Anil, Ashwini Pokle, Kaiqu Liang, Johannes Treutlein, Yuhuai Wu, Shaojie Bai, J Zico Kolter, and Roger B Grosse. Path independent equilibrium models can better exploit test-time computation. Advances in Neural Information Processing Systems, 35:7796-7809 2022a.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowyez, 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-38555 2022b.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1607 2016.
Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Logical extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Ta-Chung Chi, Ting-Han Fan, Peter Ramadge, and Alexander Rudnicky. Kerple: Kernelized relative positional embedding for length extrapolation. In Advances in Neural Information Processing Systems, 2022.
Ta-Chung Chi, Ting-Han Fan, Alexander Rudnicky, and Peter Ramadge. Dissecting transformer length extrapolation via the lens of receptive field analysis. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), pages 13522- 13537, 2023.
Artur Back de Luca and Kimon Fountoulakis. Simulation of graph algorithms with looped transform- ers. arXiv preprint arXiv:2402.01107, 2024.
10

---

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations. 2018.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, et al. Faith and fate: Limits of transformers on compositionality. arXiv preprint arXiv:2305.1860 2023.
Jonas Geiping and Tom Goldstein. Cramming: Training a language model on a single gpu in one day. In International Conference on Machine Learning, pages 11117-11143, PMLR, 2023.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, pages 1 11398-11 11 1442. PMLR, 2023.
Siavash Golkar, Mariel Pettee, Michael Fickenberg, Alberto Bietti, Miles Cranmer, Geraud Krawezik, Francois Lanusse, Michael McCabe, Ruben Ohana, Liam Parker, et al. xval: A continuous number encoding for large language models. arXiv preprint arXiv:2310.022 10.09999, 2023.
Borja Ibarz, Vitaly Kurin, George Papamakarios, Kyriacos Nikiforou, Mchdi Bennami, Róbert Csordás, Andrew Joseph Dudzik, Matko Bošnjak, Alex Vitvitskyi, Yulia Rubanova, al. A generalist neural algorithmic learner. In Learning on graphs conference, pages 2-1. PMLR, 2022.
Samy Jelassi, Stéphane d' Ascoli, Carles Domingo-Enrich Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306.1540 15400, 2023.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization transformers. arXiv preprint arXiv:2305. 1..30551946 19466, 2023.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020. URL https: http:///ppeeeeeeeeeeeeeeooeee//e H1eA7AEtvS.
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.03381, 2023.
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. arXiv preprint arXiv:2310004418, 2023.
John Loeber. #16: Notes on Arithmetic in GPT-4, February 2024. URL https:// /oebber.
substack. wttcc//o////// p/ p/16-notes: on- arithmeticc in--p-t--....
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209. 10655, 2022.
Sean McLeish, Avi Schwarzschild, and Tom Goldstein. Benchmarking chatgpt algorithmic reasoning. arXiv preprint arXiv:2404.0344), 2024.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv: 11895, 2022.
OpenAl. Gpt-4 technical report. ArXiv, abs/2303.08774, semanticscholar. .ow. wrt.:/orurrrg/ooouuIDDD..I :12.ppu21::255528155 257532815.
2023.
URL https://api.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yamn Efficient context window extension of large language models. International Conference on Learning Representations, 2024.

---

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
Jing Qian, Hong Wang, Zekun Li, Shiyang Li, and Xifeng Yan. Limitations of language models in arithmetic and symbolic induction. arXiv preprint arXiv:2208.0500 2022.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 21(140):1-61 2020.
Gleb Rodionov and Liudmila Prokhorenkova. Discrete neural algorithmic reasoning. arXiv preprint arXiv:2402./1668, 2024.
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. 16843, 2023.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. arXiv preprint urXiv:/9044 01557, 2019.
Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks. Advances in Neural Information Processing Systems, 34, 2021.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint arXiv: 10X3:0/800302/55: 2018.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202. 2020.
Ruoqi Shen, Séhastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Positional description matters for transformers arithmetic. arXiv preprint arXiv:2311 1.123/111.77,,,. 2023.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiig, 568: 127063, 2024.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention
span in transformers. In Anna Korhonen, David Traum, and Lluís Mårquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 331 331-335, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.1 10.18653/v1/P19/ 1032. URL https: / /aaaaaacaannttooooo/e otg/p19-1133
Alberto Testolin. Can neural networks do arithmetic? a survey on the elementary numerical skills of state-of-the-art deep learning models. Applied Sciences, 14(2), 2024. ISSN 2076-3417. doi: 1..3300/app140220774 URL https: //ww mddi. /ttdpp..cmmm00-3347////////1/ 3417/14/2/7444
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXhv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Petar Veličković, Adrià Puigdomènech Badia, David Budden, Razvan Pascanu, Andrea Banino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The clrs algorithmic reasoning benchmark. In International Conference on Machine Learning, pages 22084 22102.. PMLR, 2022.
Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang. Dongdong Zhang, and Furu Wei. DeepNet: Scaling Transformers to 1,000 Layers. arXiv:2203.00555 [cs], March 2022. URL hrggabbs/2255 00555.
Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. arXiv preprint arXiv:2311.12424 2023a.
12

---

Zhen Yang, Ming Ding. Qingsong Lv, Zhihuan Jiang, Zehai He, Yuyi Guo, Jinfeng Bai, and Jie Tang. Gpt can solve mathematical problems without calculator. arXiv preprint #xXiv2239.03224. 2023b.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12104-12113, 2022.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization. arXiv preprini arXiv:2310.16022 2023.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:240.003371, 2024.
13

---

A Appendix
Author Contributions
Sean McLeish* - Led the project, developed the idea, contributed to code, organized majority of experiments and contributed to writing.
Arpit Bansal* - Contributed large amount the idea, contributed to code, organized experiments for sorting arrays, and contributed to writing.
Alex Stein Contributed to code, and helped plan the experiments.
Neel Jain - Contributed large amount to writing, and helped plan the experiments. John Kirchenbauer Contributed large amount writing.
Brian R. Bartoldson, Bhavya Kailkhura - Helped set up large scale parallel addition evaluations, contributed to writing.
Abhinav Bhatele Contributed to writing.
Jonas Geiping, Avi Schwarzschild, Tom Goldstein Developed the idea, helped plan and organize the experiments, and contributed large amount to the writing.
A.1 Extended Related Works
A.1.1 Positional Embeddings.
FIRE embeddings additive embeddings in the attention mechanism: AnrE(X)
XWQ(XWx)" + B where Bu,- et al. [2023] show empirically that these embeddings allow for length generalization and theoretically show they are capable of representing many other embedding types. Ruoss et al. [2023] propose using a random subset of a larger set of possible positions during training so that larger positional
logfe((-1)+11 and L learnable parameters. Li [og(emax((.2)+1)
embeddings trained. Zhou al. [2024] use randomized FIRE [Ruoss et al., 2023, Li al., 2023]
embeddings to achieve length generalization on arithmetic tasks, which use randomized positions as input to the small multi layer perceptron used in FIRE embeddings.
A.2 Datasets
Addition: We sample equally, with replacement, from all possible operand lengths up to maximum dataset size of 20 million, we call this a dataset of size in the main text. For evaluation we sample 100 samples for each pair of operand lengths evaluated.
The input for this problem is two binary vectors, the longer input vector is all zeros and the shorter input contains a one. The output should be the length of the longer vector with the one in the same position as in the shorter vector. If the inputs are the same length, the one can be in either vector. E.g. 001 @ 00000 00100. For training, we exhaustively sample the space of all vectors of sizes less than or equal to the predefined maximum input vector size.
Bitwise OR:
Sorting: Given list of reversed integers indexed by characters, output the characters in ascending
order E.g. a 64957, b: 99963, 10218, d 7141, e : 05781 = d, b, a, We implement the grid like manor. We query each "square" of [1. [1, n] grid sampling process for sorting until the maximum size has been reached for the dataset. When querying "square" (i,j) we randomly sample integers of size less than or equal to j digits. We randomly sample consecutive indices for the natural numbers in our list at both train and test time.
We implement the multiplication datasets for both training and testing the exact same manor as for addition, only changing the operation used to calculate the answer.
Multiplication:
A.3 Bitwise OR on Binary Vectors
A necessary condition to perform addition is aligning digits of the same significance. We begin by examining positional embeddings for exactly this task. To this we analyze the bitwise OR task, where the model has to output left aligned position wise OR of two binary vectors. We present samples from the dataset in Section 1, these are left aligned to be representative of the task of aligning digits for reversed addition.
14

---

Exact Match Accuracy
100
83.1
71.8
80
55.6
60
40
20
10.5
7.1
3.7
6.8
6.4
4.7
LT
ST
ST w/
Architecture Type
Abacus, OOD
FIRE, OOD
NoPE, OOD
Figure 8: Accuracy of models on the bitwise OR task when trained on data with size up to 20, varying different positional embeddings and architectures. Abacus Embeddings heavily improve performance on this task.
We train standard transformer, standard transformer with input injection and looped transformer models on the position wise or task, on a dataset where the maximum length of either input vector is twenty. This result is shown in Figure 8. Here to generalize further on this task than the other embeddings which prior work for addition focuses on. As with addition, we that looped transformers perform better than the standard architectures with FIRE or NoPE embeddings. We do note that these accuracies are not as high we report for addition. We hypothesize this is because the model is having to repeatedly predict the same token multiple times, this has been thought to be the cause of errors in prior addition work[Qian [iian et al., 2022]. When analyzed the errors in this task found they were predominantly caused by the model outputting one too few or too many zeros.
that the Abacus Embeddings allow all models
A.3.1 Example Data
000010 00000000000000 00001000000000
000100 @ 0000000 0001000
001 00000 00100
A.4 Addition Models Trained on Varying Data Sizes
Across Figure 9, we see that increasing the size of the operands in the training set allows for better generalization above one hundred digits for all models. This partially due the sampling method for training Abacus Embeddings. As the offset randomization hyperparameter k = 100 is fixed across experiments, there are more embeddings trained if the operands seen during training are longer. The size of the OOD set below 100 is reduced as the size of the operands seen during training increases, as the ID category now includes this data. However, this does still show that the size of the operands seen during training directly impacts the generalization, with larger training sizes allowing for better generalization.
A.5 Extreme Length Generalization for Addition
Absolute positional embeddings must be learned during training otherwise they are unusable at test time. This limits our Abacus Embeddings which are trained with the offset randomization hyperparameter & = 100. One possible way resolve this generalization problem is to increase the value of during testing. Figure 10 (left), we show the exact match accuracy of five looped transformer models, with eight layers in the recurrent block and two recurrences trained on size 20
15

---

Train Data Size: 10
Train Data Size: 20
Exact Match Accurney
92.9
97.9
100
79.8
68.9
50
35.88
21.7
50
26.7
30.6
13.7
3.6
4.3
LT
ST
ST w/ II
LT
ST
STw/ II
Architecture Type
Architecture Type
Train Data Size: 30
Train Data Size: 40
Exact Match Accuracy
8
96.
99.6
99.4
990
99.9
99.3
100
100
60.4
64.1
62.5
48.0
50
32.5
50-
24.3
16.1
15.8
23
15.3
11.4
11.0
6.2
7.6
LT
ST
ST w/ II
LT
ST
STw/ w/ II
Architecture Type
Architecture Type
Abacus, OOD
FIRE, OOD
NoPE, OOD
Abucus, 1001 OOD
FIRE, 100 OOD
NoPE, 100 OOD
Figure Mean exact match accuracy of three models effective depth sixteen, varying the training data and architecture. We omit from the plot the in distribution accuracies as these are all 100% or very close to 100% for all models, this can be verified by the dark blue inside of all the red squares in Section A.6. Models trained on larger operands achieve higher OOD accuracy.
data with Abacus Embeddings and & 101, generalizing 120 digit addition. We only show the accuracy for operands of the same length in Figure 10 (left). seeing these models consistently achieve accuracies of 95% and above. We see this across the paper this method is much more robust than that presented by Zhou et al. [2024].
Accunicy
100
90
Run
Expact Match Accurscy
100
90
k-100
80
70-
60
Run
k-50
50
Run
1-25
40
Run
Distribution
30
Distrbuttio
20
10-
110
120
110
Operand Length
Operand Length
Exact match accuracy of five models trained size 20 data, generalizing well to 120 digit addition, an extrapolation 6x. Right: Exact match accuracy of five models trained on size 20 data, offset randomization hyperparameter Only showing the accuracy for operands of the same length.
Figure 10: Left:
25, 50, 75 and 100.
In Figures 10 (right) and 11 we continue varying the maximal offset randomization hyperparameter and size of the numbers in the training data. In Figure 10 (right), we show that varying the maximal offset randomization hyperparameter (k) changes the amount of extrapolation as we increase as expected, This allows to generalize operands over googol. In Figure 11 we show models trained on size and 40 data with larger values for k, maximum 6.8x length generalization from training. We see the models struggle to use the largest embeddings, e.g. embedding 214 in Figure 11 (right), this due to the stochastic training of embeddings, meaning the very largest embeddings are updated infrequently. This can be remedied by longer training but to remain consistent with other results we only train for 24 hours on a single A4000. Hence, we can easily increase larger values and perform arithmetic with far more digits, with suitable training data.
100,
16

---

Match Accuracy
100
90
Euace Match Accuracy
100
50
k=175
80-
50
1-150
70-
60
k-125
k-125
in Distributoo
Ic Distribttoo
30
20
160
180
200
160
180
200
Operand Length
Operand Length
Figure 11: Left: Exact match accuracy of five models trained size 30 data, offset randomization
hyperparameter = 125, 150 and 175. Right: Exact match accuracy of five models trained on size 40 data, offset randomization hyperparameter k 125, 150 and 175.
Only showing the accuracy for operands of the same length. These results are from models which have 8 layers in recurrent block and 2 recurrences and are trained on size 30 data with varying each line is the average of 3 models.
A.6 Addition Full 100 100 Plots
Here present the mean accuracy heatmaps for the main addition experiments shown throughout the paper. Figure 12 (left) corresponds to Top Left of Figure Figure 12 (right) corresponds to Top Right of Figure 9 and Left of Figure 3. Figure 13 (left) corresponds to Bottom Left Figure 9. Figure 13 (right) corresponds to Bottom Right Figure 9 and Right of Figure 3. Figure 14 corresponds to Figures 4 and 17. All of these figures show the Abacus Embeddings ability generalize in both dimensions of the addition problem.
L,T,
LT,
IT,
LT, NoPE
100
100
50
100
100
ST.Abacus
ST, FIRE
ST. NoPE
Lengih of Opernd One
ST. A.Abaauu
ST. FIRE
ST. NoPE
Lengh Operand
40
100
100
ST w/ IL, Abacus
ST w/ I1, FIRE
ST w/ II, NoPE
STT w/ IL, Abacus
ST w/ I1, FIRE
ST w/ 11, NoPE
20
20
100
100
100
100
100
Length Opperadd
Length Operand
Figure 12: Full 100 x 100 exact match accuracy plots, taking the mean over three models. Left: Size 10 training data, corresponding to Top Left of Figure 9; Right: Size 20 training data, corresponding to Top Right of Figure 9 and Left of Figure 3.
A.7 Addition Ablations
A.7.1 Analyzing the Intermediate Properties of Recurrence
Thanks the looped transformer architecture, we can extract intermediate solutions from the models, allowing us to plot the models outputs over iterations of the recurrent block. We present an example Figure 15 and suggest that this level of interpretability could be leveraged future work. The model presented is a x 16 model, one decoder layer and sixteen recurrences. We do not show the full 16 iterations in this plot for readability but these models do maintain fixed point iterations and beyond.
17

---

LT, Abacus
FIRE
LT. NoPE
0
LT. Abacus
LT, FIRE
LT. NoPE
100
100
100
80
100
80
ST.Abacus
ST. FIRE
SL. NoPE
Length of Operand One
ST. Abacus
ST. FIRE
ST. NoPE
Length of Operand
60
50
40
100
ST w/ Abacus
ST w/ II, FIRE
ST wI NoPE
ST II, Abacus
FIRE
w/ NoPE
50
50
100
100
100
50
100
100
30
100 0
50
100
50
Length of Operand Two
Length of Operand Two
Figure 13: Full 100 100 exact match accuracy plots, taking the mean over three models. Left: Size 30 training data, corresponding to Bottom Left Figure 9; Right: Size 40 training data, corresponding Bottom Right Figure and Right of Figure 3.
A.7.2 Removing Masking Before Equals
We mask all tokens before the equals sign all of our experiments, we hypothesize that with more training time this constraint may be able to be removed. In Figure 16, we show the effect of training with the same amount of flops the other addition experiments without masking before the equals sign.
A.7.3 Varying Effective Depth
We begin in Figure 17 by showing replica of Figure 4, this time including comparisons FIRE and NoPE embeddings. Seeing, yet again, the improvements Abacus Embeddings give for addition.
In Figure 18, we present models with effective depths 8 and more than 16, respectively. In Figure 18 (left), we see that the effective depth 8 models under perform the models with 8 layers in the recurrent block and two recurrences shown in Figure 4, demonstrating the benefit of recurrence in this case. We see very high accuracy from all models in Figure 18 (right). Again, the depth 32 recurrent models outperform the standard models with input injection, a quarter of the parameters and achieves the highest OOD mean accuracy of all models presented. These ablations show that with Abacus Embeddings the addition task can be learned across many effective depths to varying degrees of accuracy.
though it only has approximately
In Figure 19 (left), we remove the input injection to the intermediate layers in the recurrent block, only keeping input injection to first layer of the recurrent block. In Figure 19 (right) we divide the gradients in the recurrent block by the number of recurrences for the looped transformer models during training. We see very minor performance changes for all models shown in Figure 19, with the 2 x 8 model improving its performance slightly in left plot and the x 4 model improving slightly in the right plot. We ablate this design choices as we have to remove the input injection inside of the recurrent and divide the gradients in the recurrent block by the number of recurrences for the multiplication models show in Figure 6. Hence, we can conclude there would only be very minor performance changes in this case for addition.
A.7.4 Adding randomized Padding
Abacus Embeddings give strong priors for numerical tasks but without them, looped transformers perform better than the standard transformer architectures we present. The result shown in Figure 20 aligns well with the hypothesis that with fewer priors the looped transformer models are able to generalize better. In this case the priors reduced as the training data is noised with random pad
18

---

16x1, Abacus
16x1, FIRE
16x1, NoPE
0
50
100
100
8x2, Abacus
8x2, FIRE
8x2, NoPE
0
80
50
100
Length of Operand One
4x4, Abacus
4x4, FIRE
4x4, NoPE
0
60
50
100
0
2x8, Abacus
2x8, FIRE
2x8, NoPE
40
50
100
20
1x16, Abacus
1x16, FIRE
1x16, NoPE
0
50
100
50
100 0
50
100 0
50
100
Length of Operand Two
Figure 14: Full 100x100 exact match accuracy plots, taking the mean over three models, relating to Figures 4 and 17.
symbols, a method which was shown to improve length generalization in prior work [Shen al., 2023].
A.7.5 Index Hints
Zhou al. [2023] "randomly sample consecutive index hints from a pre-defined ordered set of hints with 102 symbols," for example a667c5 + a166c3 = a763c9. We implement this method two ways. Firstly, cyclic, here we treat the list as cyclic when sampling. Secondly, non-cyclic, this reduces the number of samples which receive the embeddings later in the ordering as we only sample from the list order. We see similar results for models trained on up to twenty digits as Zhou al. [2023]. We do note that our format of taking the mean exact match accuracy does highlight robustness as one of the three models tested were to not generalize well, this would impact reported accuracy heavily. We only show comparison to size 20 training data due to the increased cost of evaluating these index hint models, as the inputs and outputs are approximately double the length of regular questions the inference time is heavily increased. Due to the robustness issues highlighted by Zhou
19

---

Intermediate Outputs Over Recurrences
Loc:
Loc: guesse
Loc:
Loc: 98 gresso
99 guessi
Loc: 100
gueesss9
paweospppoeeepe preeejete
Octput Token
Fiernntons
fierations Loct 1 pu0ces1
Fiernsions
Biernnions
Eierniions
Dremitions
Loct: 1
Loct: 2
Lace 98
Lace 999
303
ans: gucce:0
1704
Probsbility
40%
605
601
60%
601
50%
374
A0%
Diccusions
Nerrations
lierstions
Nerratinns
leeritions
Htrratinns
Figure 15: Plot showing the improvement of the prediction over "thinking" iterations on 100 digit addition problem.
Input Prompt:
Answer: 00560022255888140000000777771110
586064772262384957011118887
(Note that the plot is truncated.)
Exact Match Accuracy
100
84.9
80
$3..0
60-
40
20
0.8
0.7
1.3
3.5
2.1
LT
ST
ST w/ II
Architecture Type
Abacus, OOD
FIRE, OOD
NoPE, OOD
Figure 16: Effect of removing the masking of the loss before the "=" sign in the addition task. All models perform worse when trained for 24 hours on single Nvidia RTXA4000 if we do not mask the input question in the loss function.
al. [2024] with their methods, we try to the best of our abilities faithfully reproduce their work within our experimental set up, noting that perhaps better random seed or initialization may be able produce better results for these models.
A.8 Additional Experimental Information
In this work, we consider three different model types, the classical standard transformer, standard transformer with input injection, and looped transformers. We visually describe these in Figure 22.
20

---

I 6666
97.9
86
97.3
Exact Match Accuracy
100
799
80
60
40
30.6
31.3
30.1
29.1
13.7
20
3.66
3.7
5.5
4.8
12.8
0
0
16x1
8x2
4x4
2x8
1x16
Layers in Recurrent Block X Number of Recurrences
Abacus, OOD
FIRE, OOD
NoPE, OOD
Abacus, 100+ OOD
FIRE, 100+ OOD
NoPE, 100+ OOD
Figure 17: Continuation of Figure 4, including FIRE and NoPE embeddings. We the Abacus
Embeddings perform best for all models.
EaachMath Acunay
Mashh AuucMMaah.commm Aconcyy
80
60
40
Architauure Type FIRE, L00 OCD
Recurances Block
Abacus
FIRE, FIRE, 100+ OCD NoPE, 100 00D
Abocus, 100- OCD
NoFE, 100 OOD
Abacus, 100- DOD
Figure 18: Left: Effective depth models, trained on size 20 data. These models under perform the models with eight layers in the recurrent block and two recurrences shown in Figure showing the benefit of recurrence for addition. Right: Effective depth > 16 models, trained on size 20 data. The models contain many more parameters than all other models we present, showing more that an effective depth of more than 16 does not necessarily improve accuracy in this setting.
97.1
68.3
08.3
Accuracy
100
Diaeet Match Accuracy
80
60+
60-
263
239.
313
20
20
8x2
4x4
2x8
8x2
4x4 Number
2x8
Layers
Recurrences
Layers
Abacus,
Abacus,
Abacus,
Replicas of the looped transformer models shown in Figure 4, to check the modifications train addition models do not adversarially impact addition training, taking the mean of three models in each case. Left: without the input injection to the layers inside of the recurrent block, only the first layer of the recurrent block. Right: dividing the gradients in the recurrent block by the number of recurrences.
Figure 19:
Due to the looped transformer architecture the number of recurrences at train time can be different to the number of recurrences at test time, although do not make use of this in this work.
21

---

Exact Match Accuracy
100
72.8
80
60
40
20
1.5
3.6
3.3
2.9
3.4
LT
ST
ST w/ II
Architecture Type
Abacus, OOD
FIRE, OOD
NoPE, OOD
Figure 20: Effect of adding randomized padding into training data only for the addition task. Looped transformer models are able maintain high accuracy when random padding is added into the data.
Exact Match Accuracy
100
80
60
40
13.8
14.6
12.3
20
6.0
9.9
LT
ST
ST w/ II
Architecture Type
FIRE Rand, OOD
FIRE Rand Circular Hints, OOD
Figure 21: Using index hints and randomized FIRE embeddings, presented by Zhou et al. [2024], training on size 20 data with our methodology, such as masking before the equals sign. This would be comparable to "1 to 20" in Figure 13 presented by Zhou et al. [2024] and Figure 3 of our work.
Architectures
Standard
Standard
Looped
With Input Injection
33
+
"1 + 22
Figure 22: Visualization the three architectures we study.
As Abacus Embeddings are a variant of absolute embeddings, reused only for numbers, they could be combined with relative embeddings being deployed current models. If digits input to model are tokenized individually, we can perform a linear time operation to find and assign relative embeddings to all numbers an input, which lower than the quadratic cost incurred by attention. Training a small number of Abacus Embeddings may be enough to handle all numerical inputs for
22

---

Table 3: Number of parameters, to the nearest million, in model with Abacus Embeddings and input injection.
Layers in Recurrent Block
Recurrences
Parameters (Millions)
16
1
122
8
2
64
4
4
34
2
8
19
16
12
Table 4: Default number of Nvidia GPU hours used to train a model.
Dataset
Number of GPU Hours (training)
Number of GPU Hours (testing)
Addition
24 RTXA4000
65.8 V100
Bitwise OR
RTXA4000
45 V100
Sorting
24 RTXA4000
64 RTXA4000
Multiplication
192 RTXA4000
0.83 RTXA4000
addition as they are reused. To fully implement our methodology all numbers also have to be reversed, this can be implemented with simple regular expressions on all inputs and outputs.
We use a character level tokenizer for all experiments and greedy decoding in all testing. We train all models with local batch size which is the maximum batch size that is power of two that will fit into the sixteen gigabytes of GPU memory. For multiplication models we first take the mean loss across samples before taking the mean across all samples in a batch, instead of taking the mean loss across all token in batch; we find this leads to slightly more stable training. We note that training models solve multiplication requires more hyperparameter tuning than addition, perhaps implying is trickier task to learn. Also, FIRE models require much greater compute budget for hyperparameter search as compared Abacus models for multiplication. In Table 3, we present the approximate parameter counts for models trained with input injection and Abacus Embeddings.
Compute Usage. We detail the default use of GPUs for each experiment in Table For some experiments, such extreme length generalization (Figure 10) and index hints (Figure 21) more GPU hours are required for testing, these are included the total number of GPU hours used. Our testing pipeline for addition and Bitise OR uses Nvidia V100 GPUs. Due to technical problem, 'torch.compile' cannot be used on the V100 GPUs we use, therefore others may be able to reduce this compute time in future studies. All compute exploratory phase of this project, we used more GPU hours to test and design the experiments shown, using approximately 1.5 terabytes of storage of the entire project. An estimate of the total compute required for all the results presented the main paper is 10, 039 GPU hours. The appendix results require a further 18, 278 GPU hours.
provided by internal resources. During the
A.8.1 Hyperparameters
We detail what we believe to be an important subset of the default hyperparameter values in Table 5. A full list of all hyperparameters and model configurations contained in the code release. For multiplication models with FIRE embeddings, the learning rate 0.00006, due to large instabilities in higher learning rates which were not experienced for the Abacus Embeddings.
A.8.2 Code Release
We will release all code and datasets on GitHub with an MIT License.
23

---

Table 5: Default hyperparameter values.
Hyperparameter
Default Value
Hidden Size
1024
Intermediate Size
2048
Embedding Size
1024
Number of Attention Heads
16
Progressive Loss Alpha [Bansal al., 2022]
1.0
Data Type
float I 6/float32
Optimizer
AdamW [Loshchilov and Hutter, 2017]
Global Batch Size
8192
Batch Size Ramp
0.6
Learning Rate
0,0001
Learning Rate Scheduler
Trapezoid [Zhai et al., 2022]
Activation Function
GELUglu [Shazeer, 2020]
Normalization Layer
LayerNorm [Ba et al., 2016]
Normalization Type
Post
Offset Randomization Hyperparameter (k) Initialization
100
Deepnorm [Wang et al., 2022]
24

---

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

---

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

---

Answer: [Yes]
Justification: We will upload our implementation code and datasets on Github. During the submission eycle, provide anonymized implementation that be found in the supplementary material of this submission. Our implementation is licensed under the MIT license.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
//nips. /nips.ccppuuliii wwtpcccppppccccctccccccccccccccc cosplliccguieee guides/eee/e// sonddeSuummsssooooo11ioo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please see Sections 3 and A.8.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[No]
Justification: While we average over several trials, we do not report exact error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
27

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please see Section A.8.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurIPS Code of Ethics https://neurip:: http::////ourrss...cccuppb/////a https/////cpccctttcccotoeeeeeee.
didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification:
We have read and follow the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The societal impact of improving transformer capability on simple arithmetic is limited to the possible effects of slightly improving the community's 's understanding. These particular tasks are far away form the frontiers of potential harm/benefit society.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release models with potential to misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Please Section A.8.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
29

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We relay the details of constructing samples but we do not release any actual datasets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "oCBSc2CCpn", "reviewer": "Reviewer_uGwk", "summary": "The paper studies the arithmetic capabilities of transformers and the problem of length generalization, specifically the ability to solve problems larger than the ones seen during training. It introduces Abacus Embeddings, a novel positional embedding that encodes the position of each digit relative to the start of the number. For multi-digit addition, Abacus Embeddings result in state-of-the-art generalization to sequences six times longer than the training sequences. Additionally, the paper explores the benefits of incorporating recurrent blocks, leading to further improvements. Finally, the paper demonstrates the effectiveness of Abacus beyond addition, showing success with in-distribution multiplication and array sorting tasks.", "questions": "3. Differently from addition, when considering multiplication, Abacus Embeddings achieve high in-distribution accuracy but struggle out-of-distribution (OOD), even when one operand is of unitary/short length. Do you have any insights about what might cause this significant difference? Do you have any ideas or potential modifications to the method that could improve OOD generalization in this context?\n4. Could you elaborate more on how your embeddings could be integrated into settings that involve mixing arithmetic with natural language data?", "limitations": "The paper adequately discusses its limitations. I do not foresee any potential negative societal impacts arising from this study.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **Originality:** While previous studies have noted that positional encoding can negatively impact the arithmetic generalization capabilities of transformer architectures, to the best of my knowledge, the introduced embeddings, the analyses, and the results presented in this paper are original.\n- **Quality and clarity:** The work is technically sound. The experiments are well-designed and convincing, and the code for their implementation is provided in the supplementary material. The paper is clearly written.\n- **Significance:** The goal of improving the extrapolation and reasoning abilities of transformers is both timely and significant. The results are very good, achieving state-of-the-art performance for length extrapolation in multi-digit addition.", "weaknesses": "1. The paper does not discuss the choice of the value used for the maximal offset randomization parameter $k$, which determines the distribution of the starting position of the first digit of the numbers during training. How was the value $k = 100$ chosen? Additionally, could higher values further improve extrapolation performance?\n2. The sentence “our methods perform so well that we look beyond addition” (line 239) does not sound appropriate for a scientific paper. Please consider rephrasing it, for instance, “Given the strong performance of our method in the multi-digit addition task, we extend etc.”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable time, your comments have greatly improved our draft. \n\nWe respond to your weaknesses and answer your questions below in order:\n1. We find ~100 to be roughly optimal when the training data has a maximum of 20 digits; this already allows for addition of numbers larger than a googol. We describe and discuss the experimental evidence we gathered during the rebuttal period for varying the value of the maximal offset randomization hyperparameter (k) in the general rebuttal response and pdf. From this we conclude that the value of k can be varied with larger numbers in the training data allowing for larger values of k to be used.\n\n2. We agree with the reviewer and have updated our draft.\n\n3. While our results for multiplication are currently state-of-the-art compared to prior work [1], the reviewer is correct that we do not observe out-of-domain generalization. We do not know how to solve this problem at this time, but we do have several hypotheses on why do not see larger gains:\n    * Our multiplication models are small compared to the multi-billion parameter models we have become accustomed to in the open source community. With more compute or larger models, it may be possible to improve upon our results.  This direction is motivated by the observation that multiplication models, even with their increased compute budget compared to addition, do not achieve as low of a training loss as addition models. Furthermore, for addition models the loss plateaus at low values at the end of training, allowing for a period of training on low loss values. This does not occur for multiplication models, leading us to speculate that there is something to be gained from larger-scale training.\n    * Multiplication requires more abacus embeddings during training and testing, due to the increased output length.\n\n    We have updated our draft to include these as future directions of research.\n\n4. Provided a model tokenizes numbers by their individual digits, which is done in most new open source models (e.g. llama and gemma), a plug in and play version of Abacus Embeddings (which is our sorting implementation in the supplementary material) can be used alongside other positional embeddings for language tasks. This code simply identifies all digit tokens in the input sequence and calculates the Abacus Embeddings in parallel for a batch, these can then be added to the input embeddings. Abacus Embeddings also rely on a least significant digit first format which can be done during encoding and decoding in the tokenizer with a simple regex to reverse all numbers.\n\nShould you have any further questions or require additional clarification, please do not hesitate to ask.\n\n[1]  Shen, Ruoqi, et al. \"Positional description matters for transformers arithmetic.\" arXiv preprint arXiv:2311.14737 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9BWvGMdN1d", "reviewer": "Reviewer_mAhx", "summary": "This paper studies a well-known problem, the length generalization issue of transformers in terms of doing arithmetic. This paper solves this problem via two natural strategies: (i) separate two operands via a newly proposed embeddings (Abacus Embeddings), and (ii) using looped Transformer architecture.", "questions": "Why does recurrency in terms of the model architecture help with length generalization?", "limitations": "See weakness.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The problem is well-motivated.\n- The conjectures are very natural, and confirmed via extensive experiments.\n- Experiments are well-designed and complete.\n- The proposed solutions enjoy great performance.\n- Considers diverse downstream tasks, including addition, multiplication, and sorting.", "weaknesses": "[Medium] The reason why looped transformer or recurrency helped with length generalization is still unclear, and in-depth analysis is needed. For instance, does the number of recurrence related to the length of the digits?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable time, your comments have greatly improved our draft. \n\nWe answer your question below:\n\nWe do not find that the number of recurrences is meaningfully linked to the length of the numbers in this study and we do a small visual analysis of the intermediate properties during recurrence in Appendix Figure 13. While we do find that a small amount of recurrence can lead to performance improvements throughout our work, we also show that we can achieve highly accurate out of distribution performance with standard decoder architectures using Abacus Embeddings.\n\nWe hypothesize looped transformers may improve performance because they force the model to learn an algorithm that relies on an iterative process. This aligns their strategy with that of a human-designed algorithm, e.g. traditional addition algorithms repeat the same process iteratively for each pair of digits. The good inductive bias of recurrent models is more widely referred to as the theory of algorithmic alignment [1] and has motivated many algorithmic reasoning results, for example [2]. \n\nWe believe the weaknesses are addressed in the reviewer questions. Should you have any further questions or require additional clarification, please do not hesitate to ask.\n\n[1] Xu, Keyulu, et al. \"What can neural networks reason about?\" International Conference on Learning Representations (2020). https://openreview.net/forum?id=rJxbJeHFPS\n\n[2] Ibarz, Borja, et al. \"A generalist neural algorithmic learner.\" Learning on graphs conference. PMLR, 2022. https://arxiv.org/pdf/2209.11142", "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": "FacXw89sU0", "reviewer": "Reviewer_GBbn", "summary": "This paper introduces a simple yet effective encoding scheme that can be used to address the limitations of transformers at representing positional information, which is crucial in many algorithmic tasks such as those involving arithmetic operations. The authors propose an ad-hoc positional embedding, called “abacus embedding”, which encodes the location of each digit relative to the start of the current number and thus provides an explicit signal that the transformer can use to align digits. The effectiveness of the method is tested on addition, multiplication and sorting problems, with a particular focus on out-of-distribution test cases.", "questions": "- How could we address the fact that the hyperparameter k needs to be fixed a priori?\n- Can we implement a unified model that can learn all these tasks simultaneously?\n- The proposed method achieves impressive OOD accuracy for addition, but only works “in distribution” for multiplication. It would be important to investigate this phenomenon more in depth.\n- At least for addition, it would be useful to test OOD generalization by also adding more operands besides increasing the length of each operand.\n- How does the present method compare to other recent proposals such as xVal (https://arxiv.org/abs/2310.02989)?\n- I agree that addition, multiplication and sorting are good benchmarks because they are simple yet challenging; however the authors could better stress that these tasks are part of a broader class of elementary numerical skills that transformers struggle to learn (for a comprehensive review, see https://www.mdpi.com/2076-3417/14/2/744).", "limitations": "The authors properly addressed the limitations of their study.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "I think that this work is interesting and relevant. Although addition, multiplication and sorting problems might be considered trivial test cases because they can easily be solved with symbolic algorithms, they constitute an important benchmark to evaluate the algorithmic reasoning skills of neural networks, as also attested by the increasing interest of the deep learning community on mathematical tasks. The paper is well-written, and the method is clearly presented. The generalization achieved in the addition task is quite impressive, showing that the abacus embeddings enable a generalization factor of 6x in the OOD regimen. Although simple and straightforward, the proposed method seems original.", "weaknesses": "The abacus embeddings are defined according to the hyperparameter k, which is fixed a priori (e.g., k = 100). This limits the flexibility and generalizability of the proposed encoding scheme.\n\nThe authors deploy different architectures / hyperparameters to learn different problems (addition vs. multiplication vs. sorting). Since they argue that their architecture modification “improves performance on multiple algorithmic reasoning tasks simultaneously” it would be important to show that different tasks can really be learned simultaneously, without the need to build ad-hoc models for each algorithmic task that needs to be solved.\n\nIt is true that arithmetic operators are binary and thus “both addition and multiplication accept only two operands”. However, we can have a sequence of additions / multiplications, and it is well-known that also increasing the number of terms in arithmetic expressions causes troubles to transformers.\n\nBecause of these key issues, I think that the impact and significance of this work are not strong enough for a top-tier venue like NeurIPS.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable time, your comments have greatly improved our draft. \n\nWe answer your questions below in order:\n1. While choosing k a priori could be a barrier to adoption of Abacus Embeddings in the community, we have shown that we can scale to at least one hundred more digits at test time than seen at training time for addition in our original submission. We believe that use cases with distribution shifts larger than this are more limited. Furthermore, in our general rebuttal response and in Rebuttal Figues 2 and 3, we show we can increase the distribution shift using a larger k and larger training data if a distribution shift of more than 100 digits is required. \n\n    For any language model a fixed context length is also decided a priori, in the extreme case one could choose an extremely large value of k, for example context_length/3, so that all two argument additions can be completed by model.\n\n2. To see whether multi-skill training is possible, we trained a model on both addition and subtraction simultaneously without any hyperparameter changes from the addition models presented in the original draft. We show the results of this experiment in Rebuttal Figure 4, in the rebuttal pdf. Here we show that, even without hyperparameter tuning due to rebuttal time constraints, the models are able to extrapolate for both the symmetric addition operation and anti-symmetric subtraction operation simultaneously using a single model. We stress that we are training tiny transformers models in this paper, up to a maximum of 122M distinct parameters. We hypothesize that scaling either of these axes would allow for more simultaneous learning, however this is not the regime we analyze within this report.\n\n3. We emphasize that we also achieve state-of-the-art performance on multiplication with Abacus Embeddings, when compared to prior work [1]. We agree that further improving performance for multiplication out of distribution is important future work and have updated the future directions section of our draft to emphasize this.\n\n4. While this is a form of algorithmic generalization, we chose not to study generalization in the number of operands for addition or multiplication in this study, in a compute restricted regime. We instead highlight the generalization of operands in the sorting section. \nFor sorting we have two dimensions of generalization. Firstly, “OOD in number length,” this is similar to addition and multiplication where we increase the number of digits in the numbers in the array being sorted. Secondly, “OOD in array length,” here we increase the length of the array so there are more numbers that need to be sorted. We analyze both of these dimensions in the reported “ALL OOD” accuracies where we scale each of these dimensions concurrently. Specifically, the “All OOD'' accuracies shown in Table 1 highlight that when both the number of operands and number of digits is varied during testing, the models trained with Abacus and FIRE Embeddings achieve the highest accuracy for the sorting task.\n\n5. xVal, which embeds all real numbers by scaling a single fixed token-embedding, is a very important contribution to the language modeling community. We reference xVal in the Related Works section of our paper as it improves mathematical accuracy when the operands are small real numbers, up to 10e+8. However, it does not resolve the algorithmic generalization problems which we are working to solve in this paper, involving numbers much larger than 10e+8 (for example in Rebuttal Figure 3 we analyze up to 10e+214). The authors of the xVal paper highlight this in Section 4 - Future Directions, “Very large numbers saturate the normalization, as discussed in Sec. 2, and very small numbers are negligible from the model’s perspective.” This is because the authors “normalize numbers in the text corpus such that they fall within the range [−5, 5] as a preprocessing step before training.” Hence, we believe xVal would not be a compelling baseline for our study. We choose FIRE and RoPE as our main comparisons based on prior work and directly compare to the previous state-of-the-art (Randomised-FIRE with index hints) in Appendix Section: Addition Ablations - Index Hints.\n\n6. We agree with and will act on this feedback for a future version of the paper, including the paper cited above, we have updated this in our draft.\n\nWe believe the weaknesses are addressed in the reviewer questions. Should you have any further questions or require additional clarification, please do not hesitate to ask.\n\n[1]  Shen, Ruoqi, et al. \"Positional description matters for transformers arithmetic.\" arXiv preprint arXiv:2311.14737 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Our revised manuscript will include the following new experiments and discussion, which better clarify the broad utility and flexibility of Abacus Embeddings. We thank the reviewers for their questions and suggestions that led to these new positive results!


**Varying the value of the maximal offset randomization hyperparameter (k)**

In Rebuttal Figure 1, in the rebuttal pdf, we show models trained with k=25, 50, 75 and 100 trained on size 20 data; with the k=100 models being taken directly from Figure 4 of the current paper. We show the average of 3 models in each case analyzing accuracy only where both operands are the same length, similarly to Figure 9 of the current submission, to save on computation time during the rebuttal period. We see in the plot that these smaller values of k allow for good extrapolation and the amount of extrapolation depends on the value of k as expected. We find that increasing k to values more than 100 for models trained on data with a maximum of 20 digits, leads to diminishing returns within our experimental setup. 

However, this can be resolved by including larger numbers in the training data. In Rebuttal Figures 2 and 3, in the rebuttal pdf, we show models trained on size 30 and size 40 addition data respectively. We show that larger values of k allow for much larger distribution shifts, with a maximum distribution shift of up to 175 digits being shown, this is a 6.8x length generalization from training to testing. Hence, we can easily increase k to larger values and perform arithmetic with far more digits than we did in our original state-of-the-art submission, with suitable training data.

**Learning addition and subtraction simultaneously**

In Rebuttal Figure 4, we train a model with exactly the same hyperparameters used to train the addition models in the submitted paper but this time also include subtraction examples in the data. We see that these small language models can simultaneously learn to extrapolate for both the symmetric operation of addition and the anti-symmetric operation of subtraction using Abacus Embeddings.

Meta-review: The paper addresses the important and timely problem of length generalization in arithmetic tasks by introducing a simple yet effective embedding strategy that leverages the significance of each digit. This approach demonstrates impressive performance in solving length generalization challenges, showcasing the potential of straightforward methods in complex problem domains. This embedding also facilitates the use of looped TF models with input injection, resulting in a further improvement in length generalization accuracy. 

The reviewers commend the paper for its clarity and quality of writing, which effectively conveys the concepts and results to the reader. The simplicity and effectiveness of the proposed method are particularly noteworthy. While the approach is somewhat tailored specifically to arithmetic tasks, this focus does not detract from its contribution, as it provides valuable insights that could inspire future research in related areas. Considering these strengths, the consensus is to accept the paper for publication.
Decision: Accept (poster)