Transformers learn in-context by gradient descent

By Johannes von Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, Max Vladymyrov. International Conference on Machine Learning, 2022

821 citations recorded by the source corpus.

At present, the mechanisms of in-context learning in Transformers are not well understood and remain mostly an intuition. In this paper, we suggest that training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We start by providing a simple weight construction that shows the equivalence of data transformations induced by 1) a single linear self-attention layer and by 2) gradient-descent (GD) on a regression loss. Motivated by that construction, we show empirically that when training self-attention-only Transformers on simple regression tasks either the models learned by GD and Transformers show great similarity or, remarkably, the weights found by optimization match the construction. Thus we show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass. This allows us, at least in the domain of regression problems, to mechanistically understand the inner workings of in-context learning in optimized Transformers. Building on this insight, we furthermore identify how Transformers surpass the performance of plain gradient descent by learning an iterative curvature correction and learn linear models on deep data representations to solve non-linear regression tasks. Finally, we discuss intriguing parallels to a mechanism identified to be crucial for in-context learning termed induction-head (Olsson et al., 2022) and show how it could be understood as a specific case of in-context learning by gradient descent learning within Transformers. Code to reproduce the experiments can be found at https://github.com/google-research/self-organising-systems/tree/master/transformers_learn_icl_by_gd .

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

View Paper: https://www.semanticscholar.org/paper/525d93a382f6e7873b5d8a2e0713eb3dff7fb250

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

In Computer Vision and Pattern Recognition

Citations in the corpus

Towards Revealing the Mystery behind Chain of Thought: a Theoretical Perspective (2023, background): “Subsequent works further revealed that Transformers can implement learning algorithms such as linear regression [1], gradient descent [1, 54, 16], and Bayesian inference [59].”
A Survey for In-context Learning (2023, methodology): “Also based on simple regression tasks, von Oswald et al. (2022) showed that linear attention-only Transformers with hand-constructed parameters and models learned by gradient descent are highly related.”
Larger language models do in-context learning differently (2023, neutral): “…that models primarily rely on semantic prior knowledge), other work has shown that transformers in simple settings (without language modeling pretraining) implement learning algorithms such as ridge regression and gradient descent (Akyürek et al., 2022; von Oswald et al., 2022; Dai et al., 2022).”
Jailbreak and Guard Aligned Language Models With Only Few In-Context Demonstrations (2023, neutral): “…y 1 ) , . . . , ( x k , y k ) | f ( x i ) = y i } where x i are query inputs and y i are their corresponding labels in this task, a model can learn this mapping and successfully predict the label y new = f ( x new ) of a new input query x new by prompting [ x 1 , y 1 , . . . , x k , y k , x new ]…”
Looped Transformers as Programmable Computers (2023, methodology): “In a similar vein, (von Oswald et al., 2022) argues that training Transformers on auto-regressive tasks is closely related to gradient-based meta-learning formulations.”
Transformers learn to implement preconditioned gradient descent for in-context learning (2023, methodology): “Empirically, von Oswald et al. (2023) have observed that for linear regression task, softmax activations generally degrade the prediction performance, and in particular, softmax transformers typically need more attention heads to match their performance with that of linear transformers.”
A Survey of Large Language Models (2023, neutral): “They typically analyze from the perspective of gradient descent and consider ICL as implicit fine-tuning [60, 238].”
Language Models can Solve Computer Tasks (2023, background): “Recently, large language models (LLMs) have shown remarkable in-context learning capabilities across a variety of domains and tasks [10, 53, 4, 14, 20, 49, 7, 35, 5].”
Otter: A Multi-Modal Model With In-Context Instruction Tuning (2023, neutral): “3.3.”