Elamur Research⌕
☰
home / papers / 2023 / Transformers learn to implement preconditioned gradient descent for in-context learning

Transformers learn to implement preconditioned gradient descent for in-context learning

Paper · 2023 · Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand et al.

Summary

The paper asks whether transformers merely can represent optimization algorithms or whether ordinary training actually drives them to do so. Earlier work built weights by hand to show this. Here the authors study the loss landscape of linear (softmax-free) self-attention networks trained on randomly sampled linear regression prompts, where each prompt holds labeled examples plus an unlabeled query.

For one attention layer, they fully characterize the global minimizer. It performs a single step of preconditioned gradient descent, and the preconditioner reflects both the covariance of the inputs and the extra variance from having only finitely many in-context samples. For deeper models, under a sparsity restriction on the weights, they show that certain critical points correspond to L steps of preconditioned or adaptive-stepsize gradient descent. Without that restriction, a critical point yields a richer method that also transforms the data to improve conditioning; in the isotropic case it reduces to GD++.

Experiments with a three-layer model match the predicted weight structure and reach near-zero loss. This gives early theoretical support for learned in-context algorithms.

Key points

  • Setup: Covariates are Gaussian N(0, Σ) and regression weights are Gaussian. Training minimizes the in-context squared error on the query's label with attention of the form P Z M (ZᵀQZ).
  • Theorems 1–2 (global minimizers):
    • Theorem 1: a single layer yields one preconditioned GD step.
    • Theorem 2: a two-layer model on isotropic data with symmetric sparse weights yields GD with adaptive stepsizes.
  • Theorems 3–4 (critical points):
    • Theorem 3: multi-layer critical points implement L steps of GD preconditioned roughly by Σ⁻¹.
    • Theorem 4: without the sparsity constraint, a critical point recovers GD++ of von Oswald et al. (2023) when the data are isotropic.
    • Lemma 4 shows that multiple critical points can exist.
  • Theorem 5 (ReLU attention): For isotropic data, a single ReLU attention layer has a global minimizer whose structure resembles the linear case.
  • Empirical check: In a three-layer model, learned layer norms grow with depth, meaning small early stepsizes and large later ones. Code is available at github.com/chengxiang/LinearTransformer. The paper appeared at NeurIPS 2023.

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

 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 Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand and Suvrit Sra. 2023. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv:2306.00297 · 37th Conference on Neural Information Processing Systems (NeurIPS 2023).

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

@misc{ahn2023transformers,
  title = {{Transformers learn to implement preconditioned gradient descent for in-context learning}},
  author = {Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand and Suvrit Sra},
  year = {2023},
  doi = {10.48550/arxiv.2306.00297},
  eprint = {2306.00297},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2306.00297}
}

References

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

  1. Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082, 2023.
  2. Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. International Conference on Learning Representations, 2022.
  3. Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673, 2023.
  4. Noga Alon and Joel H Spencer. The probabilistic method. John Wiley & Sons, 2016. find
  5. Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, et al. Gpt-neox-20b: An open-source autoregressive language model. Proceedings of BigScience – Workshop on Challenges & Perspectives in Creating Large Language…
  6. 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. Neural Information Processing Systems, 2020.
  7. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volum…
  8. John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011. find
  9. Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning (ICML), 2022. find
  10. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Ja… find
  11. Murat A Erdogdu, Lee H Dicker, and Mohsen Bayati. Scaled least squares estimator for glms in large-scale problems. Advances in Neural Information Processing Systems, 29, 2016. find
  12. 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.
  13. Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv preprint arXiv:2301.13196, 2023. find
  14. Alex Graves, Greg Wayne, and Ivo Danihelka. arXiv:1410.5401, 2014. find
  15. Neural turing machines. arXiv preprint Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 1997. find
  16. Stanisław Jastrzebski, Devansh Arpit, Nicolas Ballas, Vikas Verma, Tong Che, and Yoshua Bengio. Residual connections encourage iterative inference. In International Conference on Learning Representations, 2018. URL https://openreview.net/forum?id=SJa9iHgAZ. ↗
  17. Ke Li and Jitendra Malik. Learning to optimize. In International Conference on Learning Representations, 2017. find
  18. Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evaluation. White Paper. AI21 Labs, 2021. find
  19. Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023.
  20. Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2021. find
  21. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amode… — In-context Learning and Induction Heads
  22. Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing complete. The Journal of Machine Learning Research, 2021. find
  23. Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446, 2021. find
  24. Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355–9366. PMLR, 2021.
  25. Hava T Siegelmann and Eduardo D Sontag. On the computational power of neural nets. In Proceedings of Workshop on Computational learning theory, 1992. find
  26. Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 1999. find
  27. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 2017.
  28. Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151–35174. PMLR, 2023.
  29. Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35:12071–12083, 2022. find
  30. Mitchell Wortsman, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Replacing softmax with relu in vision transformers. arXiv preprint arXiv:2309.08586, 2023. — [arXiv:2309.08586](https://arxiv.org/abs/2309.08586) find
  31. Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. International Conference on Learning Representations, 2021.
  32. Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023. — [arXiv:2306.09927](https://arxiv.org/abs/2306.09927) find
  33. Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08117, 2023. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ahn-2023-transformers-learn-implement-preconditioned-gradient/