PaperlayerView exact revision on arXiv

Source-linked AI summary

What learning algorithm is in-context learning? Investigations with linear models

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, Denny Zhou

arXiv:2211.15661v3cs.LGcs.CL

TL;DR

The paper asks whether transformer in-context learning implicitly implements standard learning algorithms, rather than merely producing predictions from example sequences. It studies linear regression through constructive transformer analyses, behavioral comparisons, and probes of intermediate representations. The evidence shows implementable and empirically matched gradient-descent, ridge, least-squares, and Bayesian predictors, while the scope remains centered on linear problems.

  • Problem

    The paper asks which learning algorithms transformers can implement and actually discover during in-context learning, a question left unclear by prior work focused mainly on what functions ICL learns.

  • Method

    The authors study transformer ICL on linear regression using constructive proofs, comparisons with standard predictors, and probes of algorithmically meaningful intermediate quantities.

  • Results

    Trained learners closely match gradient descent, ridge regression, ordinary least squares, and Bayesian predictors under varying depth, width, and dataset noise.

  • Takeaways & Limitations

    The findings support understanding at least some in-context learning in algorithmic terms, including possible rediscovery of standard estimation algorithms.

  • Takeaways & Limitations

    The experiments focus on the linear case, although the authors note that extensions to richer function classes are possible.

Abstract

from arXiv · show

Neural sequence models, especially transformers, exhibit a remarkable capacity for in-context learning. They can construct new predictors from sequences of labeled examples $(x, f(x))$ presented in the input without further parameter updates. We investigate the hypothesis that transformer-based in-context learners implement standard learning algorithms implicitly, by encoding smaller models in their activations, and updating these implicit models as new examples appear in the context. Using linear regression as a prototypical problem, we offer three sources of evidence for this hypothesis. First, we prove by construction that transformers can implement learning algorithms for linear models based on gradient descent and closed-form ridge regression. Second, we show that trained in-context learners closely match the predictors computed by gradient descent, ridge regression, and exact least-squares regression, transitioning between different predictors as transformer depth and dataset noise vary, and converging to Bayesian estimators for large widths and depths. Third, we present preliminary evidence that in-context learners share algorithmic features with these predictors: learners' late layers non-linearly encode weight vectors and moment matrices. These results suggest that in-context learning is understandable in algorithmic terms, and that (at least in the linear case) learners may rediscover standard estimation algorithms. Code and reference implementations are released at https://github.com/ekinakyurek/google-research/blob/master/incontext.

1 INTRODUCTION

The paper asks whether transformer in-context learning implicitly implements recognizable learning algorithms, studying this question through linear regression. It combines constructive theory, behavioral comparisons, and probing of internal computations.

  • ICL maps labeled in-context pairs to predictions on new inputs without updating the model’s parameters.
  • The paper studies whether learners encode a context-dependent model in hidden activations and train it as examples arrive.
  • Transformers can implement one gradient-descent step using O(d) hidden size and constant depth.
  • Transformers can update a ridge-regression solution with one new observation using O(d2) hidden size and constant depth.
  • Trained learners match several linear predictors, transition with depth and noise, and behave like Bayesian predictors at large hidden sizes and depths.
  • The results suggest that transformers can discover familiar algorithms from sequence-modeling tasks alone.

2 PRELIMINARIES

The paper frames ICL as implicit meta-learning: transformers learn to predict from example sequences without an explicit inner procedure. Linear regression provides a controlled setting for comparing possible algorithms.

  • Meta-learning traditionally trains an outer model to set parameters for an inner learning procedure.
  • ICL instead trains a sequence model to map example sequences to predictions, with the learning procedure implicit in its parameters.
  • Prior work established what functions ICL can learn, but left which learning algorithms transformers discover largely unclear.
  • A transformer maps input-vector sequences through layers whose capacity depends mainly on depth, hidden size, and number of attention heads.
  • The authors directly train transformers on an autoregressive ICL objective defined by a function class, input distribution, function distribution, and loss.
  • The experiments use linear functions f(x) = w⊤x with squared-error loss, comparing a model problem that admits multiple algorithmic solutions.

3 WHAT LEARNING ALGORITHMS CAN A TRANSFORMER IMPLEMENT?

The theory constructs transformer parameterizations that implement standard linear-learning updates. These constructions use simple layer-level operations to realize gradient descent and iterative closed-form ridge regression.

  • A transformer decoder layer can implement move, matrix multiplication, division, and affine operations on its hidden-state matrix.
  • These primitives provide building blocks for implementing gradient descent and closed-form ridge-regression algorithms.
  • 3.2 GRADIENT DESCENT: A single gradient-descent update repeatedly adjusts parameters using each labeled example before predicting on a new input.
  • 3.2 GRADIENT DESCENT: A transformer implements one gradient-descent prediction with constant depth and O(d) hidden space.
  • 3.2 GRADIENT DESCENT: The constructed output contains w′⊤x_n at the position corresponding to the new input x_n.
  • 3.3 CLOSED-FORM REGRESSION: Sherman–Morrison reduces the inverse covariance calculation to rank-one updates, enabling iterative closed-form ridge-regression computation.
  • 3.3 CLOSED-FORM REGRESSION: A transformer implements one Sherman–Morrison update with constant layers and O(d2) hidden space.

4 WHAT COMPUTATION DOES AN IN-CONTEXT LEARNER PERFORM?

The paper compares trained in-context learners with standard linear-regression algorithms using behavioral agreement metrics, then examines uncertainty and computational constraints. ICL matches OLS on noiseless problems, Bayesian-optimal ridge under noise, and shifts among gradient descent, ridge, and OLS regimes as capacity changes.

  • Behavioral metrics: The analysis compares ICL predictions with gradient descent, stochastic gradient descent, ridge regression, OLS, and nearest-neighbor predictors using SPD and ILWD.SPD measures output agreement, while ILWD compares implied linear weights or their closest linear approximations.
  • Noiseless datasets: On noiseless linear regression, ICL closely matches OLS, with normalized SPD and ILWD corresponding to squared error below 0.01.Other predictors, especially nearest neighbors, agree substantially less well.
  • Noiseless datasets: In underdetermined problems, ICL predictions consistently agree with OLS’s minimum-norm solution among multiple models that fit the context exactly.This behavior is observed when the number of examples is smaller than d = 8.
  • Uncertainty: For noisy datasets, ICL matches the minimum-Bayes-risk ridge predictor across prior variances and noise levels.The best-fitting ridge parameter increases with prior variance and equals the parameter minimizing Bayes risk.
  • Uncertainty: As dataset noise approaches zero, the Bayesian predictor converges to OLS, connecting the noisy and noiseless findings.The Bayesian estimator has the same form as ridge regression with regularization parameter σ2/τ2.
  • Computational constraints: As model depth increases, ICL transitions from gradient descent in 1-layer models to ridge regression at 2–4 layers and OLS at 8 or more layers.Comparable phase shifts occur with hidden size; ridge-like behavior appears around H ≥ 16 for d = 8 and H ≥ 32 for d = 16.

5 DOES ICL ENCODE MEANINGFUL INTERMEDIATE QUANTITIES?

The paper probes whether trained in-context learners encode intermediate quantities associated with standard regression algorithms. Probes recover moment and least-squares quantities from hidden representations, with distinct nonlinear encoding and layerwise timing.

  • Probe design: Probes target the moment vector X^T Y and least-square solution wOLS after each prefix of in-context examples.The learner’s hidden states are frozen while auxiliary probes attempt to recover these quantities.
  • Encoding structure: A 2-layer MLP probe outperforms a linear probe for both targets, indicating nonlinear encoding in the learner representations.Probe success is inaccurate at the input layer but becomes accurate deep in the network.
  • Layerwise computation: X^T Y becomes predictable around layer 7, while w becomes predictable later, around layer 12.This ordering resembles the intermediate computations expected from the paper’s regression constructions.
  • Temporal localization: Probes attend to the correct timestamps when decoding targets, primarily using the current example’s representation.For prefixes of length n, the target is decoded mainly from y_n; after n = d examples, decoding is distributed across y_n≥4.
  • Control comparison: Control-task probes perform significantly worse at recovering moment matrices than probes applied to models requiring in-context learning.The control task uses a fixed weight vector w = 1 and does not require ICL.

6 CONCLUSION

The paper characterizes linear in-context learning through theory, behavioral comparisons, and probing of intermediate quantities. It concludes that some ICL can be understood as familiar learning algorithms discovered by transformers, while noting that the experiments focus on linear functions.

  • Main conclusion: The experiments show that transformers can implement multiple linear regression algorithms in theory, empirically implement a range of them, and expose their intermediate quantities to probes.The empirical behavior transitions with model capacity and dataset noise.
  • Scope: The study’s experiments focus on the linear case, although the methodology can extend to richer function classes with nonlinear feature computation.The authors present this extension as a possibility rather than as an evaluated result.
  • Implication: The results provide initial evidence that apparently mysterious ICL can be analyzed with standard machine-learning tools and may rediscover solutions developed by machine-learning researchers.The broader extension to larger-scale ICL, including language models, remains future work.

A THEOREM 1

Theorem 1 constructs a transformer implementation of a gradient-descent update for linear regression. The construction computes prediction errors, forms an update, writes the new weight vector, and can repeat across examples.

  • Update construction: The constructed operations read inputs, compute w^T x and its residual, multiply by x, and write the updated weight vector.The operator chain includes affine transformations, multiplication, movement, and writing steps.
  • Output: The construction produces the updated prediction in the last hidden unit of the third time-step.The prediction is obtained after applying the update to a subsequent input.
  • Multiple steps: Repeating the update across examples requires O(n) layers for one pass through the dataset.The procedure reuses the written weight vector for subsequent iterations.

B THEOREM 2

Theorem 2’s construction uses transformer primitives to implement the matrix operations needed for closed-form ridge regression. The RAW operator provides a common read–arithmetic–write mechanism for these computations.

  • Theorem 2: The ridge-regression construction requires O(d2) hidden space because matrix multiplications must be represented explicitly.Generalizing the procedure to multiple iterations requires at least O(n) layers.
  • RAW operator: The RAW operator reads hidden units with attention, performs elementwise arithmetic, optionally reduces the result, and writes values into the current hidden state.Its parameters specify arithmetic, index sets, output weights, and a causal timestep map.
  • Primitive decomposition: The mov, aff, and dot-product primitives can each be defined using a single RAW operator.The construction treats matrix multiplication and division separately.
  • Transformer implementation: A single transformer layer can implement RAW, using embeddings, positional information, attention, and feed-forward computations.The proof constructs scratch space and parameter settings for the required operations.
  • Parallelization: The construction can be extended from vector-vector dot products to general matrix-matrix products required by mul.This extension is obtained by parallelizing the relevant dot products.

D DETAILS OF TRANSFORMER ARHITECTURE AND TRAINING

The experiments use Jax on P100 GPUs with Adam, cosine learning-rate scheduling, warmup, and learned absolute position embeddings; Table 1 lists the ICL hyperparameters.

  • Experiments use the Jax framework on P100 GPUs with Adam and a cosine learning-rate scheduler with warmup.Warmup steps equal one-fifth of total iterations.
  • The transformer uses learned absolute position embeddings.
  • Table 1 reports the hyperparameters used for in-context learning and highlights the best setting for each hyperparameter.
  • Phase-shift plots hold the x-axis value constant while selecting the best combination of layers, hidden size, heads, and learning rate.

E DETAILS OF PROBE

The probe extracts information from hidden states using learned position attention and compares actual-task probing with a control task to test whether the measured accuracy is non-trivial.

  • Probe architecture: The probe and task model are distinguished explicitly in the probing setup.
  • Probe architecture: The probe uses learned position scores to form attention weights over sequence positions for each target variable.It applies a linear projection before averaging hidden states.
  • Probe diagnostics: Figure 6 reports R2 for linear weight estimation on the d = 8 problem.
  • Control experiment: The control task fixes w to the all-ones vector and does not require in-context learning.During probing, the control model receives data generated with normally sampled w.
  • Control experiment: The control probe has a significantly higher error rate than probing the actual task model, indicating that the actual probing accuracy is non-trivial.Detailed control-probe errors are reported in Figure 5.

F LINEARITY OF ICL

The paper evaluates how linearly ICL represents implicit weights and finds that its behavior becomes gradually linear in the underdetermined regime despite matching the min-norm least-squares solution under other measures.

  • Implicit weight comparison: The implicit linear weight difference measure compares ICL with linear algorithms without assuming that the predictors themselves are linear.For non-linear predictors, it measures differences between their closest linear predictors.
  • Linearity across regimes: ICL becomes gradually linear in the underdetermined regime, as measured by R2 of the estimator.The analysis considers both underdetermined and determined linear problems.
  • Linearity across regimes: Although ICL matches the min-norm least-squares solution in both metrics, its hypothesis class is not purely linear.
  • Nonlinear computation: Figure 7 includes approximations of multiplication using different nonlinearities.

G MULTIPLICATIVE INTERACTIONS WITH OTHER NON-LINEARITIES

The paper constructs multiplication approximations from nonlinear activation functions, using Taylor expansions for smooth functions and piecewise-linear constructions for ReLU, with differing approximation complexity.

  • Smooth nonlinearities: For a smooth nonlinearity, Taylor expansion provides the basis for approximating multiplication from function evaluations.The expansion converges within a sufficiently small neighborhood.
  • Smooth nonlinearities: The identity 1/(2a2)(f(x + y) − f(x) − f(y) + a0) = xy + O(x3 + y3) approximates multiplication when the second-order term is usable.
  • Smooth nonlinearities: GeLU can instantiate the generic multiplication approximation, which is plotted against x2 over the range [−0.1, 0.1].
  • Smooth nonlinearities: When the second-order coefficient is zero or negligible, numerical derivatives can use a third-order term to obtain a multiplication approximation.
  • Alternative smooth nonlinearities: Tanh requires six terms instead of three GeLU terms in the construction, increasing the hidden-space requirement.
  • Non-smooth nonlinearities: ReLU yields a piecewise-linear approximation of x2 using shifted ReLU terms, plotted over [−0.1, 0.1].

H EMPIRICAL SCALING ANALYSIS WITH DIMENSIONALITY

The analysis sweeps transformer dimensions and optimization settings to measure when in-context learning fits least-squares OLS better than Ridge regression. A single attention head suffices across tested input dimensions, while layer and hidden-size requirements vary stepwise with input size; some hyperparameters were not independently optimized.

  • The study tests input dimensions d = {1, 2, 4, 8, 12, 16, 20}.
  • The sweep varies layers, hidden size, heads, and learning rate while optimizing the remaining parameters.The tested ranges are L = {1, 2, 4, 8, 12, 16}, H = {16, 32, 64, 256, 512, 1024}, M = {1, 2, 4, 8}, and learning rates {1e-4, 2.5e-4}.
  • A single head is sufficient for all tested problem dimensions, whereas other capacity parameters show step-function-like dependence on input size.
  • The threshold is defined by when SPD(OLS, ICL) becomes lower than SPD(Ridge(λ = ϵ), ICL), with ϵ = 0.1 in the plotted experiment.
  • Hyperparameters such as weight initialization were not optimized independently for each input dimension.
  • The empirical requirements compare SPD(Ridge(λ = 0.1), ICL) against SPD(OLS, ICL) while other parameters are optimized.
Loading 2211.15661v3…