PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformers as Algorithms: Generalization and Stability in In-context Learning

Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, Samet Oymak

arXiv:2301.07067v2cs.LGcs.CLstat.ML

TL;DR

The paper asks how to analyze and explain transformers that perform in-context learning as implicit algorithms. It formalizes ICL through multitask learning and algorithmic stability for independent and dynamical prompts, obtaining generalization and transfer results. Experiments show near-optimal regression behavior, increasing stability with longer prompts, and predictable transfer dependence on task complexity and the number of training tasks.

  • Problem

    The paper studies how transformers can perform in-context learning by constructing prediction rules from prompt examples, including when prompts contain independent data or dynamical trajectories.

  • Method

    The paper models a transformer as a learning algorithm and relates multitask and transfer generalization to algorithmic stability across i.i.d. sequences and dynamical-system trajectories.

  • Results

    ICL implements near-optimal algorithms on flagship regression problems, becomes more stable with longer prompts, and has transfer risk governed by task complexity and the number of multitask-learning tasks.

  • Takeaways & Limitations

    The results support viewing ICL as on-the-fly algorithm learning whose generalization can be studied through stability and task-level inductive bias.

  • Takeaways & Limitations

    The guarantees include full-state observations for dynamical systems, while extending them to partial observations and other dynamic settings remains an open question.

Abstract

from arXiv · show

In-context learning (ICL) is a type of prompting where a transformer model operates on a sequence of (input, output) examples and performs inference on-the-fly. In this work, we formalize in-context learning as an algorithm learning problem where a transformer model implicitly constructs a hypothesis function at inference-time. We first explore the statistical aspects of this abstraction through the lens of multitask learning: We obtain generalization bounds for ICL when the input prompt is (1) a sequence of i.i.d. (input, label) pairs or (2) a trajectory arising from a dynamical system. The crux of our analysis is relating the excess risk to the stability of the algorithm implemented by the transformer. We characterize when transformer/attention architecture provably obeys the stability condition and also provide empirical verification. For generalization on unseen tasks, we identify an inductive bias phenomenon in which the transfer learning risk is governed by the task complexity and the number of MTL tasks in a highly predictable manner. Finally, we provide numerical evaluations that (1) demonstrate transformers can indeed implement near-optimal algorithms on classical regression problems with i.i.d. and dynamic data, (2) provide insights on stability, and (3) verify our theoretical predictions.

1 Introduction

The paper formalizes in-context learning as a transformer-based algorithm learning problem and studies its generalization, stability, and transfer across tasks. Experiments show ICL can implement near-optimal algorithms on regression and dynamic-data problems while exhibiting theoretically predicted stability and transfer behavior.

  • ICL uses a transformer to adapt to new tasks from instructions and a small number of in-context examples without updating model weights.The paper interprets this adaptation as on-the-fly inference from prompt examples.
  • ICL matches Bayes-optimal weighted linear ridge regression on noisy linear-regression tasks and outperforms autoregressive least-squares estimators on partially observed dynamics.The model leverages task priors learned during training to implement competitive machine-learning algorithms.
  • The paper abstracts a transformer as a learning algorithm that constructs a hypothesis function from the prompt and predicts outputs for new inputs.This formalism supports an implicit optimization view in which the model learns from prompt data at inference time.
  • The analysis derives multitask generalization bounds for prompts consisting of i.i.d. examples or trajectories generated by dynamical systems.The bounds relate excess risk to algorithmic stability and address temporally dependent prompts.
  • ICL predictions become more stable as prompt length increases, and training with noisy data further promotes stability.Theoretical analysis identifies conditions under which self-attention has favorable stability properties by controlling token influence.
  • Transfer risk is governed predictably by task complexity and the number of multitask-learning tasks, with little dependence on transformer architecture complexity.The paper interprets this pattern as an inductive bias arising from source-task learning.

2 Problem Setup

The problem setup represents ICL as an algorithm that maps prompt sequences to prediction functions across multiple tasks. It covers independent labeled examples and dynamical trajectories, then defines empirical and population risks for multitask and transfer learning.

  • A length-m prompt contains m−1 in-context examples and a final input whose output the transformer must predict.The prompt is represented as a sequence of examples z_i followed by x_m.
  • For independent data, prompts contain i.i.d. (input, label) pairs followed by a new input, with the transformer predicting its label.This setting uses samples from a distribution over X × Y.
  • For dynamical systems, prompts are observed trajectories in which x_m is paired with the next state generated by noisy dynamics.The state-observed setting uses x_m+1 = f(x_m) + noise and treats x_m+1 as the label.
  • The model is trained across T tasks, each supplying an independently sampled sequence of n examples and its subsequences.Training feeds task sequences to the transformer autoregressively through multitask learning.
  • The transformer is abstracted as an algorithm Alg that maps a data sequence to a prediction function, and training searches for the optimal algorithm in a hypothesis class.The resulting objective is expressed through task-level and task-averaged empirical losses.
  • Generalization analysis compares empirical and population risks, while transfer risk averages population risk over a newly sampled target task.The setup distinguishes multitask risk from evaluation on previously unseen tasks.

3 Generalization in In-context Learning

The paper derives multitask generalization guarantees for in-context learning by treating transformers as algorithms and controlling their sensitivity through algorithmic stability. The analysis covers i.i.d. and dynamic sequences, characterizes stability for transformer architectures, and shows that full-sequence training is important for transfer performance.

  • Generalization framework: ICL generalization is analyzed by modeling a transformer as an algorithm that infers functional relationships from prompt sequences.The framework studies multitask learning with sequences of examples and extends to dynamical-system trajectories.
  • Generalization framework: Algorithmic stability controls the effect of one prompt example on all later predictions, enabling generalization analysis despite temporal dependencies.The paper uses stability conditions and covering or chaining arguments to obtain uniform concentration bounds.
  • Empirical verification: Training on all n=40 prompts lets ICL learn linear regression on new tasks, whereas training on separate n/4=10 trajectory segments fails even within the target ranges.The comparison uses T=6.4 million random linear regression tasks and evaluates transfer risk.
  • Generalization bounds: Theorem 3.5 gives high-probability excess MTL-risk bounds for stable algorithm classes, with dependence on sequence length, task count, and algorithm-space complexity.The bounds use covering numbers, while stronger pairwise stability yields a Dudley-entropy-integral bound.
  • Generalization bounds: For richer task function classes, the algorithm distance and covering number become larger, making the generalization bound task-dependent.For realizable tasks, the distance depends on the function class, such as linear models or neural networks.
  • Generalization bounds: The main statistical feature is a 1/√n factor obtained despite temporal dependencies, while multiple independent sequences improve the rate to 1/√(nMT).Task-count dependence is comparatively straightforward because tasks are sampled independently.

4 Generalization and Inductive Bias on Unseen Tasks

The paper studies transfer risk when ICL is evaluated on unseen tasks and finds that performance depends predictably on source-task coverage, task complexity, and source-target distance.

  • Transfer risk: Transfer risk measures ICL performance on a target task drawn after multitask pretraining on T source tasks.The analysis compares transfer risk with multitask-learning risk under distribution shift.
  • Transfer risk: Transfer risk decays as 1/poly(T), while additional samples n or sequences per task M typically cannot eliminate the unseen-task distribution shift.
  • Inductive bias: Experiments across dimensions d=5, 10, 20 show nearly aligned risk curves as T varies over {1, 2, 5} × d^2, indicating predictable dependence on task complexity and task count.The experiments use isotropic features, noiseless labels, and Gaussian task vectors.
  • Inductive bias: The proposed explanation is an optimization-induced inductive bias toward an algorithm determined by the span of the source-task parameters.The paper presents this explanation as a conjecture and leaves the theoretical analysis of the observed d^2/T behavior for future work.
  • Covariance estimation: ICL may require T proportional to d^2 for strong entrywise covariance control, whereas weighted least squares requires T=Ω(d) for spectral control in isotropic tasks.This contrast suggests ICL training may not be sample-optimal in the number of source tasks.
  • Source-target distance: For linear tasks, transfer performance worsens as the target moves farther from the source tasks, with smaller cosine similarity corresponding to poorer generalization.Source and target task vectors are normalized, so Euclidean distance and cosine similarity encode equivalent proximity here.

5 Extension to Stable Dynamical Systems

The paper extends ICL generalization analysis from i.i.d. examples to dependent prompts generated by stable dynamical systems, using system stability to control perturbation effects over trajectories.

  • Problem setting: Prompts are trajectories generated by task-specific dynamical systems with bounded states and i.i.d. bounded process noise.Each task has a system f_t and a sequence generated recursively from an initial state.
  • Stability assumptions: Generalization requires exponential stability of the dynamical systems, which controls distances between trajectories initialized from different states.This system stability condition is distinct from the algorithmic stability imposed on the transformer-induced algorithm.
  • Stability assumptions: The analysis also assumes algorithmic stability with respect to Euclidean perturbations caused by swapping a trajectory noise term.This extends the stability framework used for i.i.d. prompts to dependent dynamic data.
  • Guarantee: Theorem 5.4 gives the same generalization bound as Theorem 3.5 after updating the stability constant to account for dynamical-system effects.The updated constant is expressed using the system-stability parameters and bounded state and noise scales.
  • Proof strategy: The proof controls the impact of a perturbation by tracking how it propagates through the future trajectory under system stability.

6 Interpreting In-context Learning as a Model Selection Procedure

The paper interprets ICL as an algorithm that can select among hypothesis classes from data, linking transformer generalization to adaptive model selection rather than training over one large union class.

  • Model-selection view: ICL can be viewed as an implicit model-selection procedure in which a transformer implements an algorithm that competes with ERM.The framework assumes such an algorithm can approximate ERM across sample sizes and hypothesis classes.
  • Model-selection view: The competing hypothesis classes may represent different model families, such as convolutional networks, fully connected networks, decision trees, or constrained linear classes.Nested classes can also encode capacity control, including covariance-indexed linear classes.
  • Statistical formulation: The minimum achievable risk within each class is combined with its Rademacher-complexity-controlled estimation risk to characterize model-selection performance.
  • Generalization benefit: Adaptive selection across sample sizes avoids the less favorable complexity dependence associated with training over the union of all hypothesis classes.The result applies even though the algorithm class can have VC-dimension n log H when H sample-size-specific ERM algorithms are considered.
  • Scope: The discrete-family presentation uses H<∞ classes, while the paper's Section 3 theory also allows continuous settings.

7 Numerical Evaluations

The numerical evaluations test ICL on regression, dynamical-system, and stability settings. Across these experiments, ICL matches or outperforms conventional baselines while showing risk changes that decrease with longer prompts.

  • Linear regression: ICL matches optimally tuned ridge or weighted regression and outperforms least-squares solutions on linear regression tasks.These results provide evidence that transformers learn task priors and can perform automated model selection.
  • Partially-observed dynamical systems: ICL outperforms least-squares predictors using history windows H=1, 2, 3, 4 for partially observed dynamical systems.Against optimally tuned ridge regression, the performance gap is narrower but remains favorable to ICL.
  • Stability analysis: The stability experiments evaluate linear models, 3-sparse linear models, depth-4 decision trees, and 2-layer ReLU networks with 100 hidden units.Each function class uses input dimension 20 and is trained with a GPT-2 architecture.
  • Stability analysis: Risk changes decrease after a certain sample size when an in-context example is replaced and its label is flipped.The experiments are consistent with the stability assumption, although they do not establish whether stability is O(1/m).
  • Stability analysis: Adding label noise to regression tasks during multitask training can improve stability.

8 Conclusions

The paper frames ICL as algorithm learning and connects its generalization to algorithmic stability for both example sequences and dynamical trajectories. It identifies open questions about task-level guarantees, broader dynamical settings, and the inductive bias governing transfer risk.

  • Conclusions: The paper presents multitask generalization bounds for transformers trained on sequences containing n examples across T tasks.The analysis builds on algorithmic stability verified for transformer architectures theoretically and empirically.
  • Conclusions: The generalization and stability guarantees extend to dynamical systems, reflecting the autoregressive nature of transformers.
  • Open questions: Transfer risk is governed by the number of multitask-learning tasks and task complexity, while appearing independent of model complexity.The paper identifies characterizing the learned algorithm as an open question.
  • Open questions: The bounds mainly capture average multitask-learning risk rather than generalization on individual tasks or prompts with specific lengths.The authors pose controlling these more specific forms of generalization as an open problem.

A.1 Supporting Experiments for Section 4

Supporting experiments examine transfer and MTL risk, greedy task selection, and transformer stability. Results show predictable task-scaling behavior, near-optimal regression performance, and a stabilizing effect of label noise.

  • Transfer and MTL risk: Transfer risk is mostly independent of model complexity dim(A), unlike its dependence on task complexity d.This comparison uses tiny, small, and standard GPT-2 models across varying numbers of five-dimensional linear-regression tasks.
  • Transfer and MTL risk: T=3d source tasks already yields performance very close to optimal weighting with the true covariance Σ.The comparison concerns empirical-covariance-weighted ridge regression versus Σ-weighted ridge regression.
  • Greedy task selection: A greedy search over T source-task parameters identifies the best-fitting task in n≲log(T) samples, while MTL risk changes little as T varies among d, 2d, and 3d.The experiments compare greedy MTL with transformer MTL and least-squares behavior.
  • Transfer and MTL risk: With M=1, increasing T from 2000 to 50000 reduces MTL risk, while transfer and MTL risks remain almost perfectly aligned.The experiment fixes d=5 and n=10 and varies the number of linear-regression tasks.
  • Stability experiments: Adding label noise improves stability: noisy regression is more stable than noiseless regression, whose perturbation peaks occur around n=d.The proposed explanation is that noise induces weighted ridge regression, regularizing the learned algorithm; Figure 12 compares noisy and noiseless settings.
  • Stability experiments: Under bounded unit-ball inputs and bounded attention weights, attention perturbations satisfy ∥P∥2,1≤(2Γ+1)e2Γ∥E∥2,1.The bound combines perturbations from changed attention weights and changed values.

B.1 Proof of Theorem 3.2

The proof establishes transformer stability by controlling how token perturbations propagate through attention and layers. It also identifies settings where this guarantee is tight or fails, particularly for large attention weights or perturbations to the final token.

  • Stability guarantee: A D-layer transformer satisfies stability with K=((1+Γ)eΓ)D CρL under the stated dynamical-system assumptions.The result applies when trajectories are bounded and the system satisfies the stable-dynamics assumption.
  • Failure modes: Stability can fail when attention weights are large, showing that the bounded-Γ condition is operationally important.The paper explicitly presents instability for large Γ as a complement to the stability theorem.
  • Failure modes: The theorem’s earlier-token perturbation condition is tight because each token strongly influences its own self-attention output.This explains why the final token is excluded from the same stability guarantee.
  • Failure modes: Perturbing the last token can change its output by at least 0.2 even when the sequence length n varies.The construction uses a single self-attention layer with W=V=I and sets the perturbation amount to ε=1.

C.1 Proof of Theorem 3.5

The proof derives MTL and transfer generalization bounds by concentrating empirical and population risks over stable algorithm hypotheses. It uses martingale concentration for prompt dependence and covering arguments for continuous algorithm spaces.

  • MTL generalization: The excess MTL-risk bound is obtained by controlling sup_Alg|L(Alg)−bL(Alg)| over the algorithm class.The proof decomposes excess risk and bounds concentration for empirical-risk minimization.
  • MTL generalization: Continuous algorithm classes are handled with ε-covers, chaining, and union bounds over progressively refined covers.The proof bounds both approximation to the cover and deviations on the discrete cover.
  • MTL generalization: For training on all task samples, the same bounds extend by replacing T with T_M.This is stated as an identical-bound result under the corresponding empirical-risk formulation.
  • Transfer learning: A greedy source-task search needs n∝log(T) samples, contrasting with empirically closer-to-n∝d behavior for MTL.The greedy method searches over source parameters and returns the source task best fitting the prompt.
  • Transfer learning: Transfer risk is bounded by MTL risk under task diversity: R_T(dAlg)≤R_MTL(dAlg)/ν+2ε.The statement applies when the target task is (ν,ε)-diverse over the T source tasks.

D Proof of Theorem 5.4

The dynamical-system proof adapts the MTL concentration argument to dependent trajectories. Stability of the dynamics and bounded states control perturbation propagation, enabling the corresponding excess-risk bound.

  • Concentration argument: Changing one noise term is controlled through dynamical stability, which bounds its effect on later states and losses.The proof applies the perturbation lemma and the stable-dynamics assumption to the affected suffix of the trajectory.
  • Final bound: The resulting excess MTL test-risk theorem assumes bounded Lipschitz loss and stable dynamical systems.The theorem is stated for the ERM solution in the dynamical setting with probability at least 1−2δ.
  • Dynamical-system setup: The proof treats each task as an independent trajectory generated by a stable dynamical system with process noise.The trajectories obey xᵗ_i=fᵗ(xᵗ_{i−1})+wᵗ_i.
  • Concentration argument: Martingale concentration is rebuilt using filtrations generated by the initial state and successive process-noise terms.The resulting sequence exposes the trajectory dependence increment by increment.

E Model Selection and Approximation Error Analysis

The analysis assumes transformers can approximately perform in-context ERM and studies model selection over hypothesis classes, relating approximation and test-risk bounds to hypothesis complexity and Rademacher complexity.

  • Model-selection assumption: The analysis assumes transformers can approximately perform empirical risk minimization over a family of hypothesis classes.This assumption lets transformer-induced algorithms compete with ERM over general hypothesis classes.
  • Model-selection assumption: For prompt length m, the analysis searches only the hypothesis class F_h^m rather than the full union F_all.The selected class satisfies dim(F_h^m) ≤ dim(F_all).
  • Model-selection assumption: The hypothesis family contains H countable classes and can represent discrete or continuous model-selection settings.Examples include norm-bounded linear models akin to ridge regression and covariance-weighted linear models akin to weighted ridge.
  • Risk analysis: Classical ERM generalization control bounds test risk using Rademacher complexity under bounded, Lipschitz loss assumptions.The loss takes values in [0, B] and is L-Lipschitz in Euclidean norm.
  • Risk analysis: The resulting analysis bounds excess test risk relative to the best hypothesis or algorithm, with approximation error controlled through the transformer’s ERM capability.The formal results define the optimal risk over the hypothesis class or algorithm family and combine this with concentration and stability-related terms.
  • Risk analysis: The framework connects ICL to meta-learning and multitask learning because training uses many related prompt instances while adaptation occurs implicitly through the prompt.The paper develops excess multitask-learning risk bounds for training across T tasks.
Loading 2301.07067v2…