Elamur Research⌕
☰
home / papers / 2023 / Length Generalization in Arithmetic Transformers

Length Generalization in Arithmetic Transformers

Paper · 2023 · Samy Jelassi, Stéphane d'Ascoli, Carles Domingo-Enrich et al.

Summary

The paper asks whether transformers can learn integer arithmetic on short operands and still be correct on much longer ones. This matters because large numbers rarely appear in text, so language models must extrapolate from small-number examples. The authors train encoder-only models, both a standard transformer and a universal transformer, on four tasks: addition, multiplication, and their modular versions. Training uses 5-digit operands (5 × 3 for multiplication), and testing uses longer inputs. Swapping absolute position embeddings (APE) for relative ones (RPE) is enough for addition: 5-digit training carries over to sums of 15–20 digits. For multiplication, RPE alone does not work. Fine-tuning on long examples needs many samples and causes catastrophic forgetting. The authors instead propose train set priming, which mixes a small fixed set of long examples into the ordinary training data. With this, models extrapolate to 35 × 3 products. The technique uses very little target-distribution data and avoids forgetting. The authors suggest it may also help adapt language models outside arithmetic.

Key points

  • Setup: the vocabulary has 15 tokens (digits, four operator symbols, and <PAD>). The first operand comes from a fixed pool of N_train = 5000 values, and the second is sampled fresh during training. Padding keeps input and output lengths constant.
  • Priming with just 50 long examples out of 5000 (ε = 1%) lets an RPE model trained on 5 × 3 multiplication handle 35 × 3. The number of priming samples needed grows logarithmically with training set size.
  • A single model primed with 500 examples generalizes across first-operand lengths from 6 to 35 digits. Priming also works for APE models, but they need roughly a 10% priming rate.
  • On 20-digit addition, failures cluster on inputs with at least three carries, including two consecutive ones. Wrong answers usually differ in only a few digits, mostly the leading ones.
  • Per-digit learning curves show that models that generalize (RPE, or primed models) learn output positions simultaneously. Models that fail tend to learn one digit position at a time.

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

🏢GoogleMountain View · United States533 works · 443,885 citations🎓Stanford UniversityStanford · United States567 works · 335,190 citations🎓Carnegie Mellon UniversityPittsburgh · United States633 works · 195,218 citations🎓New York UniversityNew York · United States233 works · 104,003 citations🎓Princeton UniversityPrinceton · United States137 works · 53,849 citations🎓École Polytechnique Fédérale de LausanneLausanne · Switzerland56 works · 8,189 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

Cite

Cite as Samy Jelassi, Stéphane d'Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li and François Charton. 2023. Length Generalization in Arithmetic Transformers. arXiv:2306.15400.

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

@misc{jelassi2023length,
  title = {{Length Generalization in Arithmetic Transformers}},
  author = {Samy Jelassi, Stéphane d'Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li and François Charton},
  year = {2023},
  doi = {10.48550/arxiv.2306.15400},
  eprint = {2306.15400},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2306.15400}
}

References

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

  1. Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, et al. Solving rubik’s cube with a robot hand. arXiv preprint arXiv:1910.07113, 2019. find
  2. Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint arXiv:2207.04901, 2022.
  3. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021.
  4. Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Logical extrapolation without overthinking, 2022.
  5. Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
  6. Mirelle Bueno, Carlos Gemmel, Jeffrey Dalton, Roberto Lotufo, and Rodrigo Nogueira. Induced natural language rationales and interleaved markup tokens enable extrapolation in large language models. arXiv preprint arXiv:2208.11445, 2022. find
  7. François Charton. Linear algebra with transformers. arXiv preprint arXiv:2112.01898, 2021.
  8. Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 2021. find
  9. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
  10. Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019. find
  11. Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819, 2018. find
  12. Aurélien Dersy, Matthew D. Schwartz, and Xiaoyuan Zhang. Simplifying polylogarithms with machine learning, 2022. find
  13. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
  14. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
  15. Yann Dubois, Gautier Dagan, Dieuwke Hupkes, and Elia Bruni. Location attention for extrapolation to longer sequences. arXiv preprint arXiv:1911.03872, 2019. find
  16. Kaden Griffith and Jugal Kalita. Solving arithmetic word problems with transformers and preprocessing of problem text. arXiv preprint arXiv:2106.00893, 2021. — [arXiv:2106.00893](https://arxiv.org/abs/2106.00893) find
  17. Christopher Hahn, Frederik Schmitt, Jens U Kreber, Markus N Rabe, and Bernd Finkbeiner. Teaching temporal logics to neural networks. arXiv preprint arXiv:2003.04218, 2020.
  18. Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Curtis Hawthorne, Andrew M Dai, Matthew D Hoffman, and Douglas Eck. An improved relative self-attention mechanism for transformer with application to music generation. 2018. find
  19. Zhiheng Huang, Davis Liang, Peng Xu, and Bing Xiang. Improve transformer models with better relative position embeddings. arXiv preprint arXiv:2009.13658, 2020. find
  20. Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34, 2021. find
  21. Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. Advances in neural information processing systems, 28, 2015. find
  22. Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. arXiv preprint arXiv:1511.08228, 2015. find
  23. Shun Kiyono, Sosuke Kobayashi, Jun Suzuki, and Kentaro Inui. Shape: Shifted absolute position embedding for transformers. arXiv preprint arXiv:2109.05644, 2021. find
  24. Guillaume Lample and François Charton. Deep learning for symbolic mathematics. arXiv preprint arXiv:1912.01412, 2019.
  25. Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv:1909.11942, 2019.
  26. Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. arXiv preprint arXiv:2206.14858, 2022.
  27. Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016. find
  28. Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017. find
  29. Decoupled weight decay regularization. arXiv preprint Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pages 109–165. Elsevier, 1989. find
  30. Yuanliang Meng and Anna Rumshisky. Solving math word problems with double-decoder transformer. arXiv preprint arXiv:1908.10924, 2019. — [arXiv:1908.10924](https://arxiv.org/abs/1908.10924) find
  31. Kenton Murray and David Chiang. Correcting length bias in neural machine translation. arXiv preprint arXiv:1808.10006, 2018.
  32. Benjamin Newman, John Hewitt, Percy Liang, and Christopher D Manning. The eos decision and length extrapolation. arXiv preprint arXiv:2010.07174, 2020. find
  33. Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2102.13019, 2021.
  34. Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021.
  35. Theodoros Palamas. Investigating the ability of neural networks to learn simple modular arithmetic. 2017. find
  36. Stanislas Polu and Ilya Sutskever. Generative language modeling for automated theorem proving. arXiv preprint arXiv:2009.03393, 2020.
  37. Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177, 2022.
  38. Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021. find
  39. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(140):1–67, 2020.
  40. Gregory Rogez, Philippe Weinzaepfel, and Cordelia Schmid. Lcr-net: Localization-classification-regression for human pose. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3433–3441, 2017. find
  41. Jan Rosendahl, Viet Anh Khoa Tran, Weiyue Wang, and Hermann Ney. Analysis of positional encodings for neural machine translation. In Proceedings of the 16th International Conference on Spoken Language Translation, 2019. find
  42. Rasmus Rothe, Radu Timofte, and Luc Van Gool. Dex: Deep expectation of apparent age from a single image. In Proceedings of the IEEE international conference on computer vision workshops, pages 10–15, 2015. find
  43. Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604–609, 2020. find
  44. J Schulman, B Zoph, C Kim, J Hilton, J Menick, J Weng, JFC Uribe, L Fedus, L Metz, M Pokorny, et al. Chatgpt: Optimizing language models for dialogue, 2022. find
  45. Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint arXiv:1803.02155, 2018.
  46. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864, 2021.
  47. Andrew Trask, Felix Hill, Scott Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neural arithmetic logic units. arXiv preprint arXiv:1808.00508, 2018. find
  48. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998–6008, 2017.
  49. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022.
  50. Emily Wenger, Mingjie Chen, François Charton, and Kristin Lauter. Salsa: Attacking lattice cryptography with transformers. arXiv preprint arXiv:2207:04785, 2022. find
  51. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language process…
  52. Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301, 2022. find
  53. Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. arXiv preprint arXiv:2211.09066, 2022.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/jelassi-2023-length-generalization-arithmetic-transformers/