PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shivam Garg, Dimitris Tsipras, Percy Liang, Gregory Valiant

arXiv:2208.01066v3cs.CLcs.LG

TL;DR

The paper asks whether Transformers can genuinely learn unseen tasks from in-context examples rather than rely on memorized training tasks. It trains Transformers from scratch on synthetic function classes and finds that they learn linear, sparse linear, neural-network, and decision-tree functions, including under distribution shifts, while the implications for language models remain open.

  • Problem

    It is unclear whether language-model in-context learning reflects learning new tasks from examples or indexing tasks present in training data.

  • Method

    The paper trains standard Transformers from scratch to in-context learn well-defined function classes from function-generated input-output prompts.

  • Results

    Transformers learn linear functions with performance comparable to least squares, remain robust to distribution shifts, and also learn sparse linear functions, neural networks, and decision trees.

  • Takeaways & Limitations

    The results show that standard Transformers can encode non-trivial learning algorithms that use in-context examples, including algorithms for function classes usually handled by iterative methods.

  • Takeaways & Limitations

    How these controlled-function results relate to in-context learning in large language models remains open.

Abstract

from arXiv · show

In-context learning refers to the ability of a model to condition on a prompt sequence consisting of in-context examples (input-output pairs corresponding to some task) along with a new query input, and generate the corresponding output. Crucially, in-context learning happens only at inference time without any parameter updates to the model. While large language models such as GPT-3 exhibit some ability to perform in-context learning, it is unclear what the relationship is between tasks on which this succeeds and what is present in the training data. To make progress towards understanding in-context learning, we consider the well-defined problem of training a model to in-context learn a function class (e.g., linear functions): that is, given data derived from some functions in the class, can we train a model to in-context learn "most" functions from this class? We show empirically that standard Transformers can be trained from scratch to perform in-context learning of linear functions -- that is, the trained model is able to learn unseen linear functions from in-context examples with performance comparable to the optimal least squares estimator. In fact, in-context learning is possible even under two forms of distribution shift: (i) between the training data of the model and inference-time prompts, and (ii) between the in-context examples and the query input during inference. We also show that we can train Transformers to in-context learn more complex function classes -- namely sparse linear functions, two-layer neural networks, and decision trees -- with performance that matches or exceeds task-specific learning algorithms. Our code and models are available at https://github.com/dtsip/in-context-learning .

1 Introduction

The paper formalizes in-context learning as learning a function class from examples and studies whether Transformers can acquire this ability from scratch. It finds successful learning for linear and more complex functions, robustness to distribution shifts, and improved scaling with model capacity.

  • Motivation: The paper asks whether a model can learn new tasks from in-context examples rather than merely index tasks seen during training.This question motivates a controlled study using well-defined function classes.
  • Approach: The study trains Transformers from scratch on prompts containing function-generated input-output examples and a query input.The model predicts the function value at the query from preceding examples.
  • Linear functions: Transformers can in-context learn linear functions with error comparable to the optimal least squares estimator.The result is reported for isotropic Gaussian inputs and Gaussian-distributed linear-function weights.
  • Distribution shifts: The trained model remains quite robust under shifts between training and inference prompts and between in-context examples and query inputs.These shifts include noisy example outputs and examples and queries occupying different orthants.
  • More complex function classes: Transformers also learn 3-sparse linear functions, two-layer neural networks, and depth-4 decision trees in-context.Performance is better than least squares for sparse functions, comparable to same-architecture gradient-trained networks, and competitive against tree methods in the studied distribution.
  • Scaling: Increasing model capacity improves performance and enables in-context learning of higher-dimensional linear functions.Greater capacity often also improves performance under distribution shifts.

2 Training models for in-context learning

The paper trains models by sampling functions and inputs, constructing prefixes of in-context examples, and minimizing prediction loss across those prefixes. It uses decoder-only Transformers, from-scratch optimization, and curriculum learning to make training practical.

  • Prompt construction: Training prompts are formed by sampling a function, independently sampling inputs, and evaluating the function on those inputs.The resulting sequence contains input-output examples followed by a query input.
  • Training objective: Each training prefix contains preceding input-output examples and a next input whose function value the model must predict.The model is optimized over prefixes containing different numbers of in-context examples.
  • Model: The implementation uses a decoder-only GPT-2-family Transformer with 12 layers, 8 attention heads, and 9.5M parameters.The architecture processes vector sequences and predicts the next vector.
  • Optimization: Models are trained from scratch with squared error by repeatedly sampling random prompt batches and applying gradient updates.The reported setup uses batch size 64 and 500k training steps without text or pretrained-language-model fine-tuning.
  • Curriculum learning: Curriculum learning gradually increases function complexity, often speeding training drastically.Training begins with simpler function distributions before moving to more complex ones.

3 In-context learning of linear functions

A Transformer trained from scratch can in-context learn linear functions, matching least squares on the training prompt distribution and behaving like a learned regression algorithm.

  • Prompt distribution: The trained Transformer in-context learns linear functions under isotropic Gaussian inputs and weights, using prompts constructed from labeled examples and a query.The input dimension is d = 20, with all inputs and weight vectors sampled independently from N(0, I_d).
  • Baselines: Least squares computes the minimum-norm linear fit and is the optimal estimator, while nearest neighbors and averaging provide simpler suboptimal baselines.These baselines contextualize whether the Transformer matches optimal or simple learning procedures.
  • Results: The Transformer performs comparably to least squares for every tested number of in-context examples, substantially outperforming simpler baselines.This indicates that the model encodes a more complex algorithm than the simple comparison methods.
  • Results: 0.02 is the Transformer’s normalized squared error at d = 20 examples, decreasing to 0.0006 at 2d examples while least squares reaches 0.The error is averaged over 1280 prompts with 90% confidence intervals from 1000 bootstrap trials.
  • Beyond memorization: The model’s performance cannot be explained by memorizing training prompts or weight vectors, because it reaches error below 0.001 where the best observed training vector gives about 0.2 error.Training on only 10,000 distinct weight vectors still yields similar performance, while the best seen vector would give about 0.5 error.
  • What the model learns: With fewer than d examples, the ideal prediction is the ground-truth weight projected onto the span of the in-context inputs; the model closely follows this target.The model’s query-direction predictions and gradients align with the projected weight for all k, and with the true weight when k ≥ d.

4 Extrapolating beyond the training distribution

The trained Transformer generalizes its in-context linear-learning behavior beyond the training prompt distribution, including shifts in input distributions, noise, and relationships between examples and queries.

  • Distribution shifts: The study evaluates whether the model generalizes when training and test prompt distributions differ or when query inputs differ from in-context examples.These shifts probe the learning algorithm encoded by the model rather than only its performance on training-like prompts.
  • Overall result: The model performs reasonably accurately across these shifted prompt distributions, indicating that it learned linear regression to some generality.Inputs are generally normalized to preserve the expected squared norm seen during training.
  • Input-distribution shifts: Under skewed covariance, the model matches least squares through k = 10, then plateaus, while remaining below half the nearest-neighbor error in most cases.This shift samples inputs from a Gaussian with eigenvalues proportional to 1/i^2.
  • Input-distribution shifts: In a random 10-dimensional subspace, the model achieves errors 0.036, 0.0014, and 0.00057 at 10, 20, and 40 examples, closely matching least squares.Examples become linearly dependent between 10 and 20 examples, unlike the training prompts.
  • Label noise: With noisy labels, the model closely tracks least squares away from the input dimension and exhibits the least-squares double-descent error curve.The noise is Gaussian with unit variance, and the optimal estimator uses appropriate ℓ2-regularization.
  • Example-query shifts: The model remains close to least squares when in-context examples share an orthant different from the query, achieving errors 0.062 and 0.004 at 20 and 40 examples.The query mismatch does not substantially affect performance in this setting.
  • Example-query shifts: When the query is orthogonal to the in-context span, the model predicts near zero and achieves error close to 1 because the examples contain no relevant query information.When the query exactly matches an example, errors are 0.001, 0.001, and 0.0005 for 10, 20, and 40 examples.

5 More complex function classes

Transformers trained from scratch can in-context learn sparse linear functions, decision trees, and two-layer ReLU networks, often matching or outperforming task-specific algorithms. The neural-network-trained model also generalizes to linear functions it was not explicitly trained on.

  • Sparse linear functions: Transformers nearly match Lasso on sparse linear functions and outperform minimum norm least squares.For k = 5 and 10, the Transformer errors are 0.58 and 0.09, versus Lasso errors of 0.62 and 0.08.
  • Decision trees: With k = 100 examples, the Transformer achieves error 0.12 on decision trees, versus 0.80 for greedy learning and 0.62 for XGBoost.Providing the baselines with coordinate signs improves them to 0.50 and 0.31, but they remain worse than the Transformer.
  • Decision trees: The decision-tree results suggest that Transformers can discover an algorithm for the considered prompt distribution.The authors connect this possibility to reverse engineering the algorithm encoded by a Transformer.
  • Two-layer ReLU neural networks: Transformers can in-context learn two-layer ReLU networks with performance comparable to a neural network trained on the in-context examples.For k = 100, both methods achieve error 0.17.
  • Two-layer ReLU neural networks: The same model trained on two-layer neural networks also learns linear functions, despite not being explicitly trained on them.Its error for k = 20, 50, and 100 is 0.34, 0.05, and 0.01, compared with 0.37, 0.04, and 0.003 for the two-layer baseline.

6 Investigating what matters for in-context learning

The paper examines how problem dimension, model capacity, curriculum, and training-data diversity affect in-context learning. Capacity is especially helpful for out-of-distribution prompts, while curriculum avoids a prolonged initial period of stagnant loss.

  • Problem Dimension and Capacity: The experiments vary problem dimension across 10, 30, 40, and 50 while changing model capacity.The standard setting uses matching training and inference prompt distributions, with additional out-of-distribution evaluations.
  • Problem Dimension and Capacity: Model capacity helps in-context learning across most settings, especially for out-of-distribution prompts.The study varies dimension and capacity using 2d in-context examples and reports only small absolute gains in some in-distribution settings.
  • Curriculum: Curriculum learning gradually increases subspace dimension and prompt length until reaching the ambient dimension d and length 2d + 1.Training begins in a fixed 5-dimensional subspace with prompt length 11 and increases every 2,000 steps.
  • Curriculum: Without curriculum, training can show a relatively long period of stagnant loss followed by a sharp decrease.The duration varies with training randomness and appears to increase on average with problem dimension.
  • Number of distinct prompts or functions seen during training: The paper uses ablation studies to estimate how much training data is required for in-context learning.One study limits the number of distinct prompts seen during training.

7 Related work

The related work situates in-context learning among studies of Transformer capabilities, meta-learning, and data-driven algorithm design. This paper connects these areas by training Transformers to learn function classes from prompts.

  • In-context learning: Prior in-context-learning work includes efforts to improve and understand the capability demonstrated by GPT-3.The paper identifies Bayesian explanations and related studies as especially relevant.
  • Transformers: Transformer research spans capabilities, limitations, applications, and internal workings.The paper places its architectural focus within this broader literature.
  • Meta learning: Training a model for in-context learning is framed as a form of meta-learning or learning-to-learn.The broader paradigm includes learning how downstream learners update or learning useful parameter initializations.
  • Data-driven algorithm design: The work also relates to data-driven algorithm design, which seeks algorithms that perform well on input distributions rather than worst-case guarantees.The paper presents its approach as part of this line of work.

8 Discussion

The paper shows that Transformers can learn several function classes in-context, while leaving the implications for language models and the encoded algorithms open. It identifies model complexity, curricula, and inductive bias as directions for further study.

  • Discussion: Transformers in-context learn linear, sparse linear, neural-network, and decision-tree function classes, including tasks typically addressed with iterative algorithms.The paper reports performance comparable to least squares for linear functions and successful learning of the other classes.
  • Discussion: The results show that standard Transformers can encode nontrivial learning algorithms that use in-context examples.The authors emphasize that this behavior arises with standard Transformer architectures and optimization procedures.
  • Discussion: Whether large language models exhibit the same nontrivial in-context learning behavior remains open.The authors state that implications for language models require further investigation.
  • Future directions: Model capacity improves in-context learning accuracy and robustness in most studied cases, motivating analysis of function-class complexity, model capacity, and training-prompt count.The appropriate notion of function-class complexity may depend on the model family.
  • Future directions: Curriculum learning gradually increases function-class complexity and prompt difficulty, producing drastic training speed-ups whose cause remains unknown.The paper also asks whether similar speed-ups could apply to large language models.
  • Future directions: The framework enables comparisons of inductive bias across model families and investigation of the learning algorithms encoded inside Transformers.The authors suggest reverse engineering Transformers could yield better algorithms for some problems.

A Experimental setup

The experiments use GPT-2-family Transformers with configurations varying in capacity, adapted to scalar-output function-learning prompts. Predictions at input positions use only preceding examples and the current input.

  • Architecture: The experimental models come from the GPT-2 family as implemented by HuggingFace, with configurations varying in embedding size, layers, heads, and total parameters.The standard model is used for most experiments, while smaller models support capacity explorations.
  • Input representation: Because the tasks map low-dimensional vectors to scalars, scalar inputs and outputs are converted into vectors in the Transformer embedding space.This representation allows prompts to interleave inputs and function values.
  • Prediction rule: The model prediction at the position corresponding to x_i is treated as the prediction of f(x_i).That prediction depends only on earlier input-output pairs and the current x_i; predictions at output positions are ignored.

A.2 Training

Training samples random functions and Gaussian prompt inputs, optimizes squared prediction error with Adam, and uses curriculum learning to increase dimensionality and context length. Evaluation compares Transformers with task-specific estimators.

  • Training objective: Each training prompt samples a random function, draws Gaussian inputs, and trains predictions of function values using squared error.The loss is averaged over randomly generated prompts containing different functions and inputs.
  • Optimization: Training uses Adam for 500,000 steps with batch size 64 and learning rate 10^-4, while most runs require 5–20 hours on one RTX 3090 GPU.The standard linear model takes 17 hours for d = 50, 7 hours for d = 20, and 5.5 hours for d = 10.
  • Curriculum learning: Curriculum learning gradually increases the active input dimension and number of prompt inputs during training.For neural networks and decision trees, d_cur starts at 5 and k_cur at 26, then increases every 2,000 steps until d_cur = d and k_cur = 5d + 1.
  • Function classes: Linear functions use f(x) = w^T x with w sampled from N(0, I_d), while sparse linear, neural-network, and decision-tree classes use separately specified random constructions.The main linear setting uses d = 20; the neural-network construction uses r = 100 and the sparse setting uses d = 20 and k = 3.
  • Baselines: Least squares, nearest neighbors, averaging, Lasso, tree learning, tree boosting, and gradient-descent neural-network learning provide task-specific baselines.Minimum-norm least squares is the optimal linear-regression estimator; the other listed estimators serve as alternative baselines for relevant function classes.
  • Baselines: The decision-tree baseline greedily partitions examples by coordinates to minimize within-set output variance, while tree boosting uses ensembles of decision trees.The boosting baseline is implemented with XGBoost, and its hyperparameters are selected through a stated search.
  • Baselines: The neural-network baseline fits a two-layer network to in-context examples with ReLU activations using Adam optimization.The baseline uses batch size 10, 5,000 optimization steps, and r = 100.

B.1 Robustness to query scale

The trained model is robust to substantial query-scale changes and generally degrades gracefully on shifted prompts, often tracking least squares. Capacity improves accuracy in most settings, with exceptions and saturation depending on the shift.

  • Query-scale robustness: The model’s error increases little when query inputs are scaled up by up to 2 or down by up to 16, then degrades slowly.The experiment fixes the in-context examples while scaling only the query input.
  • Distribution shifts: Out-of-distribution prompts alter covariance, subspace, noise, orthants, query relations, or prompt scale relative to training prompts.Prompt inputs are normalized to match the training expected norm before input scaling is studied separately.
  • Query relations: Queries matching in-context examples can achieve zero error, while orthogonal queries have optimal normalized error 1 because the examples contain no relevant information.The query-matching setup makes the correct prediction directly available in the prompt.
  • Distribution shifts: The model generally degrades gracefully under shifted prompt distributions and closely tracks the least squares estimator.The shifts include skewed covariance, low-dimensional subspaces, noisy outputs, different orthants, orthogonal queries, and queries matching examples.
  • Prompt scaling: Scaling prompt inputs or weights by factors of 1/3, 1/2, 2, or 3 tests robustness to train–inference scale mismatch.The reported figure states that performance is robust to scaling weights but degrades when inputs are scaled by a factor of 3.
  • Capacity: Accuracy improves with model capacity in most settings, but scaling x shows no clear trend.With noisy outputs, accuracy nearly saturates at 1.2M parameters; with orthogonal queries, even the 0.2M-parameter model reaches error 1.

B.4 Training variance

Across three random seeds, standard Transformer error is generally concentrated, with higher variance emerging for different-orthants and skewed-covariance prompts at dimensions d ≥ 30.

  • Three models with different random seeds were trained for each dimension to measure error variability.
  • Error is concentrated in the standard setting and for most out-of-distribution prompts.
  • Higher-dimensional different-orthants and skewed-covariance settings show high variance at d ≥ 30.

B.5 Curriculum

Curriculum training substantially accelerates optimization for in-context linear-function learning, while usually leaving final performance similar. Additional experiments examine training-data requirements and whether memorization explains performance.

  • Curriculum: Curriculum often trains a low-error model 4 times faster for functions in 20 or more dimensions.
  • Curriculum: Without curriculum, the loss can remain flat within 500k steps, including one d = 30 run and all d = 50 runs.
  • Curriculum: Without curriculum, training shows an initial loss lull whose length varies and appears to increase with dimension.
  • Curriculum: For d = 20, curriculum produces no major qualitative final-error difference in most settings, while without curriculum performs slightly better under skewed covariance.
  • Training data: Models trained with limited distinct prompts can approach unrestricted-model error with 1M distinct prompts or functions, according to Figure 14.
  • Memorization: A memorization explanation is weakened because similar prompts occupy an 800-dimensional space, while trained models outperform nearest-training-weight-vector estimates.
Loading 2208.01066v3…