Abstract: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.

Polynomial Regression as a Task for Understanding In-context Learning Through Finetuning and Alignment

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

Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions

On the Role of Depth and Looping for In-Context Learning with Task Diversity

Learning Dynamics of LLM Finetuning

Explaining Emergent In-Context Learning as Kernel Regression

In-Context Learning with Representations: Contextual Generalization of Trained Transformers

MLPs Learn In-Context on Regression and Classification Tasks

Context-Parametric Inversion: Why Instruction Finetuning May Not Actually Improve Context Reliance

Understanding In-Context Learning with a Pelican Soup Framework

Enhancing In-context Learning via Linear Probe Calibration

Toward Understanding In-context vs. In-weight Learning

Does learning the right latent variables necessarily improve in-context learning?

Exploring Universal Intrinsic Task Subspace Via Prompt Tuning

Asymptotic theory of in-context learning by linear attention

Exploring Universal Intrinsic Task Subspace for Few-shot Learning Via Prompt Tuning

In-context Learning in Presence of Spurious Correlations

In-Context Learning Functions with Varying Number of Minima

How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?

Trained Transformers Learn Linear Models In-Context

Re-examining learning linear functions in context