In-context learning of simple function classes

Context. Language models pick up a task from a few examples in the prompt, and it was unclear whether they learn from the examples or retrieve a task seen in training. This paper makes the question well posed: train a transformer from scratch on prompts $(x_1,f(x_1),\dots,x_k,f(x_k),x_{\rm query})$ with $f$ drawn from a known function class, and compare its predictions at $x_{\rm query}$ with estimators whose behavior is understood. Training on many random tasks turns the forward pass into a learning algorithm, and for linear functions that algorithm comes out close to least squares, the Bayes rule for the training prior; how the model behaves away from that prior tells you which algorithm it actually runs. The target is the minimum-norm interpolant of the benign-overfitting papers, and the model inherits its double-descent peak when label noise appears at test time.

Setting. Training minimizes the squared loss averaged over every prefix of the prompt, which causal masking computes in one forward pass. The model is a 12-layer GPT-2 decoder (9.5M parameters) in which each $f(x_i)$ is its own token, trained for 500k steps on fresh prompts (32M tasks) with a curriculum that starts in a 5-dimensional subspace with 11 examples and grows. The main class is linear: $d=20$ , $x_i,x_{\rm query},w\sim N(0,I_d)$ , $y_i=w^\top x_i$ , with error $(M(P)-w^\top x_{\rm query})^2/d$ , so the zero predictor scores 1. The baselines are minimum-norm least squares $\hat w=X^+y$ , the averaging estimator $\hat w=\frac1k\sum_iy_ix_i$ , and 3-nearest neighbors. Everything is empirical; nothing is proved.

Step 1

On in-distribution prompts the model is least squares, which is the Bayes rule here. Conditioning the noiseless prior $w\sim N(0,I_d)$ on $Xw=y$ gives $w\mid\text{data}\sim N(X^+y,\ I-P_X)$ , with $P_X$ the projection onto the span of the $x_i$ , so the posterior mean is the minimum-norm interpolant and its risk is the leftover posterior variance, $\mathbb E\,\frac{(x_{\rm query}^\top X^+y-w^\top x_{\rm query})^2}{d}=\frac{(d-k)_+}{d},$ a straight line to zero at $k=d$ . The transformer lies on it, with residual error 0.02 at $k=d$ and 0.0006 at $k=2d$ , what an approximate, regularized solver would leave. The cheap estimators that attention can easily represent are far off: the averaging estimator, one unpreconditioned gradient step from zero, has error $(d+1)/k$ (my computation, not in the paper, checked numerically) and beats the zero predictor only once $k>d+1$ , and 3-NN is still near 0.7 at $k=2d$ (by simulation; the paper gives no 3-NN numbers).

Step 2

The model does not retrieve training tasks. Even a perfect lookup is too coarse: an oracle that picks, for the test $w^\star$ , the closest of the 32M training weight vectors has error about 0.216 (estimated by sampling), while the model is below 0.001 at $2d$ examples. With only 10k distinct training weights the oracle gets 0.505, yet a model trained on those 10k comes close to the unrestricted model, and 1M fixed training prompts give near-unrestricted error. Twenty Gaussian dimensions are too many to cover by lookup, so the forward pass must estimate from the prompt.

Step 3

With the examples fixed, the model computes the right linear function of the query. Along a random ray $\lambda x$ the model follows $w^\top(\lambda x)$ for $k\ge d$ and $(P_Xw)^\top(\lambda x)$ for $k=d/2$ , as long as $\lambda$ stays near the norm of typical training inputs, and its gradient in $x_{\rm query}$ has normalized inner product near 1 with $P_Xw$ at every $k$ . So it is locally the projected least-squares predictor and, for $k\lt d$ , does not guess outside the span of the data, which is what the posterior mean does.

Step 4

Distribution shifts separate the candidate algorithms. In-context inputs confined to one random orthant, with the query elsewhere, leave the error near least squares (0.062 and 0.004 at 20 and 40 examples), which rules out nearest-neighbor style local averaging. Inputs in a random 10-dimensional subspace, linearly dependent as never in training, give 0.0014 at 20 examples, and queries orthogonal to all examples get predictions near 0, the optimum. With $N(0,1)$ label noise at test time the model tracks least squares away from $k=d$ and shows the same double-descent peak at $k=d$ ; since the Bayes rule for noisy data is ridge with penalty $\sigma^2$ , the model is running the ridgeless rule that was optimal for its noiseless training data.

In-context inputs all lie in one orthant and the query in another, and the model's error stays near least squares. Which family of explanations does this rule out?

Step 5

The failures point to a step size tuned to the training inputs. Least squares is invariant to the input scale and equivariant in $y$ , yet rescaling the inputs by $1/3,1/2,2,3$ gives errors 0.30, 0.013, 0.043, 0.58 at 40 examples, while rescaling $w$ by the same factors gives at most 0.028. With covariance eigenvalues $\propto1/i^2$ the model matches least squares up to $k=10$ and then plateaus. So the model is close to linear in the labels but carries a normalization tied to the training input distribution, which is what a fixed step size or preconditioner would produce; the paper does not test that reading, and the gradient-descent papers that follow make it precise.

Step 6

The same training learns richer classes, and the learned algorithm carries their inductive bias. For 3-sparse $w$ the model matches Lasso (0.09 against 0.08 at 10 examples), well below least squares. For depth-4 decision trees it gets 0.12 at 100 examples against 0.80 for greedy trees and 0.62 for XGBoost; with no Bayes-optimal reference, this shows only that the learned algorithm beats tuned off-the-shelf learners on this prior. For two-layer ReLU networks it matches a same-architecture network trained by Adam on the prompt, and run on linear functions this model learns at the network's slower rate instead of the least-squares rate, so the in-context algorithm carries the inductive bias of the class it was trained on.

Step 7

Larger models are more robust, and the curriculum skips a plateau. Error at $2d$ examples falls with capacity (0.2M to 9.5M parameters) and rises with dimension, with the largest gains out of distribution. Without the curriculum the loss sits flat for a long, seed-dependent stretch and then drops sharply; the stretch grows with $d$ , and at $d=50$ the loss barely moves in 500k steps. Starting in 5 dimensions, where the plateau is nearly absent, trains about four times faster for $d\ge20$ , and at $d=20$ the final model is qualitatively the same either way. The paper reports the plateau-then-jump without explaining it.

Consequences and limits. The paper turns in-context learning into a statistics question: fix a prior over tasks, compare the transformer with the Bayes rule and with candidate algorithms, and use off-distribution prompts to tell them apart. For linear regression the answer is close to minimum-norm least squares, which ties the model's risk to the spectral calculus of the interpolation papers, and the scale and covariance failures point to a built-in step size and preconditioner fit to $\Sigma=I$ , the starting point of the gradient-descent and preconditioning results that followed. The limits are that everything is empirical, for one architecture family and small $d$ ; "close to least squares" is measured by error curves and query gradients, with no inspection of the weights; the trees and networks lack an optimal benchmark; and nothing here is about how language models trained on text acquire the ability.

Terms

Prefix objective: loss averaged over all prefixes $P^i$ ; one causal pass gives every prediction
Normalized error: $(M(P)-w^\top x_{\rm query})^2/d$ ; the zero predictor scores 1
$\hat w=X^+y$: minimum-norm least squares; the posterior mean under $w\sim N(0,I)$ , noiseless
$(d-k)_+/d$: least-squares risk on in-distribution noiseless prompts
Averaging estimator: $\hat w=\frac1k\sum y_ix_i$ , one plain gradient step; risk $(d+1)/k$
$P_Xw$: projection of $w$ on the span of the examples; the right answer for $k\lt d$
Curriculum: start in a 5-dim subspace with 11 examples, grow to $d$ and $2d+1$
$n_p,\ n_w$: distinct training prompts or weight vectors; 1M or 10k nearly suffice
Oracle lookup error: best training $w$ for the test task: 0.216 with 32M, 0.505 with 10k

Quiz

Why is minimum-norm least squares the right yardstick for the model trained on noiseless prompts with $w\sim N(0,I_d)$ ?
With $N(0,1)$ label noise added at test time, the model's error spikes at $k=d$ like least squares. What does that say about the learned algorithm?
The model trained on two-layer ReLU networks learns linear functions in context at the rate of a network trained by Adam, which is slower than the least-squares rate. What does this suggest?

Sources

Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes. NeurIPS, 2022.