Elamur Research⌕
☰
home / papers / 2023 / Teaching Arithmetic to Small Transformers

Teaching Arithmetic to Small Transformers

Paper · 2023 · Nayoung Lee, Kartik Sreenivasan, Jason D. Lee et al.

Summary

The paper asks which factors let a next-token predictor acquire arithmetic. It uses a controlled setting: small decoder-only transformers (NanoGPT and GPT-2, from 10.6M up to 124M parameters) trained from random initialization with character-level tokens. Tasks include addition, subtraction, multiplication, square root and sine.

The authors run extensive ablations over:

  • Data format: plain, reversed output, and two scratchpad variants.
  • Sampling: balancing examples by digit count and number of carries.
  • Other factors: mixing in Shakespeare text, few-shot prompting, fine-tuning pretrained GPT-2 and GPT-3, and model scale.

Main results. Writing sums the conventional way is a poor teaching signal. Emitting the answer starting from the lowest digit, or spelling out intermediate steps, cuts the number of examples needed and raises accuracy, even with no pretraining. Accuracy jumps abruptly as training data grows, which the authors link to low-rank matrix completion.

Limitations. Detailed formats cost many more tokens. The models also fail on digit lengths they never saw, which suggests they learn a length-bound mapping rather than a true algorithm.

Why it matters. The work isolates how data design, rather than scale alone, drives emergent skills.

Key points

  • Four addition formats are compared: Plain, Reverse, Simplified Scratchpad and Detailed Scratchpad. Plain never hits 100% accuracy, and the sample cost of the others falls as the format becomes more detailed.
  • An n-digit addition table forms a rank-2 matrix. This makes learning addition analogous to low-rank matrix completion, and NanoGPT generalizes better than standard completion solutions.
  • Joint training on all five operations helps each individual task. Moving from zero-shot to 1-shot prompting gives a large accuracy boost, while additional examples add little.
  • Fine-tuning pretrained GPT-2 and GPT-3 (davinci) works with few samples, but non-standard formats such as Reverse can clash with the patterns learned in pretraining. Larger scale helps but is not required.
  • No length generalization was observed beyond the trained digit lengths. Code is available at github.com/lee-ny/teaching_arithmetic.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Princeton UniversityPrinceton · United States113 works · 48,612 citations🎓University of Wisconsin–MadisonMadison · United States54 works · 49,089 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 works

Cite

Cite as Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee and Dimitris Papailiopoulos. 2023. Teaching Arithmetic to Small Transformers. arXiv:2307.03381.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2307.03381 · arXiv 2307.03381

@misc{lee2023teaching,
  title = {{Teaching Arithmetic to Small Transformers}},
  author = {Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee and Dimitris Papailiopoulos},
  year = {2023},
  doi = {10.48550/arxiv.2307.03381},
  eprint = {2307.03381},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2307.03381}
}

References

71 references, 35 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring length generalization in large language models. arXiv preprint arXiv:2207.04901, 2022.
  2. Bowman, S. R. Can recursive neural tensor networks learn logical reasoning? arXiv preprint arXiv:1312.6192, 2013. find
  3. Bowman, S. R., Potts, C., and Manning, C. D. Recursive neural networks for learning logical semantics. CoRR, abs/1406.1827, 5, 2014. find
  4. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  5. Bubeck, S., Chandrasekaran, V., Eldan, R., Gehrke, J., Horvitz, E., Kamar, E., Lee, P., Lee, Y. T., Li, Y., Lundberg, S., et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712, 2023.
  6. Cai, J., Shin, R., and Song, D. Making neural programming architectures generalize via recursion. arXiv preprint arXiv:1704.06611, 2017. find
  7. Charton, F. Linear algebra with transformers. arXiv preprint arXiv:2112.01898, 2021.
  8. Charton, F. What is my math transformer doing?–three results on interpretability and generalization. arXiv preprint arXiv:2211.00170, 2022. find
  9. Chen, X., Liu, C., and Song, D. Towards synthesizing complex programs from input-output examples. arXiv preprint arXiv:1706.01284, 2017. find
  10. Chen, X., Lin, M., Schärli, N., and Zhou, D. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05128, 2023. find
  11. Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
  12. Chung, H. W., Hou, L., Longpre, S., Zoph, B., Tay, Y., Fedus, W., Li, E., Wang, X., Dehghani, M., Brahma, S., et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416, 2022. find
  13. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
  14. Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, Ł. Universal transformers. arXiv preprint arXiv:1807.03819, 2018. find
  15. Drozdov, A., Schärli, N., Akyürek, E., Scales, N., Song, X., Chen, X., Bousquet, O., and Zhou, D. Compositional semantic parsing with large language models. arXiv preprint arXiv:2209.15003, 2022. find
  16. Gadre, S. Y., Ilharco, G., Fang, A., Hayase, J., Smyrnis, G., Nguyen, T., Marten, R., Wortsman, M., Ghosh, D., Zhang, J., et al. Datacomp: In search of the next generation of multimodal datasets. arXiv preprint arXiv:2304.14108, 2023.
  17. Giannou, A., Rajput, S., Sohn, J.-y., Lee, K., Lee, J. D., and Papailiopoulos, D. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196, 2023. find
  18. Hajij, M., Zamzmi, G., Ramamurthy, K. N., and Saenz, A. G. Data-centric ai requires rethinking data notion. arXiv preprint arXiv:2110.02491, 2021. — [arXiv:2110.02491](https://arxiv.org/abs/2110.02491) find
  19. Hanna, M., Liu, O., and Variengien, A. How does gpt-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model. arXiv preprint arXiv:2305.00586, 2023.
  20. Huang, J., Gu, S. S., Hou, L., Wu, Y., Wang, X., Yu, H., and Han, J. Large language models can self-improve. arXiv preprint arXiv:2210.11610, 2022.
  21. Kaiser, Ł. and Sutskever, I. Neural gpus learn algorithms. arXiv preprint arXiv:1511.08228, 2015. find
  22. Karpathy, A. char-rnn. https://github.com/karpathy/char-rnn, 2015. ↗
  23. Karpathy, A. Andrej karpathy’s lightweight implementation of medium-sized gpts. GitHub, 2022. URL https://github.com/karpathy/nanoGPT. ↗
  24. Kim, J., Hong, G., Kim, K.-m., Kang, J., and Myaeng, S.-H. Have you seen that number? investigating extrapolation in question answering models. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7031–7037, 2021. find
  25. Király, F. J., Theran, L., and Tomioka, R. The algebraic combinatorial approach for low-rank matrix completion. J. Mach. Learn. Res., 16(1):1391–1436, 2015. find
  26. Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. arXiv preprint arXiv:2205.11916, 2022.
  27. Lake, B. and Baroni, M. Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks. In International conference on machine learning, pp. 2873–2882. PMLR, 2018. find
  28. Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let’s verify step by step. arXiv preprint arXiv:2305.20050, 2023.
  29. Ling, W., Yogatama, D., Dyer, C., and Blunsom, P. Program induction by rationale generation: Learning to solve and explain algebraic word problems. arXiv preprint arXiv:1705.04146, 2017. find
  30. Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Exposing attention glitches with flip-flop language modeling. arXiv preprint arXiv:2306.00946, 2023. find
  31. Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zettlemoyer, L. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837, 2022.
  32. MosaicML. Introducing mpt-7b: A new standard for open source, commercially usable llms, 2023. URL www.mosaicml.com/blog/mpt-7b. Accessed: 2023-05-05. find
  33. Motamedi, M., Sakharnykh, N., and Kaldewey, T. A data-centric approach for training deep neural networks with less data. arXiv preprint arXiv:2110.03613, 2021. — [arXiv:2110.03613](https://arxiv.org/abs/2110.03613) find
  34. Nogueira, R., Jiang, Z., and Lin, J. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2102.13019, 2021.
  35. Nye, M., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D., Lewkowycz, A., Bosma, M., Luan, D., et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021.
  36. Ontanón, S., Ainslie, J., Cvicek, V., and Fisher, Z. Making transformers solve compositional tasks. arXiv preprint arXiv:2108.04378, 2021. find
  37. Pérez, J., Barceló, P., and Marinkovic, J. Attention is turing complete. The Journal of Machine Learning Research, 22(1):3463–3497, 2021. find
  38. Peterson, J., Meylan, S., and Bourgin, D. Open clone of openai’s unreleased webtext dataset scraper. GitHub, 2019. URL https://github.com/jcpeterson/openwebtext. ↗
  39. Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5, 2023. find
  40. Qian, J., Wang, H., Li, Z., Li, S., and Yan, X. Limitations of language models in arithmetic and symbolic induction. arXiv preprint arXiv:2208.05051, 2022. find
  41. Radford, A. and Narasimhan, K. Improving language understanding by generative pre-training. 2018. find
  42. Rajani, N. F., McCann, B., Xiong, C., and Socher, R. Explain yourself! leveraging language models for commonsense reasoning. arXiv preprint arXiv:1906.02361, 2019.
  43. Razeghi, Y., Logan IV, R. L., Gardner, M., and Singh, S. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXiv:2202.07206, 2022. — [arXiv:2202.07206](https://arxiv.org/abs/2202.07206) find
  44. Recht, B. A simpler approach to matrix completion. Journal of Machine Learning Research, 12(12), 2011. find
  45. Reed, S. and De Freitas, N. Neural programmer-interpreters. arXiv preprint arXiv:1511.06279, 2015. find
  46. Roy, S. and Roth, D. Solving general arithmetic word problems. arXiv preprint arXiv:1608.01413, 2016.
  47. Shaw, P., Uszkoreit, J., and Vaswani, A. Self-attention with relative position representations. arXiv preprint arXiv:1803.02155, 2018.
  48. Shi, F., Suzgun, M., Freitag, M., Wang, X., Srivats, S., Vosoughi, S., Chung, H. W., Tay, Y., Ruder, S., Zhou, D., et al. Language models are multilingual chain-of-thought reasoners. arXiv preprint arXiv:2210.03057, 2022.
  49. Srivastava, A., Rastogi, A., Rao, A., Shoeb, A. A. M., Abid, A., Fisch, A., Brown, A. R., Santoro, A., Gupta, A., Garriga-Alonso, A., et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206.04615, 2022.
  50. Sun, Y., Dong, L., Patra, B., Ma, S., Huang, S., Benhaim, A., Chaudhary, V., Song, X., and Wei, F. A length-extrapolatable transformer. arXiv preprint arXiv:2212.10554, 2022.
  51. Sutskever, I., Vinyals, O., and Le, Q. V. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014.
  52. Talmor, A., Tafjord, O., Clark, P., Goldberg, Y., and Berant, J. Leap-of-thought: Teaching pre-trained models to systematically reason over implicit knowledge. Advances in Neural Information Processing Systems, 33:20227–20237, 2020.
  53. Tay, Y., Wei, J., Chung, H. W., Tran, V. Q., So, D. R., Shakeri, S., Garcia, X., Zheng, H. S., Rao, J., Chowdhery, A., et al. Transcending scaling laws with 0.1% extra compute. arXiv preprint arXiv:2210.11399, 2022.
  54. Thoppilan, R., De Freitas, D., Hall, J., Shazeer, N., Kulshreshtha, A., Cheng, H.-T., Jin, A., Bos, T., Baker, L., Du, Y., et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239, 2022.
  55. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models, 2023.
  56. Uesato, J., Kushman, N., Kumar, R., Song, F., Siegel, N., Wang, L., Creswell, A., Irving, G., and Higgins, I. Solving math word problems with process-and outcome-based feedback. arXiv preprint arXiv:2211.14275, 2022.
  57. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  58. Wallace, E., Wang, Y., Li, S., Singh, S., and Gardner, M. Do nlp models know numbers? probing numeracy in embeddings. arXiv preprint arXiv:1909.07940, 2019.
  59. Wang, C., Zheng, B., Niu, Y., and Zhang, Y. Exploring generalization ability of pretrained language models on arithmetic and logical reasoning. In Natural Language Processing and Chinese Computing: 10th CCF International Conference, NLPCC 2021, Qingdao, China, October 13–17, 2021, Proceedings, Part… find
  60. Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., and Zhou, D. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171, 2022.
  61. Webb, T., Holyoak, K. J., and Lu, H. Emergent analogical reasoning in large language models. arXiv preprint arXiv:2212.09196, 2022. find
  62. Wei, C., Chen, Y., and Ma, T. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35: 12071–12083, 2022a. find
  63. Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682, 2022b.
  64. Wei, J., Wang, X., Schuurmans, D., Bosma, M., Chi, E., Le, Q., and Zhou, D. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022c.
  65. Yuan, Z., Yuan, H., Tan, C., Wang, W., and Huang, S. How well do large language models perform in arithmetic tasks? arXiv preprint arXiv:2304.02015, 2023.
  66. Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S. J., and Kumar, S. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077, 2019.
  67. Zaremba, W. and Sutskever, I. Learning to execute. arXiv preprint arXiv:1410.4615, 2014. find
  68. Zaremba, W., Kurach, K., and Fergus, R. Learning to discover efficient mathematical identities. Advances in Neural Information Processing Systems, 27, 2014. find
  69. Zelikman, E., Mu, J., Goodman, N. D., and Wu, Y. T. Star: Self-taught reasoner bootstrapping reasoning with reasoning. 2022. find
  70. Zhou, D., Schärli, N., Hou, L., Wei, J., Scales, N., Wang, X., Schuurmans, D., Bousquet, O., Le, Q., and Chi, E. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205.10625, 2022a.
  71. Zhou, H., Nova, A., Larochelle, H., Courville, A., Neyshabur, B., and Sedghi, H. Teaching algorithmic reasoning via in-context learning. arXiv preprint arXiv:2211.09066, 2022b.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/lee-2023-teaching-arithmetic-small-transformers/