Elamur Research⌕
☰
home / papers / 2023 / Improving Length-Generalization in Transformers via Task Hinting

Improving Length-Generalization in Transformers via Task Hinting

Paper · 2023 · Pranjal Awasthi and Anupam Gupta

Summary

Transformers trained on arithmetic or algorithmic tasks often fail on inputs longer than those seen in training. This paper asks whether a general training technique can fix this without architectural tricks or any longer-length data. The authors propose task hinting: a multitask setup in which a decoder-only transformer, trained from scratch, learns a main task while also learning a simpler related auxiliary task. Training alternates gradient updates between the two.

The main case study is sorting integers from 1 to 100. Training sequences have length at most 20 and are skewed toward very short ones. With a successor-identification hint, accuracy on length-100 sequences rises from under 1% to over 92%. Other plausible hints, such as counting, help much less or even hurt.

Probing and projection-based visualization show which mechanisms the network tends to adopt. From this, the authors derive a theoretical sorting construction and a length-dependent "tempered softmax" in attention, which further improves extrapolation. The approach also helps on a second task, incrementing a number. The work matters because hints that match the model's internal biases can produce strong length generalization.

Key points

  • Sorting setup: 80% of training sequences have length 2–5 and the rest have length 6–20. No data longer than 20 is used for training or model selection.
  • Hint choice matters: the successor hint works best, while the count and fill hints did not help and sometimes hurt. Combining all three hints gave mixed or negative results.
  • Tempered softmax (depth-2 model): with no hints, length-100 accuracy rises from 0.0% to 45.2%. With successor hints, it rises from 52.4% to 64.8%.
  • Increment task: a depth-4 model was trained on 1M examples of length up to 10, with the carry sequence as the hint. At length 15, accuracy is 41% for standard training, 47.7% with hinting and 62.1% with tempered softmax. At length 20, all three methods stay at 4% or below.
  • Open problems: the authors call for a principled way to choose hint tasks and for tests of length-dependent attention parameters in larger language models.

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,884 citations🎓Carnegie Mellon UniversityPittsburgh · United States633 works · 195,217 citations

 Fields, topics and concepts

📚 Shelves 4The 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 works

Cite

Cite as Pranjal Awasthi and Anupam Gupta. 2023. Improving Length-Generalization in Transformers via Task Hinting. arXiv:2310.00726.

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

@misc{awasthi2023improving,
  title = {{Improving Length-Generalization in Transformers via Task Hinting}},
  author = {Pranjal Awasthi and Anupam Gupta},
  year = {2023},
  doi = {10.48550/arxiv.2310.00726},
  eprint = {2310.00726},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2310.00726}
}

References

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

  1. 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
  2. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  3. Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
  4. Xi Chen, Xiao Wang, Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Grycner, Basil Mustafa, Lucas Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.06794, 2022.
  5. Tao Tu, Shekoofeh Azizi, Danny Driess, Mike Schaekermann, Mohamed Amin, Pi-Chuan Chang, Andrew Carroll, Chuck Lau, Ryutaro Tanno, Ira Ktena, et al. Towards generalist biomedical AI. arXiv preprint arXiv:2307.14334, 2023. find
  6. 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.
  7. 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.
  8. 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
  9. Samy Jelassi, Stéphane d’Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306.15400, 2023.
  10. Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization on the unseen, logic reasoning and degree curriculum. arXiv preprint arXiv:2301.13105, 2023. find
  11. 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.
  12. 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.
  13. Yann Dubois, Gautier Dagan, Dieuwke Hupkes, and Elia Bruni. Location attention for extrapolation to longer sequences. arXiv preprint arXiv:1911.03872, 2019. find
  14. 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
  15. Benjamin Newman, John Hewitt, Percy Liang, and Christopher D Manning. The EOS decision and length extrapolation. arXiv preprint arXiv:2010.07174, 2020. find
  16. Eran Malach. Auto-regressive next-token predictors are universal learners. arXiv:2309.06979, 2023. arXiv preprint Tiedong Liu and Bryan Kian Hsiang Low. Goat: Fine-tuned LLaMA outperforms GPT-4 on arithmetic tasks. arXiv preprint arXiv:2305.14201, 2023.
  17. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  18. 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.
  19. Neel Nanda, Lawrence Chan, Tom Liberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023a.
  20. Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring a sequence model trained on a synthetic task. arXiv preprint arXiv:2210.13382, 2022. find
  21. Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models, 2023b.
  22. Michael Crawshaw. Multi-task learning with deep neural networks: arXiv:2009.09796, 2020. — [arXiv:2009.09796](https://arxiv.org/abs/2009.09796) find
  23. Diederik P Kingma and Jimmy Ba. arXiv:1412.6980, 2014.
  24. A survey. arXiv preprint Adam: A method for stochastic optimization. arXiv preprint Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016. find
  25. William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843–856, 2022. find
  26. Ashok K. Chandra, Larry J. Stockmeyer, and Uzi Vishkin. Constant depth reducibility. SIAM J. Comput., 13(2):423–439, 1984. doi: 10.1137/0213028. URL https://doi.org/10.1137/0213028. — [doi:10.1137/0213028](https://doi.org/10.1137/0213028) ↗
  27. Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL http://github.com/google/flax. ↗

Updated

↑ back to top

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