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_Workshop_MATH-AI
Published: 2024-09-25 01:08:07

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 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: This paper proposes a new positional embedding that can achieve strong OOD on addition problems. It would be interesting to see if this method has broader implications.

Reviews: review: title: title: Review
main_review: main_review: This paper advocates for the use of the right position embeddings + input injection + recurrence to get length generalization in arithmetic. They start with an addition experiment where they train several models on up to 20-digit addition and compare NoPE, FIRE and their position embedding Abacus. This latter consists in giving the same positional embeddings to all digits of the same significance. To avoid length generalization issues, they do the same trick as in the random position embedding where they give values to the position embedding that may be encountered at inference time. They show with this Abacus embedding that they can get 93% on 100-digit addition. Then, by addition recurrence + input injection, they show they can get 99% on 100-digit addition. 

I overall like this paper. They basically managed to get the right position embeddings for addition and something else that I find nice is that they show that input injection + recurrence have complementary effects in that they further increase the accuracy. For this reason, I strongly advocate for acceptance of this paper. 


A few remarks:

- The proposed position embedding looks like a lot to index hinting. The authors mention the connection to this and explain that index hinting is more costly (which totally makes sense). I would still be interested in comparing index hinting vs Abacus. My hypothesis is that to get length generalization one needs these “sparse hints” and giving it at the level of position embedding and at the level of the input should not matter. 
 
- One obvious criticism is that this set of techniques seem to be very specific to the case of addition. I would maybe add a limitation section to mention this. Would these tricks apply to multiplication and sorting for instance? [1] also came up with a smart position embedding that seems to work for other tasks than addition.


[1] Cho, H., Cha, J., Awasthi, P., Bhojanapalli, S., Gupta, A. and Yun, C., 2024. Position Coupling: Leveraging Task Structure for Improved Length Generalization of Transformers.
title: title: Review for Submission63
main_review: main_review: Summary:
The paper aims to solve the long-standing problem of summing up two numbers using transformer, which doesn't generalize on OOD yet. Authors showed that, with the right embedding, transformer can learn this task and also two other tasks: multiplication and sorting.



Strength:
- Authors described a new embedding, which injects domain knowledge into the model. In particular they inject the digit "relevance" when they sum two numbers and showed that with the new embedding it's possible to generalize OOD.
- Authors compared their proposed embedding against the standard embeddings, in particular NoPE and RoPE. They showed that these embeddings fail to generalize, while Abacus can solve OOD.



Weakness:
- There is one important missing point from this paper: https://arxiv.org/pdf/2405.18719. The cited work offers an opportunity to create an embedding which is learned (compared to Abacus where it's applicable for math task only) and it can provide similar results in terms of accuracy and OOD accuracy. One suggestion would be to try CoPE and see if it doesn't perform as good as Abacus. If this is the case, I think Abacus can play a much bigger role into the Embedding selection for math tasks.
- There is no discussion on how this embedding can help on other harder math tasks, where it's not possible to inject this inductive bias. Multiplications and sorting follows a similar pattern with respect to addition.



Question:
I keep reading that, due to inference, you cannot run more tests. I am curious to know if you do model.forward rather model.generate, given you are interested in exact accuracy. This should speed up your computation by a good factor.
title: title: The paper propose a different positional embedding, abacus embedding, which enables length generalization for addition
main_review: main_review: The experiments on addition are convincing. 
The paper could be more clearly written: the related work and the beginning of section 3 speak of recurrent architectures, which are not evoked in section 3.1. 

A question for the authors: replacing the absolute positional embedding by the abacus embedding means that the tokens in sequences "1254+248=" and "1258+244=" (where the last digits of each number were exchanged) are encoded the same. For addition, this has no bearing on the result, but this would be false for any other arithmetic operation (substraction, multiplication, modulus...). Isn't this a strong limitation of abacus embedding, ie isn't it too specific to the addition task?