Understanding Emergent In-Context Learning from a Kernel Regression Perspective
Paper · 2023 · Chi Han, Ziqi Wang, Han Zhao et al.
Summary
Large language models can solve new tasks from a few demonstrations in the prompt (in-context learning, ICL) without any weight updates, but the mechanism behind this ability remains unclear. Han et al. adopt the framework of Xie et al. (2022), which models the pre-training corpus as a mixture of HMMs with one per task. Within that framework, they prove that as the number of demonstrations grows, Bayesian inference over an ICL prompt approaches kernel regression: the prediction becomes a similarity-weighted average of the demonstration labels.
They test this claim empirically on GPT-J-6B by probing attention maps and the key and value vectors in each head. The model's attention behaves as the theory predicts. A single selected head can reproduce the model's ICL outputs fairly well, and kernel regression built on it outperforms kernel regression over standard sentence encoders.
The theory also accounts for several known ICL phenomena:
- Retrieving demonstrations similar to the test input improves accuracy.
- ICL is sensitive to output formatting.
- Representative, in-distribution examples help.
Key points
- Theory: Under the HMM-mixture assumptions, Bayesian ICL prediction asymptotically matches the kernel regression form ŷ = Σ yᵢK(x, xᵢ) / Σ K(x, xᵢ).
- Attention evidence: In GPT-J-6B, attention from the last token focuses mostly on label positions of demonstrations. Ridge regression probes (λ = 0.01) show that value vectors encode labels yᵢ and key vectors encode the model's prediction P(o|xᵢ).
- Reconstruction setup: The best head is chosen on 700 validation examples per task and evaluated on 300 held-out examples per task.
- Reconstruction results: Agreement with the model's ICL outputs is 68–80% on sst2, rottentomatoes and tweet_eval (hate/irony/offensive), dropping to about 60% on MNLI. Task accuracy of the reconstructed outputs is close to GPT-J-6B's own ICL accuracy and similar to or better than kernel regression over all-MiniLM-L6-v2 and bert-base-nli-mean-tokens.
- Code: Available at github.com/Glaciohound/Explain-ICL-As-Kernel-Regression.
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works
🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works
Cite
Cite as Chi Han, Ziqi Wang, Han Zhao and Heng Ji. 2023. Understanding Emergent In-Context Learning from a Kernel Regression Perspective. arXiv:2305.12766.
@misc{han2023understanding,
title = {{Understanding Emergent In-Context Learning from a Kernel Regression Perspective}},
author = {Chi Han, Ziqi Wang, Han Zhao and Heng Ji},
year = {2023},
doi = {10.48550/arxiv.2305.12766},
eprint = {2305.12766},
archivePrefix = {arXiv},
url = {https://doi.org/10.48550/arxiv.2305.12766}
}References
- Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36: 45614–45650, 2023. find
- 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. ↗
- Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Re. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Representations. Yu Bai, Fan Chen, Huan Wang, Caiming Xiong,…
- Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36:1560–1588, 2023. find
- 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.
- Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878–18891, 2022. find
- Wenlong Chen and Yingzhen Li. Calibrating transformers via sparse gaussian processes. arXiv preprint arXiv:2303.02444, 2023. — [arXiv:2303.02444](https://arxiv.org/abs/2303.02444) find
- 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. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2022. find
- Fabian Falck, Ziyu Wang, and Christopher C Holmes. Is in-context learning in large language models bayesian? a martingale perspective. In International Conference on Machine Learning, pp. 12784–12805. PMLR, 2024. find
- Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn to achieve second-order convergence rates for in-context linear regression. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. find
- Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Advances in Neural Information Processing Systems, 2023. find
- Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
- Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. In The Twelfth International Conference on Learning Representations. Yu Huang, Yuan Cheng, and Yingbin Lia… find
- Daniel Khashabi, Chitta Baral, Yejin Choi, and Hannaneh Hajishirzi. Reframing instructional prompts to gptk’s language. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 589–612, 2022. — Reframing Instructional Prompts to GPTk's Language
- Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. In Forty-first International Conference on Machine Learning. Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models… — Large language models are zero-shot reasoners
- Jannik Kossen, Tom Rainforth, and Yarin Gal. In-context learning in large language models learns label relationships but is not conventional learning. arXiv preprint arXiv:2307.12375, 2023. — [arXiv:2307.12375](https://arxiv.org/abs/2307.12375) find
- Sha Li, Ruining Zhao, Manling Li, Heng Ji, Chris Callison-Burch, and Jiawei Han. Open-domain hierarchical event schema induction by incremental prompting and verification. In Proc. The 61st Annual Meeting of the Association for Computational Linguistics (ACL2023), 2023. find
- Jiachang Liu, Dinghan Shen, Yizhe Zhang, William B Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pp. 100–114, 2…
- Robert Logan IV, Ivana Balažević, Eric Wallace, Fabio Petroni, Sameer Singh, and Sebastian Riedel. Cutting down on prompts and parameters: Simple few-shot learning with language models. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 2824–2835, 2022.
- Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 808…
- Yue M Lu, Mary I Letey, Jacob A Zavatone-Veth, Anindita Maiti, and Cengiz Pehlevan. Asymptotic theory of in-context learning by linear attention. arXiv preprint arXiv:2405.11751, 2024. — [arXiv:2405.11751](https://arxiv.org/abs/2405.11751) find
- Arvind V Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations. Sewon Min, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoy… find
- 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? arXiv preprint arXiv:2202.12837, 2022b.
- 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.
- Core Francisco Park, Ekdeep Singh Lubana, Itamar Pres, and Hidenori Tanaka. Competition dynamics shape algorithmic phases of in-context learning. arXiv preprint arXiv:2412.01003, 2024. — [arXiv:2412.01003](https://arxiv.org/abs/2412.01003) find
- Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. 2023. find
- Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Thomas Adler, David Kreil, Michael K Kopp, et al. Hopfield networks is all you need. In International Conference on Learning Representations. Allan Raventós, Mansheej Paul, Feng Chen,… find
- Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2019. URL http://arxiv.org/abs/1908.10084. ↗
- Ruifeng Ren and Yong Liu. Towards understanding how transformers learn in-context through a representation learning lens. arXiv preprint arXiv:2310.13220, 2023. — [arXiv:2310.13220](https://arxiv.org/abs/2310.13220) find
- Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2655–2671, 2022.
- Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
- 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.
- Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677, 2022.
- Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax, May 2021. ↗
- Xinyi Wang, Wanrong Zhu, and William Yang Wang. Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning. arXiv preprint arXiv:2301.11916, 2023.
- Zhenhailong Wang and Heng Ji. Open vocabulary electroencephalography-to-text decoding and zero-shot sentiment classification. In Proc. Thirty-Sixth AAAI Conference on Artificial Intelligence (AAAI2022), 2022. find
- Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022a.
- Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed H Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022b.
- Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03846, 2023.
- Yuhuai Wu, Felix Li, and Percy S Liang. Insights into pre-training via simpler synthetic tasks. Advances in Neural Information Processing Systems, 35:21844–21857, 2022. find
- Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2022.
- Wanying Xie. Gx at semeval-2021 task 2: Bert with lemma information for mcl-wic task. In Proceedings of the 15th International Workshop on Semantic Evaluation (SemEval-2021), pp. 706–712, 2021. find
- Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1–55, 2024. find
- Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning, pp. 12697–12706. PMLR, 2021.
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/han-2023-understanding-emergent-context-learning-kernel/