Teaching Arithmetic to Small Transformers

Thank you for your valuable and overall positive feedback on our paper. We have responded to your concerns below and hope that this will convince you to elevate your score.

scalability of the findings The introduction says that the novelty of this work over prior work is to pinpoint the factors that contribute to the fast emergence of arithmetic capabilities through careful ablations, but I'm not sure how well the findings from this study would transfer to LLMs at scale. Additionally, I feel like some of the findings are not properly backed by the experiments:

We apologize that this was not clear. We have a full table on scale ranging from NanoGPT, GPT2 and up to GPT3 models in Section 8 and also some additional experiments in Section F showing that most of the findings do indeed carry over when we scale up the models.

transferability to language data The Transformers in this work are almost exclusively trained on arithmetic data, and it's unclear whether the findings would transfer to models trained primarily on language with some arithmetic data weaved in. Indeed, you show that arithmetic can be learned from next-digit prediction, but that is unsurprising since the setup is simply a cross-entropy loss on the next "token"(=digit).

While we agree that the loss is simple CE loss on next "token", it is unclear whether this is sufficient to learn the underlying algorithm that is necessary for performing arithmetic. Note that arithmetic in the token space is not equivalent to arithmetic in the embedding space. Moreover, this does not explain generalization to unseen examples of arithmetic. It is this aspect that is surprising. In fact, what is even more surprising is that without using careful data formatting and sampling, NanoGPT is unable to learn arithmetic with a reasonable number of samples.

Furthermore, we have conducted experiments where arithmetic tasks are intermingled with textual data (refer to Section 7 and Appendix E). Although these experiments are on a small scale, the results suggest that our findings are transferable to a setting with the presence of text, and that few-shot promptining in these scenarios enhances the performance.

practicality of reverse format The finding that reversing the output helps also most likely doesn't explain emergent abilities in LLMs, as this doesn't happen in natural text.

While we agree that reversing the output is not something that occurs in natural text, we use it as a mechanism to describe the importance of data formatting in learning compositional functions. Reversing the output is meant as an illustration of how even simple data formatting changes can elicit large improvements in performance.

few-shot prompting on no-text model You show that few-shot prompting can be used to improve performance when text is mixed in, but you don't try few-shot prompting when text is not mixed in, so it's unclear whether the increase in performance is due to text being mixed in, or simply because of few-shot prompting. Few-shot prompting should also work when you train the model without text, or am I missing something?

Please refer to Figure 21 where we have results on few-shot prompting on models trained exclusively on text data. As the reviewer expects, few-shot prompting does work even when the model is trained without text.

trasnferring the data sampling strategy The finding that sampling strategy is important is interesting, but on its own it's again unclear whether that transfers to emergent capabilities of larger models. For example, does the distribution of numbers in pretraining dataset follow a similar one as the one produced by your sampling strategy?

Is the reviewer referring to typical pretraining datasets such as C4? If so, then it is difficult to even extract the distribution of numbers from it due to scale. However, we would expect that it does not since the arithmetic performance of most typical LLMs trained on C4 and the Pile is quite poor. We hope that our work will encourage more research in the area of dataset creation and curation. As seen in Gunasekar et al. "Textbooks are all you need", high quality tokens are becoming more and more important in order to train models efficiently.

Thank you for your feedback on our paper. We appreciate your comments and would like to address your concerns.

$-symbol I was curious about the $ symbol bit present here and there in the paper, but it became clear after reading appendix B. In my opinion, the baseline used for making the strong claims of models being able to handle reverse addition better in the paper is wrong. The authors should have either done the reverse methodology without $ or used $ for the baseline too. This is a bit concerning as data formatting techniques highlighted in the paper are touted as one of the important contributions, and this finding makes that invalid.

We thank the reviewer for their careful reading of the paper and acknowledge that this is somewhat confusing. In order to simplify the message of the paper, we decided to subsume the $ delimiter into the reverse data format. However, as the reviewer points out, even in Appendix B, it is clear that reverse does improve the performance significantly both with and without the $ delimiter. We will modify the wording in the paper to make this more clear.

Figure 9. plain reaching almost 100% Expanding on the previous point, Figure 9 specifically highlights that even plain addition is able to reach almost 100% test accuracy with the addition of the $ symbol, and reverse without $ is almost at 90% (only 2/3% difference with the baseline).

You are correct in observing that plain addition can approach 100% test accuracy with an adequate number of training examples. We have updated our manuscript to reflect that the plain format does not plateau at 85% as previously stated. We also found that introducing a \n character as a prefix while prompting the model significantly enhances performance for models trained without the $ delimiter. This result is presented in Figure 9, which can be viewed here. These results indicate that the inclusion of a delimiter during testing is a critical factor for model accuracy. We will add this result in our revision, and also include experiments on the plain format with this modification.

Organization of paper I believe some parts of the original manuscript is appendix material and some important bits in the appendix should be moved to the main paper. Specifically, the lemmas in section 4 seem a bit irrelevant given that difference of $ symbol in the baseline and reverse addition. Appendix section B.1 should be present in the main paper.

We maintain that the benefit of reversing the order of digits is significant even though it is reduced with the removal the $ delimiter. However, we will include the contents of B.1 into the main paper as it serves to clarify the setting more clearly.

Typos

Thank you for pointing these out. We have fixed it in the revised manuscript.

scratchpad starting with A=[] Figure 1, $A$ is built incrementally, like $A=[5]$, then $A=[9,5]$, and so on. What is the effect of starting with $A=[]$ (empty set) in the prompt, as the first line just starts with a carry.

We apologize for the confusion. The detailed scratchpad example in Figure 1. initially contained a typo in the first intermediate step. All the detailed scratchpad experiments were actually conducted on training examples starting with A=[]. We have fixed Figure 1 with the correct example.

Figure 2. Setting There's a slight confusion on the experiment setting in Figure 2. Is it plain addition with $ symbol and not the actual baseline? Also it seems that even with structured sampling, the overall performance on 2-digit addition is low when $n=3$.

Yes that is indeed the setting we chose. We chose to use the $-wrapped plain as a baseline since we wanted to use a method that reaches nearly 100% accuracy within a reasonable number of samples. The reviewer is indeed correct that 2-digit addition is not as good as one would hope. The balanced digit sampling strategy involves selecting 100 instances of 1-digit, 900 instances of 2-digit (constituting 9% of the total 10,000 instances for this category), and 9,000 instances of 3-digit numbers. While our strategy performs significantly better than random sampling, it perhaps needs to be tweaked to reach 100% accuracy within each category.

GPT-2, GPT-3 Do you have an equivalent diagram for Figure 9 for GPT-2 and GPT-3?

We have added an additional result on GPT-2 in Figure 9 of the revised version, which shows a similar result as the NanoGPT experiments (Updated Figure 9).

I thank the authors for their response to the review and appreciate the updates to the main manuscript. I still believe the authors have not addressed the mismatch in the baseline clearly, the plain should be with \$symbol too or reverse should not contain the$ symbol. This more fair baseline has very little difference between the plain and reverse. Due to this reason, I maintain my original rating of 5.

Also, there's a typo in Figure 9 (without instead of wihout).

commutativity Does the model learn the properties of addition, e.g. commutativity? This could be done by testing whether model predictions for A+B and B+A are the same, even early during training.

This is an excellent question! We find that while the model doesn't learn commutativity perfectly, it definitely learns it to a large extent, especially after being fully trained. We use NanoGPT fully trained on the plain data format and check if $(A+B) = (B+A)$. Out of 9900 test examples, only 8799 samples commute. While most of these samples are trivially commutative (since the model gets the answer correct), we find that out of the 164 samples where the model gets both $(A+B)$ and $(B+A)$ wrong, it still commutes on 99 of them! We summarize these findings in the table below.

Therefore, the model learns to commute 60.3% of the time even when it gets the answers wrong!

decimal numbers Can those results generalize to decimal numbers (e.g. 1.21+13.12)?

We observe that our results do generalize to addition of two decimal numbers (See Decimal Results). We convert each integer number (up to 3-digit) to decimal numbers with 1 digit + 2 floating points (ex. 1-digit number a -> 0.0a, 2-digit number ab -> 0.ab, 3-digit number abc -> a.bc) to train the model on decimal numbers. The obtained results align with our findings on integer addition, where addition is learned more efficiently with the reversed output. Somewhat surprisingly, the overall performance is better than 3-digit integer addition. However, this is likely because the decimal number formatting behaves like zero-padding since we ensure fixed digit length (Appendix B.1).

signed numbers Can those results generalize to signed numbers, by adding a sign token to all three integers?

We suspect that our results can be generalized to signed numbers as well, which is similar to training the model on both addition and subtraction. Our preliminary experiment on training the model with all signed combinations of two 3-digit numbers show the following results. We find that training with signed numbers is a more difficult task than training on addition or subtraction alone. Interestingly, the accuracy of $(+,+)$ and $(-,-)$ sign combinations is higher than $(+,-)$ or $(-,+)$ sign combination, suggesting that (as expected) learning addition is easier than learning subtraction.

learned algorithm There is a tension between your results from sections 3 and 4, which suggest that models are learning the same algorithm as us (quickly memorizing the 10x10 table, then adding successive digits and propagating carries), and the results from section 5, which frame addition as a memorization+interpolation problem (learning to interpolate a large but low-rank matrix). Is there a way to decide what algorithm the model is actually learning?

This is a keen observation and an excellent point! In fact, most of our work is aimed at trying to glean some insights into what algorithm the model is actually learning. In fact, the second part of Section 5 suggests that the generalization properties of NanoGPT when certain numbers are excluded from the training data, surpass Matrix Completion and therefore are not simply interpolating the low-rank addition matrix. While the results from Section 3 and 4 do indicate that the model learns our addition algorithm when presented with reversed outputs, the limited length generalization indicates that it does not "truly" learn this algorithm. In the end, while we do not precisely characterize the algorithm learned by the model, we do characterize its tendencies and behaviors, particularly in regard to its sensitivity to the input data format.

Thank you for your appreciative feedback on our paper. Please find the responses to your concerns inline below.

Paper being too compressed The submission is a compressed version of a very long paper. As a result, some of the claims in the introduction, e.g. those relative to length generalization and compositionality, are not discussed in the main paper. The results on other operators, and the impact of pre-training on text, are very hard to assess because the paper provides almost no description of the experimental setting. Finally, some of the figures (e.g. fig. 5) are compressed beyond legibility. This makes the main paper difficult to read, especially starting with section 7 (unless one is ready to read 25 additional pages). I would recommend that the authors either submit the full version of their paper to a journal, or limit it to their results on addition, which are significant enough, and deserve a longer discussion.

We appreciate the reviewer's feedback and are encouraged by the recognition of the significance of our results on addition. We understand that the scope of the paper is extensive, and some content might be compressed, making it difficult to read. In light of the reviewer's suggestion, we would be happy to reduce the scope of this submission to focus solely on addition, moving relevant results from the appendix to the main text to provide a clearer and more comprehensive discussion.

We believe that our work is timely and relevant to the current research landscape, and we believe that presenting our findings at this year's ICLR would be beneficial to the community. We will revise Sections 7 and 8 to limit their reliance on the appendix and improve the overall legibility and coherence of our paper.

different architectures To which extent is your finding on reversing digit order in the output specific to the decoder-only architecture you use? I believe output order would be irrelevant in an encoder-only setting (possible here because the output sequence is guaranteed to be shorter than the input sequence), what about an encoder-decoder architecture?

We intentionally focus on decoder-only architectures in order to focus on understanding the emergent properties in LLMs which are typically decoder-only transformers. For encoder-only models, is the author referring to BERT style architectures? If so, since the model is producing all the digits "simultaneously", the reviewer is right in that the reverse will likely not make any difference. However, in the case of encoder-decoder architectures, since the model is still producing digits from left-to-right, reversing the digit order will give it the opportunity to learn a simpler function (see Lemma 1 and 2). We are currently running experiments on both these architectures and hope to report the results by the end of the rebuttal period.

In the paragraph on balancing digit, you say : "For instance, in the case of 3-digit addition, random sampling results in a meager 0.01% probability of selecting a 1-digit number." The probability of selecting a 1-digit operand should be 1%, right?

It is $0.01\%$ since there are only $100$ samples containing 1-digit numbers among the $10^6$ 3-digit addition samples. The reviewer is right in that the probability of selecting a single 1-digit operand is $1\%$, but the probability that either operand is 1-digit is significantly lower.

different lengths What are the performances of the model when operands have different lengths (e.g. 2 + 312, or 546 + 7)?

We report the following result on different digit length combinations. Each row and column represent the number of digits in the first and second operand, respectively.

For plain, performance is lower for lower-digit combinations. It is also interesting to note that the performance is not necessarily symmetric in the lengths of the operands. We conjecture that this is because despite our balanced sampling approach, the number of samples of $(1, 3)$ digit or $(2,1)$ digit sums are relatively few in our training data.

For reverse, performance is fairly consistent for different digit combinations since the model learns addition almost perfectly.

We appreciate the reviewer's thoughtful feedback on our paper. We are glad to hear that you find our paper is well-written, and is technically sound. We respond to your questions inline below.

Shortcoming in terms of Novelty The study's main shortcoming in terms of novelty stems from its reliance on previously established methods and datasets, particularly the use of reasoning-augmented data, which has been prevalent in enhancing model performance. The authors explicitly acknowledge that their work doesn't break new ground in terms of the types of training data or in achieving peak performance with minimal model parameters. However, the research sets itself apart through its meticulous ablation studies and in-depth exploration of various sampling techniques, training data formats, data source mixing ratios, and model scales. Additionally, while they provide certain novel theoretical explanations for observed phenomena, their primary emphasis on arithmetic isn't for its intrinsic importance but as an easily testable emergent skill to better understand emergent phenomena in models.

We acknowledge the reviewer's concern regarding the novelty of our work. While we agree that our study does not introduce new datasets or techniques, we would like to emphasize that the novelty of our research lies in the extensive ablation studies and the thorough exploration of various factors influencing model performance in a well-specified setting. Our work's contribution is the comprehensive understanding gained through dissecting the emergence of arithmetic in small transformer models, particularly highlighting the importance of data sampling and data formatting, including seemingly simple methods like reversing the outputs. Our findings contribute to a more in-depth study of how transformers develop arithmetic skills, and have a broader implications on synthetic training data generation for LLM training, which we believe is valuable for the research community. This work is particularly timely given the current state of research in this area and will hopefully encourage further exploration and development of more sophisticated techniques to better understand emergent phenomena in models.

Q1. scaling methodology How do the authors envision scaling their methodology beyond GPT-3?

We would like to emphasize that since our approach primarily focuses on data sampling and formatting techniques, it is directly applicable to models of different scales. Our experiments on finetuning GPT-3 (Table 2) show that our findings extend from NanoGPT to GPT-2 and even GPT-3. We see no reason why it should not extend beyond that as well. We are currently running experiments on finetuning GPT-4 and hope to include the results by the end of the rebuttal period.

Q2. compatibility with subquadratic variants Additionally, do they believe their approach is compatible with subquadratic transformer variants?

We have not yet explored subquadratic Transformer variants such as sparse attention, or linformer. Since our analysis is primarily focussed on the pre-training data, we expect that it should extend to those variants as well. We are currently running our experiments for these variants and hope to include the results by the end of the rebuttal. We will surely include it in the final manuscript.

Q2. Length of CoT Could the reason for better performance of the detailed chain of thought model simply be its length?

This is a great question! We have several results that demonstrate that the improved performance of the detailed scratchpad format is attributed to the intermediate steps rather than the length of the data format. For instance, in Section B.3. in the Appendix, we highlight the importance of designing intermediate steps by comparing two different versions of the detailed scratchpad format for subtraction. In Section B.4, we examine the effect of recording correct digit-wise sum (A) and carry (C) information versus randomizing A and C in the simplified scratchpad format, showing that addition is best learned with accurate A and C values.

Furthermore, in the revised version of Section B.3., we have included an additional experiment in which we replace the intermediate steps of the detailed scratchpad format with random tokens, keeping the overall data format length constant and find that the performance is significantly worse than using the correct intermediate steps. This indicates that the intermediate steps play a crucial role in the model's performance, rather than the mere length of the data format.

Q3. Table 1. In Table 1, could you report a confidence interval? In the caption, you say in some cases it improves (when some numbers are excluded from the training data)? If the difference here is significant (it actually is a an improvement), what is the intuitive explanation? You mention a regularisation effect, could you elaborate on that?

By "some cases", we refer to the columns where we exclude 100 and 500 numbers for plain and all 3 cases for reverse, since the overall accuracy is higher than in the case when we exclude no numbers. We refer to this as a regularization effect since it can be thought of as analogous to data augmentation with cropped images used while training vision models. However, as the reviewer points out, this is somewhat handwavy.

A more reasonable explanation of this is a generalization of the matrix completion argument that we discuss in Section 5. Note that learning how to add two-digit numbers could be viewed as completing a fourth-order, rank-4 tensor (rank-4 in the CP sense). $ab + cd = 10a + b + 10c + d$, which can then be represented as the sum of fourth order rank-1 tensors (rank-1 in the CP sense): $10* T_1 + T_2 + 10*T_3 + T_4$ where $T_1$ is $\mathbf{N}\circ \mathbf{1} \circ \mathbf{1} \circ \mathbf{1}$, $T_2$ is $\mathbf{1}\circ \mathbf{N}\circ \mathbf{1}\circ \mathbf{1}$ and so on. Note that $\mathbf{N} \in \mathbb{R}^{10}$ here denotes the vector $[0, 1, 2, \dots, 9]$ and $\mathbf{1} \in \mathbb{R}^{10}$ denotes the vector of 1s. While this argument is still informal, it now shows that the sample complexity of learning the "addition-map" is O(\\#digits) rather than O(10^{\\#digits}). It also reveals that seeing each digit in each position enough times to learn $T_i$ might be sufficient.

We have updated Table 1. and Table 4. to contain the confidence interval on the performance of 5 models trained with different seeds.

Q4. Figure 4. scrachpad output In Figure 4, are the models with scratch pad also with reversed output?

No. The models with scratchpad do not have reversed output. Only the reverse format is trained with reversed output.

Q5. Matrix Completion The paper argues that the sudden jump in accuracy can be explained if addition is formulated as 2-rank matrix completion, but the generalization behaviour of the model can not be explained by this. Could you elaborate how this explanation holds even though it is not fully consistent with the behaviour of the model?

We use the phase transition as a hint that matrix completion (MC) could explain part of the behavior of the model learning addition. As the reviewer correctly points out, our discussion in section 5 on generalization shows that NanoGPT is not performing exact matrix completion. We merely remark that the $O(n)$ phase transition of NanoGPT is almost identical to that of the phase transition in MC, and therefore provides some insights into its learning behavior.

Q6. Training iterations for smaller model When comparing models with different sizes, is the smaller model trained for longer?

For plain and reverse formats, NanoGPT and GPT-2 are both trained for the same number of iterations - sufficient for the training loss to converge. The learning rate is chosen from {1e-3, 5e-4, 1e-4, 5e-5} based on validation loss. For the scratchpad format, NanoGPT is trained longer since the number of tokens per sample is higher and it requires more iterations to converge. For the precise details of hyperparameter choices, we refer the reviewer to Table 13 and 14 in the Appendix.

We thank the reviewer for their detailed feedback and for acknowledging that our experiments and analyses are extensive and interesting. Please find our responses to each of the concerns you have raised below.

Generalizability These results are based on using character level tokenizer which simplifies things a lot when it comes to such tasks and is potentially one of the biggest challenges for LLMs to perform well on these tasks.

We agree that the character-level tokenizer simplifies the setting, but the main reason we chose arithmetic, is to find a setting where it is simple enough to ablate all the different factors. Note that for GPT-2, we also ran experiments with BPE tokenization (tiktoken), as shown in Figure 25. While the results may not completely generalize, our experiments on Chain-of-Thought data and sampling on models of varying scales illustrate that most of our findings remain applicable. We hope that this will encourage more research on systematically ablating individual components.

Value beyond analysis/understanding purposes Methods like reversing the output seem a bit tacky. I agree it is interesting to see that these types of modification to the input/output impacts the results, making it easier for the model to learn the task, but I am not sure if they can have any value beyond analysis/understanding purposes.

We appreciate your perspective on the value of our work and would like to emphasize that our primary objective is to analyze and understand how these models learn arithmetic skills and, if possible, determine more efficient ways to elicit such learning. While reversing may seem somewhat tacky, the importance of coming up with data formatting techniques to accelerate the learning process is more important than ever. We believe that work along this line is especially important now with growing model scales and growing scarcity of tokens. We hope that our work encourages more research along the lines of Gunasekar et al. "Textbooks are all you need" and leads to creation of small, well curated datasets for fast LLM training.

Sparse exploration space While some of the experiments presented in the paper are conceptually very interesting they seem to be exploring the space a bit sparsely which makes it harder to make any firm conclusions.

We appreciate the reviewer's observation regarding the sparse exploration space in our study. We acknowledge that it might be challenging to draw firm conclusions, but we argue that this is primarily due to the inherent complexity of studying LLMs. Given the scale and the multitude of variables involved, isolating individual aspects of the problem is a daunting task. This is the reason we focused on a simple problem like arithmetic and conducted extensive ablations to gain insights into data formatting, sampling, and their impact on teaching skills to small transformer models. However, we would like to point out that we have results on GPT-2 and various scales of GPT-3 models in Sections 8. and Appendix F. We are currently running experiments on finetuning GPT-4 and hope to post those by the end of the rebuttal period.

In our study, we aimed to examine parameters that are commonly used in practice, such as training data formatting, data scale, and few-shot prompting. We understand the reviewer's concern about the scope of our experiments and would be open to reducing it if deemed valuable. However, we believe that each experiment in our study offers useful insights, which we have attempted to summarize in each subsection of the paper. Our work serves as a foundation for future research in understanding the learning process of language models and optimizing their performance using data formatting and other techniques.

Q1. CoT generated data When chain of thought is applied during training, does the model also generate the chain of thought during inference? Is there any correlation between the validity of the chain of thought (during inference) and the correctness of the answer?

When we train the model with CoT data from scratch, it does generate output in a chain-of-thought style during inference. As expected, there is indeed a strong correlation between the validity of the chain of thought (during inference) and the correctness of the answer. We evaluate the scratchpad match score to be the fraction of intermediate computations that are correct and plot it with the final accuracy. We observe a strong correlation between the two, as can be seen here. We observe that errors in intermediate steps propagate to subsequent steps, leading to incorrect answers. However, for some of the less accurate models (say NanoGPT trained on 250 samples), the model chances upon the correct answer even after getting some of the intermediate computations wrong.