What Can Transformers Learn In-Context? A Case Study of Simple Function Classes

By Shivam Garg, Dimitris Tsipras, Percy Liang, Gregory Valiant. Neural Information Processing Systems, 2022

825 citations recorded by the source corpus.

In-context learning refers to the ability of a model to condition on a prompt sequence consisting of in-context examples (input-output pairs corresponding to some task) along with a new query input, and generate the corresponding output. Crucially, in-context learning happens only at inference time without any parameter updates to the model. While large language models such as GPT-3 exhibit some ability to perform in-context learning, it is unclear what the relationship is between tasks on which this succeeds and what is present in the training data. To make progress towards understanding in-context learning, we consider the well-defined problem of training a model to in-context learn a function class (e.g., linear functions): that is, given data derived from some functions in the class, can we train a model to in-context learn"most"functions from this class? We show empirically that standard Transformers can be trained from scratch to perform in-context learning of linear functions -- that is, the trained model is able to learn unseen linear functions from in-context examples with performance comparable to the optimal least squares estimator. In fact, in-context learning is possible even under two forms of distribution shift: (i) between the training data of the model and inference-time prompts, and (ii) between the in-context examples and the query input during inference. We also show that we can train Transformers to in-context learn more complex function classes -- namely sparse linear functions, two-layer neural networks, and decision trees -- with performance that matches or exceeds task-specific learning algorithms. Our code and models are available at https://github.com/dtsip/in-context-learning .

Subfields: Artificial Intelligence

View Paper: https://www.semanticscholar.org/paper/de32da8f5c6a50a6c311e9357ba16aa7d05a1bc9

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Looped Transformers as Programmable Computers (2023, methodology): “Previous research, such as in (Akyürek et al., 2022) and (Garg et al., 2022), has limited in-context learning to a single inference call of a deeper transformer model than ours, which restricts the types of models that can be learned and the level of accuracy that can be achieved.”
CodeMind: A Framework to Challenge Large Language Models for Code Reasoning (2024, background): “Large Language Models (LLMs) have shown exceptional programming abilities, specifically when instruction-tuned or prompted through Chain-or Tree-of-Thoughts (CoT (Wei et al., 2022b) or ToT (Yao et al., 2023)) and in-context learning (Wei et al., 2022a; Garg et al., 2022).”
Otter: A Multi-Modal Model With In-Context Instruction Tuning (2023, neutral): “3.3.”
What learning algorithm is in-context learning? Investigations with linear models (2022, methodology): “While past work has characterized what kinds of functions ICL can learn (Garg et al., 2022; Laskin et al., 2022) and the distributional properties of pretraining that can elicit in-context learning (Xie et al.”
Larger language models do in-context learning differently (2023, methodology): “…from scratch can perform in-context learning on linear-regression datasets with performance that is comparable to the least-squares estimator (Garg et al., 2022), and recent work has shown that transformers can do so by implementing standard learning algorithms such as ridge regression and…”
Towards Understanding Chain-of-Thought Prompting: An Empirical Study of What Matters (2022, background): “We note that what we find does not in any way diminish the potential of learning to reason in context for LLMs; recent work has also shown evidence that learning in context is possible and could be powerful (Garg et al., 2022; Akyürek et al., 2023).”
Large Language Models (2023, background): “A particularly interesting task is linear regression [48]; since this is the prototypical case of statistical inference, a system which learns to do it can be said to be “learning how to learn.”
A Survey for In-context Learning (2023, background): “Learning Mechanism By learning linear functions, Garg et al. (2022) proved that Transformers could encode effective learning algorithms to learn unseen linear functions according to the demonstration samples.”
Hyena Hierarchy: Towards Larger Convolutional Language Models (2023, background): “, 2021), in-context learning (ICL) (Garg et al., 2022) and language model benchmarking (Liang et al.”
Towards Revealing the Mystery behind Chain of Thought: a Theoretical Perspective (2023, methodology): “This regression setting is convenient for manipulating numerical values and has been extensively adopted in prior works [21, 1].”
Transformers learn to implement preconditioned gradient descent for in-context learning (2023, methodology): “This property has become the focus of body of recent research that aims to shed light on the underlying mechanism of large language models (Garg et al., 2022; Akyürek et al., 2022; von Oswald et al., 2023; Li and Malik, 2017; Min et al., 2021; Xie et al., 2021; Elhage et al., 2021; Olsson et al.,…”
Prompting PaLM for Translation: Assessing Strategies and Performance (2022, methodology): “Hard prompts have also been used to analyze model capabilities (Garg et al., 2022; Li et al., 2022a), the role of data (Singh et al.”
Transformers learn in-context by gradient descent (2022, methodology): “Inspired by (Garg et al., 2022), we additionally provide results when training a single linear self-attention layer on a fixed number of training tasks.”