Elamur Research⌕
☰
home / papers / 2023 / One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention

One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention

Paper · 2023 · Arvind Mahankali, Tatsunori B. Hashimoto and Tengyu Ma

Summary

Earlier experiments by von Oswald et al. [2022] found that a transformer with a single linear self-attention layer and no MLP, trained on synthetic regression prompts, behaves like a single gradient descent update on least squares. This paper supplies the missing proof. The authors study a one-layer, single-head linear self-attention model of width d+1, where each token stacks an input x_i with its label y_i. They then characterize the parameters that globally minimize the expected squared error on noisy linear regression prompts. With standard Gaussian inputs and weight vectors, the optimal model predicts by scaling the sum of y_i x_iᵀx, which is exactly one GD step from zero. With non-isotropic Gaussian inputs, the optimum becomes a preconditioned GD step. Replacing the linear target with a broad family of nonlinear functions that is symmetric under rotations and negation still yields plain one-step GD, with only the step size changing. These results give in-context learning experiments a rigorous footing. They also suggest that the input distribution shapes the learned algorithm far more than the label distribution does.

Key points

  • Setup: Prompts take the form (x_1, y_1, …, x_n, y_n, x_{n+1}), with x_i ~ N(0, I), a fresh w ~ N(0, I) for each sequence, and y_i = wᵀx_i plus Gaussian noise of variance σ². The model never attends a token to itself.
  • Isotropic case: The global optimum outputs η·Σ y_i x_iᵀx, which matches the empirical observation of von Oswald et al. [2022].
  • Non-isotropic case: For x_i ~ N(0, Σ) and w ~ N(0, Σ⁻¹), the optimum implements one step of preconditioned GD.
  • Nonlinear targets (Theorem 3): The result also covers target functions such as fully connected networks of any depth and width with i.i.d. N(0,1) weights in the first and last layers. The optimum remains least-squares one-step GD.
  • Related work and open questions: Overlapping results were proven concurrently by Zhang et al. [2023] and, separately, by Ahn et al. [2023], so the nonlinear-target analysis is what sets this paper apart. The authors leave multi-head layers and multi-layer models with nonlinear targets as open directions.

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

🎓Stanford UniversityStanford · United States567 works · 335,190 citations

 Fields, topics and concepts

📚 Shelves 2The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 works

Cite

Cite as Arvind Mahankali, Tatsunori B. Hashimoto and Tengyu Ma. 2023. One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention. arXiv:2307.03576.

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

@misc{mahankali2023one,
  title = {{One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention}},
  author = {Arvind Mahankali, Tatsunori B. Hashimoto and Tengyu Ma},
  year = {2023},
  doi = {10.48550/arxiv.2307.03576},
  eprint = {2307.03576},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2307.03576}
}

References

17 references, 8 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, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 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. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=0g0X4H8yN4I. ↗
  3. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter,… — Language Models are Few-Shot Learners
  4. Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers, 2023. find
  5. Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? A case study of simple function classes. In NeurIPS, 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/c529dba08a146ea8d6cf715ae8930cbe-Abstract-Conference ↗
  6. Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers, 2023. find
  7. Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evaluation. Technical report, AI21 Labs, 2021. find
  8. Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. find
  9. OpenReview.net, 2023. URL https://openreview.net/pdf?id=De4FYqjFueZ. ↗
  10. Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In EMNLP, 2022.
  11. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. Technical report, OpenAI, 2019. find
  12. Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, 2022.
  13. Ben Wang and Aran Komatsuzaki. Gpt-j-6b: A 6 billion parameter autoregressive language model, 2021. find
  14. Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, and Tengyu Ma. Larger language models do in-context learning differently, 2023.
  15. Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022.
  16. OpenReview.net, 2022. URL https://openreview.net/forum?id=RdJVFCHjUMI. ↗
  17. Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context, 2023. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/mahankali-2023-one-step-gradient-descent-provably/