Elamur Research⌕
☰
home / papers / 2024 / Looped Transformers for Length Generalization

Looped Transformers for Length Generalization

Paper · 2024 · Ying Fan, Yilun Du, Kannan Ramchandran et al.

Summary

Transformers trained from scratch on arithmetic or algorithmic tasks usually fail on inputs longer than those seen during training, partly because their depth is fixed. Fan et al. define n-RASP-L, a class of tasks solvable by repeating one RASP-L step a number of times that grows with input length n. They show that Copy, Parity and Addition belong to this class. They then train a Looped Transformer, a single decoder block applied repeatedly with shared weights. Only the final answer is supervised, after the number of iterations the task requires, so no scratchpad or intermediate labels are needed. Prediction uses a full-output (FOP) scheme rather than next-token prediction. At test time, the loop count either follows the known step count or is chosen by a confidence-based stopping rule. On six synthetic tasks, the looped models extrapolate to longer inputs considerably better than fixed-depth NTP and FOP baselines and their pause-token variants. The work suggests that adaptive depth is a practical route to models that learn the actual algorithm rather than length-specific shortcuts.

Key points

  • The n-RASP-L class consists of a RASP-L pre-processing step, a RASP-L step repeated T(n) times, and a RASP-L post-processing step. The constructions use n iterations for Parity and Copy and n+1 for Addition.
  • The evaluation covers Parity, Copy, Addition, Multiplication, Binary Sum and Unique Set. The method is compared with vanilla NTP, vanilla FOP, NTP-Pause and FOP-Pause (20 pause tokens), FOP-Looped-Fixed, and a variant without input injection.
  • The looped block is a decoder-only GPT-2 with NoPE and 256-dim embeddings. Each iteration re-adds the original input embedding, following Yang et al. (2024). Block depth (1–4 layers) and head count (8–64) vary by task.
  • Training draws online i.i.d. samples, increases maximum length through a curriculum, and uses AdamW with a cosine-decayed learning rate from 1e-4, batch size 64, for 100k steps. Parity and Binary Sum also use parameter EMA (0.9999).
  • Results report exact-match accuracy on 6,400 test samples, averaged over five seeds. Because more loops are run, inference costs somewhat more than fixed-depth baselines of comparable training compute.

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

🎓University of California, BerkeleyBerkeley · United States243 works · 251,774 citations🎓Massachusetts Institute of TechnologyCambridge · United States274 works · 82,013 citations🎓University of Wisconsin–MadisonMadison · United States60 works · 49,213 citations

 Fields, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksComputational Theory and Mathematics · 496 works

Cite

Cite as Ying Fan, Yilun Du, Kannan Ramchandran and Kangwook Lee. 2024. Looped Transformers for Length Generalization. arXiv:2409.15647.

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

@misc{fan2024looped,
  title = {{Looped Transformers for Length Generalization}},
  author = {Ying Fan, Yilun Du, Kannan Ramchandran and Kangwook Lee},
  year = {2024},
  doi = {10.48550/arxiv.2409.15647},
  eprint = {2409.15647},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2409.15647}
}

References

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

  1. Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
  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. Advances in Neural Information Processing Systems, 35:38546–38556, 2022.
  3. Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in neural information processing systems, 32, 2019. find
  4. Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. In 8th ICML Workshop on Automated Machine Learning (AutoML), 2021.
  5. Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35:20232–20242, 2022. find
  6. Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In International Conference on Machine Learning, pp. 527–536. PMLR, 2017. find
  7. Hanseul Cho, Jaeyoung Cha, Pranjal Awasthi, Srinadh Bhojanapalli, Anupam Gupta, and Chulhee Yun. Position coupling: Leveraging task structure for improved length generalization of transformers. arXiv preprint arXiv:2405.20671, 2024. find
  8. Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400, 2023. — [arXiv:2309.17400](https://arxiv.org/abs/2309.17400) find
  9. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems, 2021. URL https://arxiv.org/abs/2110.14168, 2021. ↗
  10. Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819, 2018. find
  11. Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, et al. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.02098, 2022.
  12. Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Learning iterative reasoning through energy minimization. In International Conference on Machine Learning, pp. 5570–5582. PMLR, 2022. find
  13. Yilun Du, Jiayuan Mao, and Joshua B Tenenbaum. Learning iterative reasoning through energy diffusion. In Forty-first International Conference on Machine Learning, 2024.
  14. Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583–30598, 2022.
  15. Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, pp. 11398–11442. PMLR, 2023. find
  16. Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Contextual position encoding: Learning to count what’s important. arXiv preprint arXiv:2405.18719, 2024. — [arXiv:2405.18719](https://arxiv.org/abs/2405.18719) find
  17. Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint arXiv:2310.02226, 2023. find
  18. Kaiying Hou, David Brandfonbrener, Sham Kakade, Samy Jelassi, and Eran Malach. Universal length generalization with turing programs. arXiv preprint arXiv:2407.03310, 2024. — [arXiv:2407.03310](https://arxiv.org/abs/2407.03310) find
  19. Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
  20. Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. Advances in neural information processing systems, 35: 22199–22213, 2022.
  21. Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=dsUB4bst9S. ↗
  22. Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang, Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. arXiv preprint arXiv:2310.04418, 2023. find
  23. Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale generation: Learning to solve and explain algebraic word problems. arXiv preprint arXiv:1705.04146, 2017. find
  24. Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, et al. Transformers can do arithmetic with the right embeddings. arXiv preprint arXiv:2405.17399, 2024.
  25. 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.
  26. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022.
  27. 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
  28. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  29. Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint arXiv:2305.16843, 2023. find
  30. Mahdi Sabbaghi, George Pappas, Hamed Hassani, and Surbhi Goel. Explicitly encoding structural symmetry is key to length generalization in arithmetic tasks. arXiv preprint arXiv:2406.01895, 2024.
  31. Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
  32. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  33. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022.
  34. Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pp. 11080–11090. PMLR, 2021. find
  35. Zhaofeng Wu, Linlu Qiu, Alexis Ross, Ekin Akyürek, Boyuan Chen, Bailin Wang, Najoung Kim, Jacob Andreas, and Yoon Kim. Reasoning or reciting? exploring the capabilities and limitations of language models through counterfactual tasks. arXiv preprint arXiv:2307.02477, 2023. find
  36. Liu Yang, Kangwook Lee, Robert D Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. In The Twelfth International Conference on Learning Representations, 2024. find
  37. Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Joshua M. Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. In The Twelfth International Conference on Learning Representations, 2024a. URL https://openreview.net/f… ↗
  38. Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Transformers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024b. n-RASP-L PROGRAMS A

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/fan-2024-looped-transformers-length-generalization/