Transformers Can Do Arithmetic with the Right Embeddings

Authors: sean  mcleish
arpit  bansal
alex  stein
neel  jain
john  kirchenbauer
brian r bartoldson
bhavya  kailkhura
abhinav  bhatele
jonas  geiping
avi  schwarzschild
tom  goldstein
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 00:00:33

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.

Main paper: 
 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 a 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 us to study several important questions. First, we ask what architectural design choices, dataset characteristics, and training pipeline variants are required to learn a many-step reasoning process like multi-digit addition? Going deeper, we 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 a simple modification to 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-art generalization factor of 6×, with the previous state of the art being only 2.5×. To the best of our knowledge, these are the longest sequences on 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-skip connections inserted between the input layer and each decoder layer-can reduce generalization errors by 50% over the Abacus Embedding baseline. 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 near-perfect generalization on addition problems we 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-perfect in-distribution generalization. • We show that when we combine Abacus Embeddings with input injection and looped transformers performance further improves, increasing from 92.9% to 99.1% in out of distribution accuracy, an 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.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 be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Arithmetic and Algorithmic Reasoning. Solving arithmetic with next token prediction is a difficult problem that attracts a 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-based models, reversing the digits so the arguments are written with the least significant 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 token-embedding for numbers [Golkar et al., 2023] . Moreover, Dziri et al. [2023] show multiplication is a 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 a 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 as 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 et al., 2023a , de Luca and Fountoulakis, 2024] . Additionally, recent work aims to improve reasoning in LLMs [Zhou et 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 embedded 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 et al., 2022] , and Alibi [Press et 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 [e.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 et 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.5× 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 A.1.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] .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 are not discussed in the paper. • The authors are encouraged to create a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
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. 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 et al., 2023 , Lee et al., 2023] , inputs 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 et al., 2023] . We train on all combinations of operand lengths less than or equal to i 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 i = j, hence we can use one number to define the dataset i, where i is the maximum length of either operand. We sample data with replacement and we stratify the data, so that all length pairs (i, j) are equally sampled during training. To facilitate training of many models from scratch, we use a language model cramming setup [Geiping and Goldstein, 2023] and limit each training run to 8 exaFLOP of compute (a single Nvidia RTXA4000 GPU for 24 hours); for multiplication results we 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 ̸ = j 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., 2024] ). 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 categories: (i) in distribution (ID) where the models are tested on 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 are 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 at this scale.We consider two standard transformer architectures. First, we use a standard autoregressive transformer model where multiple decoder layers are stacked in a feedforward manner. Second, we enhance this standard transformer model by incorporating input injection, where the embedded inputs are 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 .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.e. each digit's position relative to the beginning of the number) is not easy for transformers to represent, and that this sub-problem 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 = a7b3c9, finding that transformers perform much better on addition with the information provided by such hints [Zhou et al., 2023 [Zhou et al., , 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 et al. [2024] find that the ability of models trained with index hints to generalize is sensitive to the particular random initialization. [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 to 0% accuracy at 90 digit addition.To address the limitations of 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 same significance, providing an explicit signal that the model can use to align digits. We visually describe these embeddings in Figure 2 . 3We take inspiration from Randomized Embeddings [Ruoss et 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 in a number, starting from a randomly chosen offset value fromU [1, k],where k is a 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 are β, β + 1, β + 2 where β ∼ U [1, 100], which are then passed through a learned embedding matrix. The value sampled from U [1, k] is the same for all numbers in a batch, meaning all digits of the same significance obtain the same positional embedding. This training scheme allows the model to see a wide range of positional embeddings, even when training sequences are short. At test time, each positional embedding begins from one, i.e. β = 1.Abacus Embeddings Solve Addition. Abacus Embeddings improve generalization performance up to 100 digits and beyond for standard transformer architectures. In Figure 3 (left), we highlight the comparative boost Abacus Embeddings have over standard transformer architectures and embeddings for performing addition, taking the mean accuracy of three models in 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 as 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 to 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 randomizes the starting offset used for each individual addition example can be increased to enable generalization by training a larger range of embeddings for a given computational budget. Relatedly, Appendix Figure 9 shows that training on larger datasets improves performance, even for operands with fewer than 100 digits.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 to achieve the effective depth.We also employ input injection, skip-connections that propagate a copy of the input to each layer in the network.The Benefits of Recurrence. In Figure 3 (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 10× 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 is a convex combination of the loss values from two forward passes, one with the nominal number of recurrences (so 16 for a 1 × 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 a model with sixteen layers in the block and a single recurrence (16 × 1, i.e. a standard transformer), through to one layer in the block but with sixteen recurrences (1 × 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, a model with two recurrences (8 × 2) incurs half the error of the purely non-recurrent model (16 × 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 3 in the appendix, we record the parameter counts to the nearest million.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA]Answer: [Yes] Justification: We will upload our implementation code and datasets on Github. During the submission cycle, we provide an anonymized implementation that can 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. 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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: The societal impact of improving transformer capability on simple arithmetic is limited to the possible effects of slightly improving the community's understanding. These particular tasks are far away form the frontiers of potential harm/benefit to 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 not 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
pushing the limits of algorithmic reasoning for transformers: 
 While there is an emphasis on addition as a difficult problem in existing work, our method's strong performance allows us to extend to even more difficult problems, including multiplication and sorting and even multiple operations at once.We train models on a dataset made up of an even mix of addition and subtraction samples. In Figure 5 , we show results from models with 8 layers in the recurrent block and 2 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. We now study a harder task, multiplication of 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 one digit more than the longest operand, multiplication has longer-distance dependency and the output length scales much faster as problem size increases.To adapt from addition to 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, down-weighing 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 us to a much smaller scale analysis.Abacus Embeddings help looped transformers reach near-perfect accuracy in-distribution for multiplication. 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 distribution accuracy on 15 digit multiplication than prior work [Shen et al., 2023] and do not require padding each operand to 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 ). In this setting, we explore two axes of generalization. First, we increase the maximum possible length of the input numbers to 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 to 30 while keeping the maximum digit length of each number at 10 and term this scenario "OOD (array length -30)." Finally, we consider a scenario where both axes are increased simultaneously, referred to as "all OOD."In Table 1 , we illustrate the performance of a standard transformer (eight layers) trained with different embeddings-FIRE, Abacus, and their combination. Again, our results demonstrate that the combined embedding approach enhances the model's ability to generalize, surpassing the performance of either 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 utility across a broader spectrum of algorithmic tasks beyond basic arithmetic. Abacus Embeddings may be instrumental in use cases requiring transformer models to perform a variety of complex positional, numerical, and/or relational reasoning tasks.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 to maintain good downstream performance on non-arithmetic 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 are 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 , we demonstrate the true potential for integrating Abacus Embeddings into a 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.  
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 a challenge. In this paper, we study mathematical reasoning tasks including addition, multiplication, and sorting to evaluate these capabilities in a controlled setting. We analyze the ability of specialized language models to learn algorithmic tasks in a zero shot setting, without access to 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 dramatically both when applied to standard transformers as well as recurrent variants. We repeatedly achieve length generalizations of at least 6× (capped by the context length) more than doubling the extrapolation demonstrations in prior work, achieving near perfect results on addition of up to 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 simultaneously. We hope that our work deepens the community'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.Sean McLeish* -Led the project, developed the idea, contributed to code, organized majority of experiments and contributed to writing. Arpit Bansal* -Contributed large amount to 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 to 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.1 Positional Embeddings.FIRE embeddings are additive embeddings in the attention mechanism:A RP E (X) = XW Q (XW K ) T + B where B i,j = f θ log(c(i-j)+1)log(c max(i,L)+1) and c, L are learnable parameters. Li 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 embeddings are trained. Zhou et al. [2024] use randomized FIRE [Ruoss et al., 2023 , Li et 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.Addition: We sample equally, with replacement, from all i × i possible operand lengths up to the maximum dataset size of 20 million, we call this a dataset of size i in the main text. For evaluation we sample 100 samples for each pair of operand lengths evaluated.Bitwise OR: 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.Sorting: Given a list of reversed integers indexed by characters, output the characters in ascending order. E.g. a : 64957, b : 99963, c : 10218, d : 7141, e : 05781 = d, e, b, a, c. We implement the sampling process for sorting in a grid like manor. We query each "square" of an [1, n] × [1, n] grid until the maximum size has been reached for the dataset. When querying "square" (i, j) we randomly sample i 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.Multiplication: 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.A necessary condition to perform addition is aligning digits of the same significance. We begin by examining positional embeddings for exactly this task. To do 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 A.3.1, these are left aligned to be representative of the task of aligning digits for reversed addition. 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 we see that the Abacus Embeddings allow all models to generalize further on this task than the other embeddings which prior work for addition focuses on.As with addition, we see 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 et al., 2022] . When we analyzed the errors in this task we found they were predominantly caused by the model outputting one too few or too many zeros.000010 ⊕ 00000000000000 = 00001000000000 000100 ⊕ 0000000 = 0001000 001 ⊕ 00000 = 00100Across 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 is partially due to 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.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 k = 100. One possible way to resolve this generalization problem is to increase the value of k during testing. In Figure 10 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 of the red squares in Section A.6. Models trained on larger operands achieve higher OOD accuracy.data with Abacus Embeddings and k = 101, generalizing to 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] . 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 k, each line is the average of 3 models.Here we present the mean accuracy as heatmaps for the main addition experiments shown throughout the paper. Thanks to 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 in Figure 15 and suggest that this level of interpretability could be leveraged in future work. The model presented is a 1 × 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 a fixed point to 16 iterations and beyond. We mask all tokens before the equals sign in 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 as the other addition experiments without masking before the equals sign.We begin in Figure 17 by showing a replica of Figure 4 , this time including comparisons to 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, even though it only has approximately 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.In Figure 19 (left), we remove the input injection to the intermediate layers in the recurrent block, only keeping input injection to the 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 × 8 model improving its performance slightly in left plot and the 4 × 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.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 are reduced as the training data is noised with random pad symbols, a method which was shown to improve length generalization in prior work [Shen et al., 2023] .A.7.5 Index Hints Zhou et al. [2023] "randomly sample consecutive index hints from a pre-defined ordered set of hints with 102 symbols," for example a6b7c5 + a1b6c3 = a7b3c9. 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 in order. We see similar results for models trained on up to twenty digits as Zhou et al. [2023] .We do note that our format of taking the mean exact match accuracy does highlight robustness as if one of the three models tested were to not generalize well, this would impact reported accuracy heavily. We only show a 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 Abacus, OOD FIRE, OOD NoPE, OODFigure 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 a single Nvidia RTXA4000 if we do not mask the input question in the loss function.et al. [2024] with their methods, we try to the best of our abilities to faithfully reproduce their work within our experimental set up, noting that perhaps a better random seed or initialization may be able to produce better results for these models.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 . 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. 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. As Abacus Embeddings are a variant of absolute embeddings, reused only for numbers, they could be combined with relative embeddings being deployed in current models. If all digits input to the model are tokenized individually, we can perform a linear time operation to find and assign relative embeddings to all numbers in an input, which is lower than the quadratic cost incurred by attention.Training a small number of Abacus Embeddings may be enough to handle all numerical inputs for 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 a local batch size which is the maximum batch size that is a 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 a batch; we find this leads to slightly more stable training. We note that training models to solve multiplication requires more hyperparameter tuning than addition, perhaps implying it is a trickier task to learn. Also, FIRE models require a much greater compute budget for hyperparameter search as compared to 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 4 . For some experiments, such as extreme length generalization (Figure 10 ) and index hints (Figure 21 ) more GPU hours are required for testing, these are included in the total number of GPU hours used. Our testing pipeline for addition and Bitise OR uses Nvidia V100 GPUs. Due to a 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 was provided by internal resources. During the 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 of the results presented in the main paper is 10, 039 GPU hours. The appendix results require a further 18, 278 GPU hours.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 is contained in the code release. For multiplication models with FIRE embeddings, the learning rate is 0.00006, due to large instabilities in higher learning rates which were not experienced for the Abacus Embeddings.We will release all code and datasets on GitHub with an MIT License. Deepnorm [Wang et al., 2022] NeurIPS Paper Checklist 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not release models with potential to misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Please see Section A.8.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.In Appendix A.3, we motivate these embeddings further with experiments demonstrating their utility in solving a bitwise OR task. 
 
 Appendix:  


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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?
2. 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.”.
questions: 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?
4. Could you elaborate more on how your embeddings could be integrated into settings that involve mixing arithmetic with natural language data?
limitations: limitations: The paper adequately discusses its limitations. I do not foresee any potential negative societal impacts arising from this study.
paper_summary: paper_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.
strength_weakness: strength_weakness: [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?
questions: questions: Why does recurrency in terms of the model architecture help with length generalization?
limitations: limitations: See weakness.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

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

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

Because 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.
questions: questions: - How could we address the fact that the hyperparameter k needs to be fixed a priori?
- Can we implement a unified model that can learn all these tasks simultaneously?
- 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.
- At least for addition, it would be useful to test OOD generalization by also adding more operands besides increasing the length of each operand.
- How does the present method compare to other recent proposals such as xVal (https://arxiv.org/abs/2310.02989)?
- 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: limitations: The authors properly addressed the limitations of their study.

Comments: comment: Thank you for your response. I will keep my initial positive evaluation. Best wishes.
I thank the Authors for having considered my comments. Having read their responses and the comments posted by the other Reviewers, I am persuaded to raise my score from 4 to 6.