Trained Transformers Learn Linear Models In-Context

By Spencer Frei, Peter L. Bartlett. Journal of machine learning research, 2023

366 citations recorded by the source corpus.

Attention-based neural networks such as transformers have demonstrated a remarkable ability to exhibit in-context learning (ICL): Given a short prompt sequence of tokens from an unseen task, they can formulate relevant per-token and next-token predictions without any parameter updates. By embedding a sequence of labeled training data and unlabeled test data as a prompt, this allows for transformers to behave like supervised learning algorithms. Indeed, recent work has shown that when training transformer architectures over random instances of linear regression problems, these models' predictions mimic those of ordinary least squares. Towards understanding the mechanisms underlying this phenomenon, we investigate the dynamics of ICL in transformers with a single linear self-attention layer trained by gradient flow on linear regression tasks. We show that despite non-convexity, gradient flow with a suitable random initialization finds a global minimum of the objective function. At this global minimum, when given a test prompt of labeled examples from a new prediction task, the transformer achieves prediction error competitive with the best linear predictor over the test prompt distribution. We additionally characterize the robustness of the trained transformer to a variety of distribution shifts and show that although a number of shifts are tolerated, shifts in the covariate distribution of the prompts are not. Motivated by this, we consider a generalized ICL setting where the covariate distributions can vary across prompts. We show that although gradient flow succeeds at finding a global minimum in this setting, the trained transformer is still brittle under mild covariate shifts. We complement this finding with experiments on large, nonlinear transformer architectures which we show are more robust under covariate shifts.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

View Paper: https://www.semanticscholar.org/paper/4a7530bbaee7563ee244f3ffed6b706bd96f08a8

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

1. Hadi Daneshmand
2. Xiang Cheng
3. Suvrit Sra
4. Kwangjun Ahn
5. Sashank J. Reddi
6. Alex Smola
7. Barnab'as P'oczos

In Computer Vision and Pattern Recognition

1. Hadi Daneshmand
2. Suvrit Sra
3. Kwangjun Ahn

Citations in the corpus

• Transformers learn to implement preconditioned gradient descent for in-context learning (2023, background): “We discuss here two related works (Zhang et al., 2023; Mahankali et al., 2023) that appeared shortly after publication of our original draft.”