Paper: Teaching Arithmetic to Small Transformers (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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).

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
Please refer to [Figure 21](https://imgur.com/GysI4fO) 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.

Reviewer comment (Experiments):
You claim you try to disentangle the factor of pretraining, but compare NanoGPT and GPT-2 from scratch to a fine-tuned GPT-3, which doesn't disentangle pretraining from scale. Why not compare to GPT-2 fine-tuned?

Author response:
We hope that the reviewer's concern regarding the comparison between NanoGPT, GPT-2 from scratch, and a fine-tuned GPT-3 is addressed in [Figure 25](https://imgur.com/xVCWPSl), where we present results for both NanoGPT and GPT-2 models initialized from scratch and with pretraining. Additionally, we use the TikToken tokenizer for GPT-2 to match GPT-3's tokenization approach. This comparison aims to isolate the effects of pretraining from model scale. We will emphasize these details in the revised manuscript.

Reviewer comment (Experiments):
In summary, I feel like the motivation of the study (to explain emergent arithmetic abilities of LLMs) does not match with the experiments.

Author response:
Thank you for your feedback and for reconsidering some aspects of our work. We understand your concerns regarding the transferability of our findings to realistic settings and the emergence of arithmetic at scale. We have worked hard to conduct extensive experiments and provide insights into arithmetic learning in language models and would appreciate the opportunity to improve our manuscript within this review cycle. We hope that our work can be considered for ICLR, as we believe it is timely and relevant to the community. We are committed to addressing your concerns and are open to modifying our manuscript accordingly to ensure its scientific rigor and relevance.

First, while it is true that the datasets used for training models at scale may not have exactly the same distribution we used in our experiments, our work highlights the importance of data formatting, sampling, and the inclusion of chain of thought steps for fast arithmetic learning. In fact, the improved abilities of the newer models such as GPT-4, Llama-2 highlight that, (as many suspect) these models are likely trained on carefully curated data. Moreover, with the growing token sparsity and prohibitively expensive training runs, it is essential that we choose the data as carefully as possible so as to ensure that these models achieve high quality with as few samples as possible. We also believe that the insights gleaned from our work can be applied to improve finetuning strategies and guide future research in synthetic dataset creation and curation, which has been an area of high interest in the related literature such as Gunasekar et al. "Textbooks are all you need".

Second, we believe that our connection to matrix completion can help explain the sharpness of the phase transition observed in our experiments. Although our data may not resemble real-world datasets exactly, the core principles and mechanisms we uncover can still provide valuable insights into the emergence of arithmetic capabilities in large language models.

We understand that some of our hypotheses may seem overreaching, and we are more than willing to tone down any claims that are not well-supported by the evidence. We would greatly appreciate any specific guidance from the reviewer on which parts of the manuscript should be modified to better align with the presented results.

Reviewer comment (Evaluation):
I don't understand why you would call the learning curve between 1k and 4k examples a phase transition? You're learning a task with a deterministic output, and 100% accuracy is possible and this model happens to learn it between 1 and 4k examples. This doesn't seem to be a fast phase transition when the model is anyway only trained for 6k examples.

Author response:
We would like to clarify that for all experiments, we are training on a fixed set of example. The sample efficiency curve shows the performance varies when training for multiple epochs on training datasets of varying sizes. The term "phase transition" is used to describe the qualitative shift in the model's capability to perform addition—from an inability to grasp addition with an insufficient number of examples to learning perfect addition as the number of samples increase.

Therefore, we refer to learning addition perfectly between 1k and 4k examples as a phase transition since the model goes from $0\%$ accuracy to almost $100\%$ in this range. We would also like to emphasize that this is an extremely small fraction of all possible 3-digit combinations. We stop at 6k since there is no point in training with more samples once the model has learned addition perfectly. However, the "phase transition" would look sharper if we extended the x-axis to the $10^6$ possible samples.

Reviewer comment (Experiments):
It seems like the main reason for the model to learn arithmetic in this setup is because every example it sees is an entirely new one (new pair of operands), so the only strategy to achieve low training loss is actual arithmetic. It would be interesting to learn whether the generalisation behaviour is different when you train on the same set of examples for multiple epochs, which in turn could in fact explain some of the emergence of arithmetic in LLMs who are often trained for only a single epoch.

Author response:
We would like to clarify that for all experiments, we are training on a fixed set of example. The sample efficiency curve shows the performance varies when training for multiple epochs on training datasets of varying sizes. The term "phase transition" is used to describe the qualitative shift in the model's capability to perform addition—from an inability to grasp addition with an insufficient number of examples to learning perfect addition as the number of samples increase.

Therefore, we refer to learning addition perfectly between 1k and 4k examples as a phase transition since the model goes from $0\%$ accuracy to almost $100\%$ in this range. We would also like to emphasize that this is an extremely small fraction of all possible 3-digit combinations. We stop at 6k since there is no point in training with more samples once the model has learned addition perfectly. However, the "phase transition" would look sharper if we extended the x-axis to the $10^6$ possible samples.

Reviewer comment (Experiments):
I'm not sure I understand the exclusion experiments; For example for 1st (LSD) digit exclusion, do you just exclude 5 in the 1st digit? meaning that something like 543 cannot occur, but 453 can? If yes, how is that exclusion? The numbers are represented as a sequence over absolute embeddings, so 5 is not excluded at all but simply held-out at a particular position in the number no?

Author response:
That is exactly right. "Exclusion" here refers to its absence from a particular ordinal position, not its complete removal from the dataset. This approach allows us to investigate the model's ability to generalize arithmetic operations on digits from one position to another where it has never seen it. It is important to note that if a digit is completely excluded from the training data, the model cannot learn the representation of that digit and even its embedding would be random.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
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).

Author response:
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](https://imgur.com/a/SKCkdPK), 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.

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Writing):
Not a weakness since the paper is fairly well written, but here are some typos in the paper:  
1. I believe a latex shortcut is used to represent $A_{3}A_{2}A_{1} + B_{3}B_{2}B_{1} = C_{3}C_{2}C_{1}$ in the paper, because all occurrences of $B_{3}B_{2}B_{1}$ are represented as $B_{3}B_{1}B_{1}$.  
2. Figure 1, in the detailed scratchpad solution, I think a carry ($C = 1$) has been missed after $[1, 2] + [3, 6], A = [5]$.  
3. Figure 4b, it should be just Number of tokens on the x axis instead of Number of unique tokens?  
4. Figure 8: Larger model scale instead of sacale.  
5. No space between $\textit{text}$ and $\textit{and}$ in the last line on page 8.

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

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Reproducibility):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Writing):
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?

Author response:
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.

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

Author response:
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.

| Plain   | 1-digit | 2-digit | 3-digit |
| ----   | ----    | ----    | ----    |
| 1-digit| 72.8%   | 42.2%   | 77.0%   |
| 2-digit| 27.8%   | 60.0%   | 91.4%   |
| 3-digit| 38.6%   | 61.2%   | 97.0%   |

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

| Reverse| 1-digit | 2-digit | 3-digit |
| ----   | ----    | ----    | ----    |
| 1-digit| 100.0%  | 99.8%   | 99.0%   |
| 2-digit| 96.4%   | 100.0%  | 100.0%  |
| 3-digit| 94.0%   | 99.4%   | 100.0%  |

Reviewer comment (Experiments):
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.

Author response:
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.

|        | P(A+B=B+A, both are correct answer) | P(A+B and B+A are both incorrect) | P(A+B=B+A \| both are incorrect) |
| ----   | ----    | ----      | --- |
|  | 0.878 (8700 samples)   | 0.016 (164 samples)  | 0.603 (99/164 samples) |

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

Reviewer comment (Experiments):
Can those results generalize to signed numbers, by adding a sign token to all three integers?

Author response:
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](https://imgur.com/a/2mTeH9m). 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.

Reviewer comment (Experiments):
Is there a chain-of-thought approach suitable to low-rank matrix completion? If so, it would greatly improve its data efficiency, and perhaps allow it to scale to larger operands.

Author response:
Perhaps the reviewer can clarify what they mean by chain-of-thought approach suitable to low-rank matrix completion? For larger operands, we find that the simplified scratchpad is a more token efficient version that is more scalable to larger operands. Furthermore, since the sample complexity improves significantly with chain-of-thought data, it remains efficient to scale to larger operands (see Appendix C. where we train on samples up to 10 digits in length)

One possible option for designing a chain-of-thought approach suitable to low-rank matrix completion is to decompose the two operands to their 1s, 10s, 100s and so on. For example $abc + def = a00 + b0 + c + d00 + e0 + f = ghi$. We tried retraining the model with data in this format and the [results](https://imgur.com/a/rBaza0R) show that this works extremely well. While it is still slightly worse than the (longer) simple scratchpad format, it far outperforms the plain format. The high performance of this format suggests that the tensor completion argument presented in our [Response to Reviewer 9wwh (2/3)](https://openreview.net/forum?id=dsUB4bst9S&noteId=5SmGFcc0O8) might be representative of the way the model learns addition.

Reviewer comment (Experiments):
Even in the case of the ability of models to learn arithmetics, as mentioned in the paper, these results are based on using a character level tokeniser 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. (A parallel work that looks into this: https://openreview.net/forum?id=OinvjdvPjp).

Author response:
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](https://imgur.com/xVCWPSl). 
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.

Reviewer comment (Novelty):
Methods like reversing the output seem a bit tacky. While it is interesting to see that these types of modifications to the input/output impact the results and make it easier for the model to learn the task, it is not clear if they can have any value beyond analysis/understanding purposes.

Author response:
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.

Reviewer comment (Evaluation):
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?

Author response:
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](https://imgur.com/a/WLdu5Ln). 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.

Reviewer comment (Experiments):
Could the reason for better performance of the detailed chain of thought model simply be its length?

Author response:
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.

Reviewer comment (Presentation):
In Figure 4, are the models with scratch pad also with reversed output?

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

Reviewer comment (Experiments):
Do you have any experiments where the different operations are not split into different tasks (where the task contains the mixture of operation)?

Author response:
Yes, jointly training on all five different operations are presented in the Appendix D.2. We also ran an experiment where we took the model trained purely on addition and then finetuned it on samples of subtraction to test whether the model is able to transfer some of its learnings. Interestingly, this model converges much faster than a randomly initialized model albeit to a slightly worse solution. We conjecture that the embeddings learned for addition transfer effectively to subtraction. (See [Figure](https://imgur.com/VdxWE3B))

Reviewer comment (Experiments):
Is there a failure point as you increase the number of digits?

Author response:
We extend our experiments up to 10 digit numbers in Appendix C. and observe no real failure point i.e., we are always able to learn addition within a reasonable number of samples. In fact, the sample efficiency gap between reverse/scratchpad and plain tends to increase with number of digits!