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.

Bypassing the Exponential Dependency: Looped Transformers Efficiently Learn In-context by Multi-step Gradient Descent

Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?

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

Looped Transformers are Better at Learning Learning Algorithms

Linear Transformers are Versatile In-Context Learners

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

Provably Transformers Harness Multi-Concept Word Semantics for Efficient In-Context Learning

Linear Transformers with Learnable Kernel Functions are Better In-Context Models

Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression

On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery

Asymptotic theory of in-context learning by linear attention

Transformers for Supervised Online Continual Learning

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

In-Context Convergence of Transformers

Transformers are Deep Optimizers: Provable In-Context Learning for Deep Model Training

Looped Transformers for Length Generalization

Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape

How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression

Trained Transformers Learn Linear Models In-Context

One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention

Insights into LLM Long-Context Failures: When Transformers Know but Don't Tell