Elamur Research⌕
☰
home / papers / 2024 / Do language models plan ahead for future tokens?

Do language models plan ahead for future tokens?

Paper · 2024 · Wilson Wu, John X. Morris and Lionel Levine

Summary

A transformer's hidden state at one position often contains information that later positions use. It is unclear whether models prepare this information on purpose. The authors compare two explanations. Under pre-caching, gradients from later losses push the model to compute features that do nothing for the current token but help afterwards. Under breadcrumbs, the features that best serve the current next-token prediction simply turn out to help future steps as well.

To separate the two, they define myopic training, which stops gradients from a position's loss from reaching the hidden states of earlier positions. The loss difference between vanilla and myopic models, called the myopia gap, measures pre-caching. On a synthetic task that can only be solved by computing ahead, transformers clearly learn to pre-cache. In natural language, GPT-2 trained from scratch on MS MARCO shows only a small gap, which supports breadcrumbs. Pythia models fine-tuned on the Pile show a gap that widens with size. The study clarifies how much LMs "plan ahead." It also suggests that myopic models might be safer or easier to interpret.

Key points

  • Theory: With weights tied across positions, the training gradient splits into an upper-triangular sum. The off-diagonal terms, which link a later position's loss to an earlier position's computation, are the signal that encourages pre-caching. Myopic training removes these terms.

  • GPT-2 setup: Models were trained from random initialization on 4.6M MS MARCO sequences of length 64.

  • GPT-2 results (validation cross-entropy):

    • Vanilla: 3.28
    • Myopic: 3.40, a gap of 0.12
    • "Local myopic": 3.26, a local bonus of only 0.02
    • Transformer bigram baseline: 5.33

    Most of the benefit from past context therefore comes from breadcrumbs.

  • Per-position analysis on PG-19: The myopic GPT-2 does better early in the sequence but falls behind as context grows. Aggregate losses are 4.67 (vanilla) and 4.77 (myopic).

  • Pythia scaling: Models from 14M to 2.8B parameters were fine-tuned from their pre-trained checkpoints on 10M 64-token Pile sequences. They were evaluated on LAMBADA, PIQA, SciQ and ARC-Easy. The myopia gap grows with scale and becomes non-negligible at Pythia 2.8B.

  • Synthetic task and code: In the synthetic setting, linear probes recover precomputed sin(b·x) features from the hidden states. Code is available in the GitHub repo wiwu2390/FutureGPT2-public.

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

🎓Cornell UniversityIthaca · United States139 works · 65,295 citations🎓University of Colorado BoulderBoulder · United States76 works · 51,396 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

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Wilson Wu, John X. Morris and Lionel Levine. 2024. Do language models plan ahead for future tokens?. arXiv:2404.00859.

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

@misc{wu2024do,
  title = {{Do language models plan ahead for future tokens?}},
  author = {Wilson Wu, John X. Morris and Lionel Levine},
  year = {2024},
  doi = {10.48550/arxiv.2404.00859},
  eprint = {2404.00859},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2404.00859}
}

References

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

  1. Maik Barthel, Sebastian Sauppe, Stephen C Levinson, and Antje S Meyer. The timing of utterance planning in task-oriented dialogue: Evidence from a novel list-completion paradigm. Frontiers in psychology, 7:1858, December 2016. doi: 10.3389/fpsyg.2016.01858. — [doi:10.3389/fpsyg.2016.01858](https://doi.org/10.3389/fpsyg.2016.01858) find
  2. David Bau, Jun-Yan Zhu, Hendrik Strobelt, Agata Lapedriza, Bolei Zhou, and Antonio Torralba. Understanding the role of individual units in a deep neural network. Proceedings of the National Academy of Sciences, 117(48):30071–30078, September 2020. ISSN 1091-6490. doi: 10.1073/pnas.1907375117. URL ht… — [doi:10.1073/pnas.1907375117](https://doi.org/10.1073/pnas.1907375117) find
  3. Amir Beck. First-Order Methods in Optimization. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2017. doi: 10.1137/1.9781611974997. URL https://epubs.siam.org/doi/abs/10.1137/1.9781611974997. — [doi:10.1137/1.9781611974997](https://doi.org/10.1137/1.9781611974997) ↗
  4. Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances, 2021.
  5. Yonatan Belinkov and James Glass. Analysis methods in neural language processing: A survey, 2019.
  6. Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens, 2023.
  7. Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der Wal. Pythia: a suite for analyzing large language models across training and… find
  8. Yonatan Bisk, Rowan Zellers, Ronan Le bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):7432–7439, Apr. 2020. doi: 10.1609/aaai.v34i05.6239. URL https://ojs.aaai.org/index.php/AAAI… ↗
  9. Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024.
  10. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. find
  11. Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. Linearity of relation decoding in transformer language models, 2024.
  12. John Hewitt and Percy Liang. Designing and interpreting probes with control tasks, 2019. find
  13. Falk Huettig. Four central questions about prediction in language processing. Brain Research, 1626:118–135, 2015. ISSN 0006-8993. doi: https://doi.org/10.1016/j.brainres.2015.02.014. URL https://www.sciencedirect.com/science/article/pii/S0006899315001146. Predictive and Attentive Processing in Perce… — [doi:10.1016/j.brainres.2015.02.014](https://doi.org/10.1016/j.brainres.2015.02.014) ↗
  14. 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, 2023. find
  15. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt, 2023. find
  16. George A. Miller. Language and communication. McGraw-Hill, New York, NY, US, 1951. doi: 10.1037/11135-000. — [doi:10.1037/11135-000](https://doi.org/10.1037/11135-000) find
  17. Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023.
  18. Yurii Nesterov. Lectures on Convex Optimization. Springer Publishing Company, Incorporated, 2nd edition, 2018. ISBN 3319915770. find
  19. Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. MS MARCO: A human generated machine reading comprehension dataset. In Tarek Richard Besold, Antoine Bordes, Artur S. d’Avila Garcez, and Greg Wayne (eds.), Proceedings of the Workshop on Cognitive Comput…
  20. CEUR-WS.org, 2016. URL https://ceur-ws.org/Vol-1773/CoCoNIPS 2016 paper9.pdf. nostalgebraist. Interpreting gpt: The logit lens, 2020. ↗
  21. Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 10.23915/distill.00024.001. https://distill.pub/2020/circuits/zoom-in. — [doi:10.23915/distill.00024.001](https://doi.org/10.23915/distill.00024.001) ↗
  22. Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: Anticipating subsequent tokens from a single hidden state. In Proceedings of the 27th Conference on Computational Natural Language Learning (CoNLL). Association for Computational Linguistics, 2023. doi: 10.18653/v1/2023…
  23. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Ann…
  24. Jacob Pfau, William Merrill, and Samuel R. Bowman. Let’s think dot by dot: Hidden computation in transformer language models, 2024. find
  25. Tiago Pimentel, Josef Valvoda, Rowan Hall Maudslay, Ran Zmigrod, Adina Williams, and Ryan Cotterell. Information-theoretic probing for linguistic structure, 2020.
  26. Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. URL https://api.semanticscholar.org/CorpusID:160025533. ↗
  27. Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL https://arxiv.org/abs/1911.05507. ↗
  28. Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Positional description matters for transformers arithmetic, 2023.
  29. Zuowei Shen, Haizhao Yang, and Shijun Zhang. Optimal approximation rate of relu networks in terms of width and depth. Journal de Mathématiques Pures et Appliquées, 157: 101–135, 2022. ISSN 0021-7824. doi: https://doi.org/10.1016/j.matpur.2021.07.009. URL https://www.sciencedirect.com/science/artic… — [doi:10.1016/j.matpur.2021.07.009](https://doi.org/10.1016/j.matpur.2021.07.009) ↗
  30. Xing Shi, Inkit Padhi, and Kevin Knight. Does string-based neural mt learn source syntax? pp. 1526–1534, 01 2016. doi: 10.18653/v1/D16-1159.
  31. Mitchell Stern, Noam M. Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. In Neural Information Processing Systems, 2018. URL https://api.semanticscholar.org/CorpusID:53208380. ↗
  32. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper files/pape… ↗
  33. Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. ArXiv, abs/1707.06209, 2017. URL https://api.semanticscholar.org/ CorpusID:1553193. ↗
  34. Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints, 2020. find
  35. Yuyu Zhang, Hanjun Dai, Kamil Toraman, and Le Song. kg 2 : Learning to reason science exam questions with contextual knowledge graph embeddings. ArXiv, abs/1805.12393, 2018. URL https://api.semanticscholar.org/CorpusID:44098100. — [arXiv:1805.12393](https://arxiv.org/abs/1805.12393) ↗
  36. Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wu-2024-language-models-plan-ahead-future/