“Teaching Arithmetic to Small Transformers”, Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, Dimitris Papailiopoulos2023-07-07 (LM tokenization, inner monologue (AI), AI emergence, math, meta-learning; backlinks)⁠:

[Twitter; code] Large language models like GPT-4 exhibit emergent capabilities across general-purpose tasks, such as basic arithmetic, when trained on extensive text data, even though these tasks are not explicitly encoded by the unsupervised, next-token prediction objective.

This study investigates how small transformers [NanoGPT & GPT-2], trained from random initialization, can efficiently learn arithmetic operations such as addition, multiplication, and elementary functions like square root, using the next-token prediction objective. We first demonstrate that conventional training data is not the most effective for arithmetic learning, and simple formatting changes can improve accuracy.

This leads to sharp phase transitions as a function of training data scale, which, in some cases, can be explained through connections to low-rank matrix completion. Building on prior work, we then train on chain-of-thought style data that includes intermediate step results. Even in the complete absence of pretraining, this approach and simultaneously improves accuracy, sample complexity, and convergence speed.

We also study the interplay between arithmetic and text data during training and examine the effects of few-shot prompting, pretraining, and model scale. Additionally, we discuss length generalization challenges.

Our work highlights the importance of high-quality, instructive data that considers the particular characteristics of the next-word prediction objective for rapidly eliciting arithmetic capabilities.

1. Introduction
2. Related Works
3. Preliminaries and Experimental Setup
4. Learning Addition in Small Models
   1. Training on Conventional Data
   2. Reversing the Output [cf. RWKV]
5. Connection to Low-Rank Matrix Completion
   1. Addition Tables are Rank-2 Matrices
   2. NanoGPT Generalizes better than Matrix Completion solutions
6. The power of Chain-of-Thought: Incorporating Intermediate Steps in Training Data
   1. Training on Chain-of-Thought Data
   2. The Importance of Intermediate Step Design: Subtraction
   3. The Effect of Noisy Inputs on Accuracy
7. Extending to Longer Digit Addition
   1. Training from Random Initialization
   2. Fine-Tuning from Pretrained Models
   3. Impact of Formats on Fine-Tuning
8. Teaching Arithmetic Operations Beyond Addition
   1. Extended Arithmetic Operations
   2. Jointly Training on All 5 Arithmetic Tasks
9. Mixing Shakespeare with Arithmetic Data
10. Fine-tuning, Scaling, and Pretraining in Larger Models
11. Token Efficiency Across Data Formats
12. Length Generalization
13. Limitations
14. Conclusion
15. Appendix

…Data format and sampling matters: We first observe that teaching a model addition (or any other operation) using standard addition samples, ie. A₃A₂A₁ + B₃B₁B₁ = C₃C₂C₁, is suboptimal, as it requires the model to evaluate the most large digit C₃ of the result first, which depends globally on all the digits of the two summands. By training on samples with reversed results, ie. A₃A₂A₁ + B₃B₁B₁ = C₁C₂C₃, we enable the model to learn a simpler function, improving sample complexity. Additionally, balanced sampling of different “variations” of addition, based on the number of carries and digits involved, further enhances learning. Even in this simple setting, we observe relatively sharp phase transitions 0–100% accuracy as a function of the size of the training data. Although this may seem surprising, we observe that learning an addition map on n digits from random samples is equivalent to completing a low-rank matrix. This connection allows us to offer a reasonable explanation for such phase transitions.

…Chain-of-thought data during training: …We found that CoT-type training data importantly improved learning in terms of both sample complexity and accuracy in agreement with CoT fine-tuning literature (Nye et al 2021; Chung et al 2022), though our observation holds even in the absence of language pretraining. We conjecture that this is because breaking down the required compositional function to be learned into individual components allows the model to learn a higher-dimensional but easier-to-learn function map. In Figure 1, we provide examples of the 4 data formatting methods explored in our work.

…Instructional data/chain-of-thought: The idea of using detailed reasoning training data predates Transformers (Vaswani et al 2017). Ling et al 2017; Cobbe et al 2021; Nye et al 2021 use natural language to generate reasoning steps while Roy & Roth2016; Reed & De Freitas2015; Chen et al 2017; Cai et al 2017 show that symbolic reasoning may suffice. Nogueira et al 2021 note that large number of samples with small digits is important for arithmetic tasks (Yuan et al 2023). Razeghi et al 2022 observe a correlation between the frequency of numbers in the dataset and the performance involving them whereas we find that transformers can learn to add numbers that were not seen during training. Chain-of-thought (Wei et al 2022c) refers to the model’s improved performance when prompted to produce rationale. Zhou et al 2022b show that this can be achieved by providing sufficiently informative exemplars as a few-shot prompt (Brown et al 2020). Zhou et al 2022a showed that least-to-most prompting can help GPT-3 solve problems that can be decomposed into simpler sub-problems. Least-to-most prompting consists of first decomposing a complex problem into easier subproblems, and then sequentially solving these subproblems. We extend this notion to simple addition and show that asking the model to output the least important bit first has a similar effect…Zaremba & Sutskever2014 show that RNNs can learn how to execute simple programs with for-loops provided they are trained with curriculum learning…encoder-decoder models have also been extensively studied in the literature in the context of learning arithmetic (Kim et al 2021; Wang et al 2021). Qian et al 2022; Lightman et al 2023; Uesato et al 2022 explore techniques to improve the arithmetic abilities of pretrained LLMs. Wallace et al 2019 on the other hand, focus on the impact of the learned embeddings…Charton2022 & Charton2021 show that Transformers can learn linear algebra operations with carefully chosen encodings. Hanna et al 2023 use mechanistic interpretability techniques to explain the limited numerical reasoning capabilities of GPT-2.

…We find that learning all arithmetic operations discussed earlier (from addition to square root) can improve the individual performance of each task, and that going from zero-shot to 1-shot prompting (showing one arithmetic example) yields a large accuracy improvement, but there is no large improvement in accuracy by showing more examples. [Consistent with the Bayesian/meta-learning interpretation of task location: there are not many possibilities so more examples don’t offer any relevant evidence—either it knows the necessary task, or it doesn’t.]

…The phase transition of LRMC offers insights into NanoGPT’s learning process. Nevertheless, further experiments clearly demonstrate that NanoGPT’s mechanism for learning addition is fundamentally different from LRMC. It can successfully learn addition even when numbers or digits are intentionally excluded from the training data, thereby exhibiting generalization capabilities that far exceed that of typical LRMC algorithms.

…the Detailed Scratchpad format, which provides even more detailed information, achieves perfect addition with just 1,000 samples. This indicates that incorporating more information enables the model to learn addition more efficiently, requiring fewer examples. We conjecture that this is because breaking down the required compositional function to be learned into individual components allows the model to learn a higher-dimensional but easier-to-learn function map. We note that while CoT-style training enhances sample efficiency, it may not necessarily be the most “token-efficient” approach. We delve into this aspect in more detail in §11. In summary, incorporating scratchpad data and decomposing the addition task into steps offer a promising strategy to improve the performance and efficiency of small models in learning addition from scratch…[but] the detailed scratchpad method uses considerably more tokens compared to other techniques.

…Anil et al 2022 suggests that models can only perform out-of-distribution tasks by combining fine-tuning, prompting, and scratchpad techniques. Nonetheless, there have been cases where length generalization was observed. Nye et al 2021 demonstrated length generalization but only for models with more than 10⁸ parameters.

…Noisy intermediate steps in the scratchpad data: We further investigate the importance of providing accurate intermediate steps in the scratchpad during the training process. While this was inspired by the findings of Min et al 2022, it is inherently different. Min et al 2022 show that using random labels in ICL [in-context learning] demonstrations caused minimal degradation when compared to the gold labels. However, those models were trained on gold labels and then evaluated on multiple downstream tasks. In our setting, the model is trained and evaluated on a single arithmetic task. Further, the final result (or label) is left untouched as the correct answer to the arithmetic operation. We only replace the intermediate steps. The goal of this study is to verify whether the model actually learns to reason using the given intermediate steps or merely uses the scratchpad to improve its expressivity. We compare the performance of training with our simplified scratchpad formatting, which includes accurate A (digit sum) and C (carry) information, with formatting that includes random A, random C, or random A and C for each intermediate step, as depicted in Figure 1.

The results in Figure 9, demonstrate that the inclusion of noisy labels can impede sample efficiency. However, with enough samples, the model ultimately achieves full accuracy. This suggests that while the model is capable of leveraging the information contained in the intermediate steps, it can also gradually learn how to perform addition while disregarding the presence of noisy intermediate steps.

…The results presented in Table 3 reveal intriguing findings. We observe that the reverse format consistently outputs a result that deviates by no more than 1 from the true answer, regardless of whether the preceding outputs O₃O₂ are subjected to random or precise perturbation. This consistency can be explained by Lemma 2, indicating that the reverse format only requires learning a straightforward function of digit-wise addition for each corresponding position, along with the carry-on (0 or 1). Therefore, even with noise in the preceding tokens, the model accurately performs digit-wise addition, albeit with occasional carry-on prediction errors. With an exact accuracy of 81.26% even in the presence of random perturbation, the reverse format demonstrates the model’s ability to rely less on the preceding output tokens, indicating a robust learned output mapping.

…Going from character-level tokenization to BPE: …Figure 20 shows that GPT-2 demonstrates high performance in addition tasks with both character-level tokenization and Tiktoken with spaces between digits. This aligns with the results by Wallace et al 2019, suggesting that character-level tokenization exhibits stronger numeracy capabilities compared to a word or sub-word methods. Furthermore, comparing the models trained from scratch and the models trained from the pretrained model, we observe that fine-tuning a pretrained model results in better performance compared to training a model from scratch.

See Also:

• Evaluating Transformer Language Models on Arithmetic Operations Using Number Decomposition

• It’s Not Just Size That Matters: Small Language Models Are Also Few-Shot Learners

• TinyStories: How Small Can Language Models Be and Still Speak Coherent English?

• Models In a Spelling Bee: Language Models Implicitly Learn the Character Composition of Tokens

Backlinks: