Trained Transformers Learn Linear Models In-Context

• Prior work showed transformers can in-context learn function classes and mimic algorithms such as ordinary least squares, but optimization dynamics remained unclear.ICL uses labeled prompt examples to predict a query from a new task without parameter updates [Gar+22].Source: §1 · p. 2
• Gradient flow from suitable random initialization converges to a global population-loss minimum for single-layer transformers with linear self-attention, despite non-convexity.Source: §1 · p. 2
• At convergence, the transformer’s encoded learning algorithm and prediction error are characterized, establishing in-context learning of linear models on new prediction tasks.The analysis includes possibly nonlinear test tasks and shows performance competitive with the best linear predictor over the test prompt distribution.Source: §1 · p. 2
• Although several distribution shifts are tolerated, shifts in the prompt covariate distribution of features x_i cause the trained linear-self-attention transformer to fail.Source: §1 · p. 2
• Allowing covariate distributions to vary across prompts preserves global convergence guarantees but does not prevent linear-self-attention failure under covariate shift.Experiments further find that larger nonlinear transformers generalize better under covariate shift, especially with varying training covariate distributions.Source: §1 · p. 2

Transformers as in-context learners

• Transformers trained on random tasks can in-context learn linear regression, two-layer ReLU networks, and decision trees, implementing distinct algorithms across data structures [Gar+22] [APG23] [AL23].Garg et al. [Gar+22] found linear-regression predictions resembling ordinary least squares.Source: §2 · p. 3

Linear-regression mechanisms

• For isotropic Gaussian linear regression, prior analyses found transformer predictions mimic one gradient-descent step, whereas this work proves gradient flow learns linear models even with anisotropic covariates [Aky+22] [Osw+22].A single vanilla gradient-descent step cannot generally achieve small error for anisotropic covariances.Source: §2 · p. 3

Interpretations of in-context learning

• Other accounts interpret in-context learning as kernel regression, implicit finetuning, Bayesian inference or selection, gradient descent, and prompt-based inference of label and input structure [Han+23] [Dai+22] [Xie+21] [WZW23] [Li+23a] [Min+22].These works provide complementary algorithmic and probabilistic interpretations of ICL across prompts and tasks.Source: §2 · p. 3

Approximation and computational power

• Approximation-theoretic studies establish that transformers can universally approximate sequence functions, learn sparse functions, achieve Turing completeness, approximate automata, and implement statistical algorithms [Yun+19; Yun+20; Ede+22; LCW21; PMB19; Deh+19; BPG20; Liu+23; Bai+23].Related work also uses pretrained transformers for sparse retrieval and interprets ICL through Bayesian model averaging [Abe+23] [Zha+23].Source: §2 · p. 3

Provable optimization guarantees

• A concurrent study identified global minima for anisotropic Gaussian data, while this work proves suitable-initialization gradient flow reaches such minima and characterizes new-task, shift, and cross-prompt covariate effects [Ahn+23].Related optimization analyses also study vision transformers and topic modeling under gradient-based training [JSL22] [LLR23].Source: §2 · p. 4

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 framework treats ordinary learning algorithms such as least squares as in-context learners and defines performance through prediction error on sufficiently long prompts.Source: §3.1 · pp. 4–5
• The paper studies whether one-layer linear self-attention models can learn linear hypotheses under squared loss and Gaussian covariate distributions.The stated questions concern optimization, context lengths, and prediction error for possibly anisotropic Gaussian marginals.Source: §3.1 · p. 5

3.2 Linear self-attention networks

• The analyzed LSA model removes softmax, merges query-key and value-projection matrices, and predicts from the bottom-right output entry.Prompts embed labeled examples as columns followed by a query column with a zero label; under this embedding, the normalization factor is ρ = N.Source: §3.2 · p. 6

3.3 Training procedure

• Training prompts use Gaussian covariates with covariance Λ and Gaussian task weights, while the population objective is optimized by gradient flow from a balanced initialization.The initialization condition is satisfied almost surely by a normalized random matrix construction and is motivated by balancing the equivalent two-layer linear network.Source: §3.3 · pp. 7–8

4.1 Convergence of gradient flow and prediction error for new tasks

• Gradient flow from suitable initialization converges to a global minimum of the population loss for the linear self-attention transformer.This result applies to the fixed-covariate setting analyzed in Section 4.1.Source: §4.1 · pp. 8–9
• For sufficiently large training-prompt length N and test-prompt length M, a transformer trained by gradient flow from in-context examples achieves prediction error competitive with the best linear model.Ignoring dimension- and covariance-dependent factors, the prediction error is at most O(1/M + 1/N2); even as M →∞, the gap does not vanish unless N →∞ as well.Source: §4.1 · p. 10
• For noiseless linear tasks, the transformer in-context learns linear models up to error η = (1 + 2d + d2κ) tr(Λ)/N2 when M ≥(d + 1) tr(Λ)ε−1.It then achieves prediction error at most η + ε; larger dimension, covariance variance, and condition number increase the error.Source: §4.1 · p. 11

4.2 Behavior of trained transformer under distribution shifts

• The trained transformer tolerates task shifts and broad query shifts, remaining competitive with the best linear model when prompt lengths are sufficiently large.Query shifts can fail when the query depends on the training data; for example, an orthogonal query receives prediction zero.Source: §4.2 · p. 12
• Covariate shifts are not fully tolerated: changing the test feature distribution can prevent the transformer from matching ordinary least squares, which is robust to constant feature rescaling.This failure also appeared in transformer experiments by Garg et al. [Gar+22].Source: §4.2 · p. 12

4.3 Transformers trained on prompts with random covariate distributions

• With random covariate distributions across training prompts, gradient flow still reaches a global minimum, but the linear self-attention transformer fails to in-context learn linear models across those distributions.The failure occurs because the limiting normalization generally does not match a new prompt’s covariance; larger nonlinear transformers are reported as more robust under covariate shifts.Source: §4 · p. 8, §4.3 · pp. 14–15
• The GPT2 experiments compare fixed identity covariance with random diagonal covariance training across test scales c = 1, 4, 9, against least squares and linear self-attention limits.Training uses prompt lengths N = 40, 70, 100 in d = 20 dimensions, with random covariance diagonals sampled i.i.d. from Exponential(1).Source: §4.3 · p. 15
• GPT2 models trained on random covariance matrices improve robustness to higher-variance test prompts, but still fail to match least squares at the largest scale, c = 9.Models trained on fixed covariance succeed only when variance is not too large, whereas random-covariance training performs better on higher-variance prompts.Source: §4.3 · p. 15, Figure 1 · p. 16
• Prediction error spikes when test prompts exceed training context length, regardless of whether covariance is fixed or random, although higher variance reduces the spike.The experiments use d = 20 and training context lengths N ∈ {40, 70, 100}.Source: §4.3 · p. 15, Figure 1 · p. 16
• The context-length failure is conjectured to result from randomly initialized absolute positional encodings for positions beyond training length, consistent with prior observations [Dai+19] [Ani+22] and concurrent findings [APG23].Position i is updated only when training encounters context length i; removing positional encoders improves larger-context evaluation in concurrent work [APG23].Source: §4.3 · p. 16

5.1 Equivalence to a quadratic optimization problem

• The transformer’s prediction is a quadratic form u⊤H_τu, reducing training to a rank-one matrix-factorization problem whose quadratic objective is non-convex.The non-convexity follows because H_τ has at least d + 1 negative eigenvalues almost surely.Source: §5.1 · pp. 17–18

5.2 Dynamical system of gradient flow

• Gradient flow follows a coupled dynamical system equivalent to optimizing a reduced objective, whose global minima satisfy u_−1U_11 = Γ^−1 and yield the transformer’s limiting predictions.The parameter coordinates u_12 and u_21 remain zero, while u_−1(t)U_11(t) approaches Γ^−1 as N grows.Source: §5.2 · pp. 18–19

5.3 PL inequality and global convergence

• A Polyak–Łojasiewicz inequality implies that, under the stated initialization condition, gradient flow converges globally and the key coordinates U_11 and u_−1 reach explicitly characterized limits.These convergence results complete the proof of Theorem 4.1 after translating the reduced coordinates back to the original transformer parameters.Source: §5.3 · p. 19

Main conclusions

• Gradient flow from suitable random initializations converges to a global population-loss minimum despite non-convexity, for transformers trained on noiseless linear-model prompts.The analysis assumes anisotropic Gaussian covariates and a single linear self-attention layer.Source: §6 · p. 19

Distribution shifts

• The trained transformer is robust to task and query distribution shifts but brittle when training and test covariate distributions differ, matching Garg et al. [Gar+22].This limitation contrasts with ordinary least squares, which can perform well for well-specified linear regression across broader covariate distributions.Source: §6 · p. 20

Future work

• Future work should test finite-step stochastic gradient descent, more general initializations, and the dynamics of nonlinear and deep transformers beyond the single-layer linear model.The current guarantees specifically concern gradient flow on the population loss with particular random initialization schemes.Source: §6 · p. 20

Distribution shifts

• When covariate distributions vary across prompts, single-layer linear transformers fail even when test covariate distributions are sampled from the same distribution; larger nonlinear transformers appear more successful but remain sub-optimal.Understanding this setting is identified as an important direction for future research.Source: §6 · p. 20

A Proof of Theorem 4.1

• Theorem 4.1 is obtained as a natural corollary of Lemmas 5.1, 5.2, 5.3, and 5.4 after translating u−1 and U11 back to W PV and W KQ.Source: A Proof of Theorem 4.1 · p. 22

A.1 Proof of Lemma 5.1

• The section restates Lemma 5.1 before providing its proof for reader convenience.Source: §A.1 · p. 22
• Lemma 5.1 states that the query prediction for a prompt embedding matrix Eτ can be written as the output of a quadratic function.Source: §A.1 · p. 22
• The proof begins by decomposing W PV and W KQ into the components used to express the prediction formula.Source: §A.1 · p. 22
• The query prediction is identified with the right-bottom entry of fLSA(Eτ) under the definition of the linear self-attention map.Source: §A.1 · p. 22
• The derivation uses algebraic identities for matrix vectorization, Kronecker products, and traces.Source: §A.1 · p. 23

Proof setup

• The proof computes gradient-flow derivatives from L using the chain rule and matrix-derivative identities.These calculations establish the differential equations used in the subsequent blockwise analysis.Source: §A.2 · p. 23

Steps One–Three

• Gaussianity and independence of wτ, prompt inputs, and query inputs simplify the derivative terms and force the relevant cross-block contributions to vanish.The argument repeatedly uses mean-zero Gaussian wτ and independence relations among covariates and query inputs.Source: §A.2 · pp. 23, 24, 26

Step Three: u12 and u21 Vanish

• Under Assumption 3.3, gradient flow preserves u12(t) = u21(t) = 0d for all t ≥ 0.The proof shows these blocks have zero derivatives whenever they are initially zero, establishing an invariant subspace.Source: §A.2 · pp. 25–26

Step Four: Dynamics of U11

• Given vanishing off-diagonal blocks, the proof derives the entrywise dynamics of U11.The calculation evaluates each Ukl derivative using the definitions of A and the covariance structure of the prompt inputs.Source: §A.2 · pp. 26–27

Step Five: Dynamics of u−1

• The final step computes the gradient-flow dynamics of u−1 using independence between prompt covariates and the defining matrix relations.This completes the blockwise dynamical characterization required by the lemma.Source: §A.2 · pp. 27–28

A.3 Proof of Lemma 5.3

• Gradient flow on L is equivalent to gradient flow on a loss function ˜ℓ with respect to U11 and u−1 under Assumption 3.3.The equivalence is established by recovering the loss from the gradient-flow dynamics and differentiating ˜ℓ.Source: §A.3 · p. 29
• The global minima of ˜ℓ are characterized in Corollary A.2, using the commutativity of Γ and Λ and a nonnegative matrix expression.The loss ˜ℓ may differ from L by an additive constant and can be negative, without affecting its global minimization.Source: §A.3 · pp. 29–30
• Lemma 5.3 follows because the loss is unchanged under replacing (U11, u−1) by (cU11, c−1u−1) for any non-zero constant c.This scaling symmetry transfers the minimum characterization from ˜ℓ to the parameters stated in Lemma 5.3.Source: §A.3 · p. 30

A.4 Proof of Lemma 5.4

• The parameters remain balanced throughout the trajectory, explaining why balanced initialization is imposed.The balance property holds for every t ≥ 0.Source: §A.4 · p. 30
• The proof keeps u−1(t) positive and bounded below by a positive constant, preventing trajectories from reaching the saddle point at the origin.Positivity follows from loss monotonicity, continuity, and the initial sign; the lower bound is then used to establish the PL inequality.Source: §A.4 · pp. 31–33
• Under Assumption 3.3 and a suitable initialization scale, gradient flow satisfies a PL inequality and converges to the global minimum of the loss.The PL constant is positive under the lemma’s scale condition.Source: §A.4 · pp. 33–34
• The PL inequality and gradient-flow dynamics imply convergence of u−1U11 toward Γ−1 entry-wise.The argument uses positive definiteness and commutativity of Γ and Λ together with the lower bound on u−1.Source: §A.4 · p. 34

B Proof of Theorem 4.2

• Under Gaussian covariates and finite-moment assumptions, the proof analyzes the trained LSA predictor’s excess risk relative to the best linear predictor on a new task.The test prompt contains i.i.d. labeled examples and a query point from the task distribution.Source: B Proof of Theorem 4.2 · p. 35
• The excess-error decomposition simplifies because all cross terms vanish, leaving only the principal terms needed to complete the risk calculation.The proof uses conditional expectation and independence of the prompt examples to eliminate cross terms.Source: B Proof of Theorem 4.2 · pp. 35–36

C Proof of Theorem 4.5

• The proof of Theorem 4.5 extends the dynamical-system argument because the relevant lemma remains valid independently of the training data and data-generating distribution.This permits analysis when the prompt covariance is random.Source: C Proof of Theorem 4.5 · p. 37

C.1 Dynamical system

• For i.i.d. random positive-definite covariance matrices with finite third moment, gradient flow follows an explicit matrix dynamical system while the off-diagonal blocks remain zero.The result is obtained by conditioning on each covariance matrix and then taking expectation.Source: §C.1 · p. 38
• When the random covariance matrices are diagonal almost surely, the matrix dynamics reduce to an equivalent entrywise system for the entries of U11.The simplification follows by rewriting each U11 equation using diagonal covariance structure.Source: §C.1 · pp. 38–39

C.2 Loss function and global minima

• In the random-covariance setting, the dynamical equations are gradient flow on an explicit loss function, up to an additive constant and under the stated initialization condition.This equivalence is verified by differentiating the loss with respect to uij and u−1.Source: §C.2 · p. 39
• The cited characterization states that every global minimum has zero off-diagonal entries in U11, and replacing (u−1, uii) with (cu−1, c−1uii) for c ≠ 0 also achieves a global minimum.Source: §C.2 · pp. 39–40

C.3 PL Inequality and global convergence

• The balanced parameter condition is preserved throughout the trajectory, ensuring u−1(t)^2 equals the relevant trace quantity for every t ≥0.This invariant is established before proving the PL inequality and global convergence.Source: §C.3 · pp. 40, 44
• With sufficiently small initial scale, u−1 remains positive for all time, so the trajectory stays away from the saddle point at the origin.The loss is non-increasing, begins negative under condition (C.10), and would be zero whenever u−1 = 0.Source: §C.3 · pp. 40–41
• A PL inequality holds for ℓrdm once u−1 is lower-bounded by a positive constant, converting gradient-flow decrease into exponential convergence toward the global minimum.The proof compares the loss gap with the squared gradient norm and applies Grönwall’s inequality.Source: §C.3 · pp. 42–44
• Gradient flow converges to the global minimum of ℓrdm, with uij and u−1 approaching the limits specified by Lemma C.8.The result follows by combining the PL inequality with gradient-flow dynamics and Grönwall’s inequality.Source: §C.3 · pp. 42–44

D Technical lemmas

• The proof uses matrix-derivative, Kronecker-product, and vectorization identities together with Gaussian moment formulas from Isserlis’ theorem.These results are collected as Lemmas D.1 and D.2 and support the technical calculations in the convergence analysis.Source: D Technical lemmas · pp. 44–45
• Von Neumann’s trace inequality bounds trace expressions through ordered singular values, while positive-semidefinite products are characterized by matrix commutation.These facts are stated in Lemmas D.3 and D.4 and are invoked in the global-convergence proof.Source: D Technical lemmas · p. 45

Model and objective

• GPT2 uses 256-dimensional embeddings, 12 layers, and 8 heads, with separate x and y tokens, while both models’ plotted errors concern the final query prediction.The GPT2 objective averages predictions across the full context, whereas linear self-attention is trained only on the final query point [Gar+22].Source: E Experiment details · p. 46

Training setup

• Covariates are 20-dimensional mean-zero Gaussians with either identity covariance or random diagonal covariance whose entries are i.i.d. Exponential(1), while task weights are independent standard Gaussians.Models were trained for 500000 Adam steps with batch size 64, learning rate 0.0001, and Garg et al.’s curriculum [KB14] [Gar+22].Source: E Experiment details · p. 46

Evaluation setup

• Testing uses ordinary least squares as the optimal noiseless-regression baseline, with random-covariance evaluations scaling the diagonal covariance by c ∈ {1, 4, 9}.Fixed-covariance evaluation uses the identity matrix, while random-covariance evaluation samples diagonal entries from Exponential(1).Source: E Experiment details · p. 46

Evaluation statistics

• Figure 1 errors average over 642 prompts, using 64 covariance matrices and 64 prompts per covariance matrix for each curve, with 90% bootstrap confidence intervals over 1000 trials.The supplied passage contains a typographical reference to each “teat” when describing the confidence-interval computation.Source: E Experiment details · p. 46