Trained Transformers Learn Linear Models In-Context

Ruiqi Zhang; Spencer Frei; Peter L. Bartlett

Source-linked AI summary

Trained Transformers Learn Linear Models In-Context

Ruiqi Zhang, Spencer Frei, Peter L. Bartlett

arXiv:2306.09927v3Jun 16, 2023stat.MLcs.AIcs.CLcs.LG

TL;DR

The paper asks how transformers acquire in-context linear prediction through gradient-based training and how robust that learned behavior is to distribution shifts. It analyzes single-layer linear self-attention under gradient flow, proving global convergence and characterizing prediction behavior. The models remain brittle to covariate shifts, whereas larger nonlinear transformers are more robust in experiments.

Problem
The paper studies how gradient-based optimization produces transformers capable of in-context learning and how that behavior changes under distribution shifts.
Method
The authors analyze gradient flow on the population loss for single-layer linear self-attention transformers trained on random linear-regression prompts, then examine generalized prompts and larger nonlinear transformers.
Results
Gradient flow converges to a global minimum despite non-convexity; linear self-attention models tolerate some shifts but fail under covariate shift, while larger nonlinear transformers generalize better there.
Takeaways & Limitations
The learned algorithm can in-context learn linear models, but robustness to covariate shifts is limited for the analyzed linear self-attention models.
Takeaways & Limitations
The analysis is scoped to simplified single-layer linear self-attention trained on random noiseless linear models over anisotropic Gaussian marginals, with brittleness under covariate shifts.

Abstract

from arXiv · show

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.

1 Introduction

This work studies how gradient flow trains simplified transformers to perform in-context learning of linear models, and examines their behavior under distribution shifts. It finds global convergence for linear self-attention models, while covariate shifts remain challenging for these models and are handled better by larger nonlinear transformers.

Research focus: The paper investigates gradient-flow learning dynamics in a simplified transformer trained on random linear-regression prompts.The architecture has a single layer with linear self-attention.
Optimization: Gradient flow from a suitable random initialization converges to a global minimum despite the population objective being non-convex.This guarantee applies to single-layer linear self-attention transformers.
In-context prediction: At convergence, the transformer’s encoded learning algorithm and prediction error can be characterized on new prompts, including prompts from possibly nonlinear tasks.The contribution concerns both the learned algorithm and its test-prompt prediction error.
Distribution shifts: Several distribution shifts are tolerated, but shifts in the covariate distribution of prompt features xi are not.The covariate-shift limitation is a central robustness finding.
Distribution shifts: Even when covariate distributions vary across training prompts, linear self-attention transformers still fail under covariate shift.The generalized setting retains global convergence guarantees but not robustness to this shift.
Distribution shifts: Large nonlinear transformers generalize better under covariate shift, especially when trained on prompts with varying covariate distributions.This finding is empirical and contrasts with the brittleness of the analyzed linear self-attention models.

2 Additional Related Work

Related work has established empirical and constructive accounts of in-context learning, while leaving open how gradient-based optimization produces such behavior. This paper connects those questions to optimization dynamics and studies broader theoretical perspectives on transformer learning.

In-context learning: Prior empirical work found that transformers trained on random linear-regression instances produce predictions similar to ordinary least squares.The same studies also reported in-context learning of two-layer ReLU networks and decision trees.
In-context learning: Akyürek et al. and Oswald et al. showed that trained transformers on isotropic Gaussian linear-regression data mimic one gradient-descent step.They also constructed transformers implementing that step.
Interpretations of ICL: The paper relates ICL to Bayesian inference, implicit fine-tuning, kernel regression, Bayesian selection, and gradient descent on softmax regression.These works provide different interpretations of sequence-based learning behavior.
Approximation theory: Approximation-theoretic work established universal sequence-to-sequence approximation and studied sparse functions learned by single-layer self-attention.These results provide a complementary perspective on transformer expressivity.
Anisotropic covariates: For anisotropic Gaussian covariates, a single vanilla gradient-descent step from zero initialization does not generally achieve small prediction error.The limitation follows because the resulting predictor incorporates the covariance matrix rather than directly recovering the target linear predictor.
Optimization theory: Other theory studies gradient-based transformer training for spatial data, topic models, and optimization landscapes of linear self-attention networks.The cited works analyze vision transformers, co-occurrence-based topic learning, and critical points or global minima.

3 Preliminaries

The preliminaries define in-context learning as prediction from labeled prompts without parameter updates, formulate its training objective, and introduce the linear self-attention architecture analyzed in the paper.

3.1 In-context learning: In-context learning uses labeled examples and a query in a prompt to predict the query output for an unknown function.The target is a prediction by(xquery) close to h(xquery).
3.1 In-context learning: Training samples independent prompts by drawing functions and feature vectors, then minimizes the population objective with stochastic optimization.The resulting model predicts query examples from sequences of training data.
3.1 In-context learning: The framework permits additive error when the model lacks sufficient complexity or the setting is non-realizable.Thus arbitrarily small prediction error is not always expected.
3.1 In-context learning: The framework asks whether models trained on prompts from a hypothesis class can learn that class with small prediction error using standard gradient-based optimization.It also asks how training and test context lengths affect prediction error.
3.2 Linear self-attention networks: Standard self-attention forms an attention matrix from key-query interactions, applies column-wise softmax, and can process sequences of arbitrary length.The linear self-attention variant used here removes the softmax nonlinearity.
3.2 Linear self-attention networks: The analyzed linear self-attention model removes softmax, merges value-projection matrices, merges query-key matrices, and parameterizes both in θ.This simplification makes the architecture more amenable to theoretical analysis while retaining linear-model ICL capability.
3.2 Linear self-attention networks: The prompt embedding stacks (xi, yi) in the first N columns and (xquery, 0) in the final column.The query prediction is read from the bottom-right output entry.

4 Main results

The paper proves global convergence for a single linear self-attention layer trained by gradient flow and characterizes when it in-context learns linear predictors. It also shows that covariate shifts remain problematic for this model, while larger nonlinear transformers are more robust in experiments.

Convergence: Gradient flow from suitable random initialization converges to a global minimum despite the non-convexity of the linear self-attention objective.The result applies to population loss and a specified initialization class.
Prediction error: For sufficiently long training and test prompts, the trained transformer achieves prediction error competitive with the best linear model.The bound scales as O(1/M + 1/N^2), so finite training-prompt length leaves a residual gap even as M grows.
Prediction error: For Gaussian linear regression, the model in-context learns up to error η := (1 + 2d + d^2κ) tr(Λ)/N^2 when M ≥ (d + 1) tr(Λ)ε^-1.The resulting prediction error is at most η + ε, with κ the condition number of Λ.
Distribution shifts: Task shifts and broad query-distribution shifts can be tolerated when prompts are sufficiently long, but covariate shifts between training and test prompts cannot be fully tolerated.An orthogonal query example can yield zero prediction, and the covariate-shift failure follows from the learned solution's dependence on the training distribution.
Distribution shifts: Training across varying covariate distributions does not resolve the failure for single-layer linear self-attention, whereas large nonlinear transformers generalize better under covariate shift.In GPT2 experiments, fixed-covariance training succeeds when variance is not too large, while random-covariance training still fails to match least squares at c = 9.
Experiments: GPT2 prediction error spikes when test prompts exceed the training prompt length, regardless of whether covariance is fixed or random.The paper conjectures that randomly initialized positional encodings for unseen positions contribute to this spike.

5 Proof ideas

The proof reduces linear self-attention training to gradient flow on a non-convex quadratic, rank-one matrix-factorization problem. A Polyak–Łojasiewicz argument then establishes global convergence and identifies the limiting parameters that yield linear in-context predictions.

Quadratic reduction: The proof rewrites the model’s query prediction as a quadratic function of parameters determined by the prompt embedding matrix.This reduces the original transformer loss to an optimization problem involving a quadratic form.
Quadratic reduction: The resulting quadratic optimization is a rank-one matrix-factorization problem, but it is non-convex because its governing matrix has negative eigenvalues.The proof establishes that the matrix has at least d + 1 negative eigenvalues.
Gradient-flow dynamics: Gradient-flow dynamics reduce to coupled differential equations for the relevant parameter blocks, with certain off-diagonal blocks remaining zero throughout training.The dynamics are then related to gradient flow on an equivalent objective differing only by parameter-independent constants.
Global convergence: Global minima satisfy a parameter relation involving Γ^-1, which approaches Λ^-1 as the training prompt length N grows.This limiting relation connects the optimization variables to the inverse feature covariance governing prediction.
Prediction at convergence: At the converged solution, the transformer’s predictions are obtained by translating the limiting parameter blocks back into the original attention parameterization.This completes the proof by combining the reduction lemmas with the convergence result.
Global convergence: A Polyak–Łojasiewicz inequality implies that gradient flow converges globally from suitable initialization and permits exact calculation of the limiting parameter values.The convergence result is formalized through the dynamics of U11 and u^-1.

6 Conclusion and future work

The analysis shows that single-layer linear self-attention transformers can learn in-context linear prediction under gradient flow, while remaining brittle to covariate shifts. The authors identify broader optimization settings and nonlinear, deep transformers as directions for future work.

Main findings: Gradient flow from suitable random initialization converges to a global minimum despite the non-convex population objective.The result concerns single-layer transformers with linear self-attention trained on random noiseless linear models over anisotropic Gaussian marginals.
Robustness: The trained transformer is robust to shifts in task and query distributions but brittle when training and test covariate distributions differ.This matches empirical observations from prior work on trained transformer models.
Future work: The analysis is restricted to population-loss gradient flow with particular random initializations and does not establish analogous results for finite-step stochastic gradient descent.More general initialization schemes and stochastic gradient descent remain open questions.
Future work: The theoretical model uses a single linear self-attention layer, whereas the authors seek dynamics for nonlinear and deep transformers.The analyzed class is sufficient for in-context learning of linear predictors but does not cover the broader architectures of interest.
Generalized ICL: When covariate distributions vary across prompts, single-layer linear self-attention transformers fail on test prompts with covariate distributions sampled from the same distribution.Larger nonlinear transformers appear more successful in this setting but remain sub-optimal.

A Proof of Theorem 4.1

This proof establishes the algebraic form of the query prediction for a single linear self-attention transformer. It connects the prediction to selected components of the model’s parameter matrices through matrix decompositions and vectorization identities.

Proof structure: Theorem 4.1 follows as a corollary of four lemmas after translating u−1 and U11 back to W PV and W KQ.The proof section is organized around Lemmas 5.1–5.4.
Prediction form: For a prompt embedding matrix Eτ, the query prediction can be written as the output of a quadratic function.Lemma 5.1 identifies the prediction for the query covariate with the right-bottom entry of the linear self-attention output.
Prediction form: The embedding matrix Eτ represents a prompt of length N together with its weight wτ.The lemma treats Eτ as an embedding matrix in R^(d+1)×(N+1).
Algebraic derivation: The derivation uses decompositions of W PV and W KQ together with matrix vectorization, Kronecker products, and trace identities.These algebraic tools are used to express the attention output in the stated quadratic form.

A.2 Proof of Lemma 5.2

The proof derives gradient-flow dynamics for the relevant parameter blocks of the single-layer transformer. It shows that off-diagonal blocks remain zero and then computes the dynamics of U11 and u−1.

Gradient-flow setup: Gradient flow is analyzed by differentiating the loss with respect to the parameter blocks using the chain rule and matrix-derivative identities.The proof repeatedly exploits independence and Gaussianity of prompt covariates and task weights.
Derivative simplifications: Independence between prompt covariates and task weights eliminates terms in the derivative calculations.The proof invokes mean-zero Gaussian weights and independence from prompt and query inputs.
Invariant blocks: If u12 = u21 = 0d initially, their derivatives vanish, so these blocks remain zero for all t ≥ 0.The argument uses the independence of task weights from prompt and query inputs.
U11 dynamics: The proof derives entrywise dynamics for U11 after imposing u12 = u21 = 0d.For each k,l ∈ [d], the resulting derivatives determine the evolution of Ukl.
u−1 dynamics: The dynamics of u−1 are computed after establishing the invariant zero off-diagonal blocks.The resulting expression includes Γ and preserves u12(t) = u21(t) = 0d for all t ≥ 0.

A.3 Proof of Lemma 5.3

The proof reduces gradient flow to an equivalent loss over U11 and u−1, then characterizes its minima. A scaling symmetry explains why the loss has non-unique equivalent parameterizations.

Loss reduction: Gradient flow on the original loss is equivalent to gradient flow on a loss function involving only U11 and u−1.The equivalence holds for initial values satisfying Assumption 3.3.
Loss reduction: The reduced loss is recovered from the gradient-flow dynamics up to an additive constant.Its gradients with respect to U11 and u−1 reproduce the relevant dynamics.
Global minima: The reduced loss can be negative, but its global minima can still be characterized.The minimum analysis uses trace and Frobenius-norm identities.
Algebraic simplification: The proof uses that Γ and Λ commute when simplifying the reduced loss and its derivatives.This commutation property appears in the derivation of the loss expression and subsequent calculations.
Global minima: The loss is invariant under replacing (U11, u−1) with (cU11, c−1u−1) for any non-zero constant c.This scaling symmetry yields equivalent parameterizations with the same loss.

A.4 Proof of Lemma 5.4

The proof establishes that balanced initialization keeps the LSA parameters balanced throughout gradient flow and prevents the scalar parameter u−1 from reaching the origin. These properties support a PL inequality and global convergence to the minimum.

Balanced parameters: Balanced initialization preserves the parameter balance relation throughout the entire gradient-flow trajectory.The proof identifies this invariant as the reason for assuming balanced parameters initially.
Positivity: The loss is non-increasing under gradient flow, and u−1 remains positive because reaching u−1 = 0 would make the loss zero.Continuity together with positive initialization preserves positivity for all t ≥ 0.
Positivity: A positive lower bound on u−1 is established for all times, keeping the trajectory away from the saddle point at the origin.This lower bound is used as a critical ingredient in proving the PL inequality.
PL inequality: The proof selects a positive PL constant under the initialization-scale condition, yielding a PL inequality for the loss.The constant is positive because the assumed initialization scale satisfies the lemma’s condition.
Global convergence: Gradient flow converges to the global minimum of the loss, with U11 and u−1 converging to the stated limiting values.The convergence follows from the gradient-flow dynamics together with the PL condition.

B Proof of Theorem 4.2

The proof characterizes the prediction error of a trained linear self-attention layer relative to the best linear predictor, first for Gaussian covariates and then for random covariance matrices. It decomposes the error and shows that cross terms vanish, while the generalized dynamics preserve the same structural simplifications.

Fixed covariance setting: The target comparison is the prediction risk of the trained LSA model against the risk of the best linear predictor.The best linear predictor is defined by minimizing expected squared prediction error over linear weights.
Fixed covariance setting: Theorem 4.2 analyzes prompts whose labeled and query examples are independently sampled from a distribution with Gaussian covariates.The assumptions require the relevant first, second, and fourth-order moments to exist and be finite.
Error decomposition: The prediction-error decomposition separates several terms, after which independence and conditional-expectation arguments eliminate all cross terms.The proof therefore reduces the calculation to the remaining principal terms.
Random covariance setting: For random prompt covariances, the dynamical system is obtained by averaging over covariance matrices, without requiring them to be diagonal almost surely.The diagonal-covariance case is treated as a special case with a simplified system.
Random covariance setting: The generalized dynamics preserve zero off-diagonal blocks, and the associated loss has global minima characterized by diagonal U11 structure up to a scaling symmetry.The scaling transformation (u−1, uii) → (cu−1, c−1uii) preserves global optimality.

C.3 PL Inequality and global convergence

For random covariance matrices, the proof establishes balanced parameters, positivity of u−1, and a PL inequality under the initialization assumptions. These results imply convergence of gradient flow to the global minimum of the generalized loss.

Balanced trajectories: The balanced condition holds throughout gradient flow for the generalized random-covariance loss.The proof uses the assumed balanced initialization and the structure of the covariance-dependent dynamics.
Trajectory control: The loss remains non-increasing along gradient flow, supporting the positivity and lower-bound arguments for u−1.The proof combines monotonicity with the loss value at initialization and trace-based inequalities.
Trajectory control: When the initial scale is sufficiently small, u−1 stays positive and is bounded below by a positive constant for every t ≥ 0.This keeps the trajectory away from the saddle point at the origin.
PL inequality: A PL inequality holds for the generalized loss once the lower bound on u−1 is combined with the gradient-norm estimates.The required PL constant is chosen using the coefficients governing diagonal and off-diagonal terms.
Global convergence: Gradient flow converges to the global minimal value of the generalized loss, while uij and u−1 converge to the stated limits.The convergence rate follows by applying the gradient-flow dynamics and Grönwall’s inequality to the PL condition.

E Experiment details

The experiments use GPT2 and controlled Gaussian covariate distributions to evaluate trained models under fixed and random covariance settings. Ordinary least squares serves as the baseline for noiseless linear regression, with errors averaged across prompts and confidence intervals estimated by bootstrap.

Model and data: The experiments use standard GPT2 with 256-dimensional embeddings, 12 layers, and 8 attention heads.The implementation follows the HuggingFace architecture and adapts the embedding method of Garg et al.
Model and data: Covariates are mean-zero Gaussian vectors in d = 20 dimensions, with either identity covariance or random diagonal covariance.Random diagonal entries are independently sampled from a standard exponential distribution.
Evaluation: Ordinary least squares is used as the test-time baseline because it is optimal for noiseless linear regression tasks.Evaluation varies the covariance distribution, including random diagonal covariances multiplied by scaling coefficients.
Evaluation: Figure 1 reports error averaged over 642 prompts, using 64 covariance matrices and 64 prompts per covariance matrix for each curve.The evaluation therefore aggregates results across multiple covariance draws and prompt samples.
Evaluation: The experiments compute 90% confidence intervals from 1000 bootstrap trials for each test.The reported intervals quantify uncertainty around the averaged errors.