Transformers learn in-context by gradient descent

Johannes von Oswald; Eyvind Niklasson; Ettore Randazzo; João Sacramento; Alexander Mordvintsev; Andrey Zhmoginov; Max Vladymyrov

Source-linked AI summary

Transformers learn in-context by gradient descent

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, Max Vladymyrov

arXiv:2212.07677v2Dec 15, 2022cs.LGcs.AIcs.CL

TL;DR

The paper asks how Transformers implement in-context learning, whose mechanisms remain poorly understood. It constructs and tests a gradient-descent interpretation, finding that trained Transformers can implement or closely match gradient-based learning in regression settings, while emphasizing the scope limits of this evidence.

Problem
The mechanisms underlying Transformer in-context learning remain poorly understood despite its importance for flexible prediction from contextual data.
Method
The paper constructs linear self-attention weights equivalent to gradient descent and compares trained self-attention Transformers with gradient-descent solutions on regression tasks.
Results
The experiments show that optimized Transformers either match the proposed construction or closely align with gradient-descent models, including for in- and out-of-distribution validation tasks.
Takeaways & Limitations
Within the studied regression settings, Transformers can be mechanistically understood as learning algorithms that perform gradient-based adaptation during the forward pass.
Takeaways & Limitations
The evidence is restricted to small Transformers and simple regression problems, and the authors expect the findings to explain only a limited part of more complex in-context learning.

Abstract

from arXiv · show

At present, the mechanisms of in-context learning in Transformers are not well understood and remain mostly an intuition. In this paper, we suggest that training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We start by providing a simple weight construction that shows the equivalence of data transformations induced by 1) a single linear self-attention layer and by 2) gradient-descent (GD) on a regression loss. Motivated by that construction, we show empirically that when training self-attention-only Transformers on simple regression tasks either the models learned by GD and Transformers show great similarity or, remarkably, the weights found by optimization match the construction. Thus we show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass. This allows us, at least in the domain of regression problems, to mechanistically understand the inner workings of in-context learning in optimized Transformers. Building on this insight, we furthermore identify how Transformers surpass the performance of plain gradient descent by learning an iterative curvature correction and learn linear models on deep data representations to solve non-linear regression tasks. Finally, we discuss intriguing parallels to a mechanism identified to be crucial for in-context learning termed induction-head (Olsson et al., 2022) and show how it could be understood as a specific case of in-context learning by gradient descent learning within Transformers. Code to reproduce the experiments can be found at https://github.com/google-research/self-organising-systems/tree/master/transformers_learn_icl_by_gd .

1. Introduction

The paper addresses the poorly understood mechanisms of Transformer in-context learning by hypothesizing that Transformers learn gradient-based algorithms within their forward pass. It provides constructions and empirical evidence linking in-context learning to gradient descent and meta-learning, while limiting the claim primarily to simple regression settings.

The authors construct explicit self-attention weights that implement one gradient-descent step and show that multiple layers can perform iterative curvature correction.
On linear regression tasks, optimized self-attention-only Transformers either converge to the proposed construction or closely align with gradient-descent solutions in in- and out-of-distribution validation.
Adding MLPs enables Transformers to solve nonlinear regression by learning linear models on deep representations, with solutions strikingly similar to meta-learned models and output-layer gradient descent.
The paper argues that Transformers can meta-learn learning algorithms on two time scales, making them mesa-optimizers during their forward computations.
The proposed mechanism explains only a limited part of in-context learning because Transformer architectures, domains, and phase transitions vary substantially.
The paper hypothesizes that Transformer in-context learning is implemented by gradient-based optimization of an implicit loss constructed from contextual data.

2. Linear self-attention can emulate gradient descent on a linear regression task

This section constructs a linear self-attention layer whose token transformation is equivalent to one gradient-descent step on linear regression. It then identifies learned task-shared and data-dependent transformations that can improve or accelerate this gradient-descent-like computation.

A linear regression model minimizes squared error, and one gradient-descent step changes its weight matrix according to the dataset inputs and residuals.
The resulting weight update can be re-expressed as a transformation of regression targets rather than an explicit update to model weights.
Context tokens represent training pairs and a query token represents the test input, allowing self-attention to update training and test tokens simultaneously.
A single-head linear self-attention layer can be constructed so that its query prediction matches the prediction after one gradient-descent update.
Experiments find that learned self-attention weights match the construction after accounting for equivalent rescalings of matrix products.
Training can learn task-shared learning rates and task-specific data transformations that exploit current samples and substantially speed up gradient-descent-like learning.

3. Trained Transformers do mimic gradient descent on linear regression tasks

The paper finds that trained self-attention-only Transformers closely implement gradient-based updates for linear regression, with deeper models learning an iterative curvature correction and MLP-augmented models extending this mechanism to nonlinear tasks.

Single-layer linear regression: A trained single linear self-attention layer generates predictions and losses that agree closely with one gradient-descent step across varied hyperparameters.The comparison includes model alignment, predictions, interpolation with constructed weights, and out-of-distribution inputs.
Repeated updates: Repeated application of a trained self-attention layer produces the same loss decrease as repeated gradient descent after correcting both algorithms’ learning rates.The experiment uses a dampening parameter λ = 0.75.
Single-layer linear regression: Training self-attention-only Transformers on linear regression produces models that closely align with gradient-descent solutions on in- and out-of-distribution tasks.The optimized weights can approximately coincide with the explicit gradient-descent construction.
Multiple steps of gradient descent vs. multiple layers of self-attention: Multiple self-attention layers generally outperform plain gradient descent and are better described by GD++, which iteratively transforms inputs using H(X) = (I −γXXT ).The trained Transformer realigns with GD++ and matches its performance on in- and out-of-distribution tasks.
Nonlinear regression: With MLPs, Transformers solve nonlinear regression by learning linear models on deep representations, including kernelized least-squares problems induced by the MLP.Experiments on sine-wave regression found similar functions and solutions for trained Transformers and meta-learned MLPs with gradient-updated output layers.

4. Do self-attention layers build regression tasks?

The paper shows that Transformers can overcome nonstandard token construction by using an initial attention layer to merge input and target tokens, enabling later in-context gradient descent.

Token construction: A construction using alternating input-only and target-only tokens can transform them into the concatenated tokens required for the gradient-descent mechanism.The construction uses key, query, value, and projection matrices in a one-head attention layer.
Learned copying: In trained two-layer attention circuits, the first layer becomes notably dependent on the neighboring token immediately before performance reaches the gradient-descent level.The layer remains independent of the other tokens, supporting a local copying operation.
Training behavior: A two-layer Transformer’s trained loss matches one gradient-descent step rather than two, while training takes an order of magnitude longer.The figure also tracks the first layer’s sensitivity to input tokens before the performance jump.
Learned copying: The authors interpret this neighbor dependence as copying that merges input and output data into single tokens for the second layer’s gradient-descent step.The first layer required softmax attention in the reported training experiments.
Mechanistic interpretation: Copying through softmax attention enables Transformers to combine data from different tokens and compute downstream input-target dot products for gradient-based in-context learning.This is identified as a second crucial mechanism alongside the later gradient-descent computation.

5. Discussion

The paper argues that Transformer in-context learning can arise from gradient-based optimization implemented within the forward pass. Its evidence connects this mechanism to copying, meta-learning, and nonlinear regression while identifying important scope limits.

Transformers may implement in-context learning by gradient-based optimization of an implicit loss constructed from in-context data.
MLPs allow Transformers to learn linear models on deep representations, supporting nonlinear regression and producing solutions similar to meta-learned MLPs adapted by one GD step.
A standard linear self-attention construction implements gradient descent by pairing neighboring sequence elements as input-target examples.
The authors restrict their strongest mechanistic interpretation to simple regression settings and identify noisy data, regularization, and logistic regression as areas needing further work.

A.1. Proposition 1

Proposition 1 constructs a single linear self-attention layer whose token transformation is identical to one gradient-descent update on regression data. The construction also applies to the test token, producing the corresponding updated prediction.

The attention formulation connects outer-product token interactions to one gradient-descent step on a mean squared error loss.
Proposition 1 constructs key, query, value, and projection matrices so one linear-attention step equals a gradient-induced update for every token.
The construction represents each token as an input-target pair and updates the target component while preserving the input component.
The same dynamics apply to the query token, whose zero-initialized prediction is transformed using the learned update.

A.2. Comparing the out-of-distribution behavior of trained Transformers and GD

Out-of-distribution experiments compare trained Transformers with GD and GD++ under rescaled and unseen input distributions. Trained models generally track these optimization procedures, with weaker matches for deeper non-recurrent models far from training conditions.

The out-of-distribution tests vary input ranges, teacher-weight scales, and input families including normal, exponential, and Laplace distributions.
Across rescaled training distributions, trained Transformers behave remarkably similarly to GD or GD++.
On distributions unseen during training, recurrent Transformers again closely match GD or GD++, while deeper non-recurrent models match less well far from training regimes.
Deep-model evaluations use clipping between [-10, 10], which restricts loss increases and produces plateaus away from training conditions.

A.3. Linear mode connectivity between the weight construction of Prop 1 and trained Transformers

Weight interpolation experiments test whether trained Transformers occupy a connected path to the explicit GD construction. Recurrent models align closely with GD++ and interpolate well, while non-recurrent models require more elaborate scaling corrections.

Interpolation rescales trained weight products using the mean diagonal of W_KQ before averaging them with the GD products.
Ten-step GD++ and trained recurrent ten-layer Transformers align closely and remain interpolatable inside and outside training situations.
Twelve-layer Transformers with MLPs and four-headed linear self-attention closely resemble twelve GD++ steps despite independent blocks and omitted LayerNorm.
Weight-matrix products of trained one-layer and three-layer recurrent Transformers match the Proposition 1 construction after scalar correction.
Non-recurrent Transformers have cross-layer scaling ambiguities that prevent reliable interpolation with the simple correction for some random seeds.

A.4. Visualizing the trained Transformer weights

The authors visualize trained Transformer weight products against the gradient-descent construction, accounting for redundancy in how the products can be realized.

The trained products W_KQ = W_T^K W_Q and W_PV = P W_V are visualized in Figure 9.The comparison is performed in weight-product space rather than directly on individual parameters.

A.5. Proof and discussion of Proposition 3

The construction uses positional encodings and attention projections to transform paired input-target tokens into the representation required by the gradient-descent construction.

The construction applies to one-head linear or softmax attention with appropriately chosen key, query, value, and projection matrices.These matrices transform all tokens into equivalents of those required by Proposition 1.
The method concatenates unit positional encodings to input tokens before applying attention projections.Tokens are represented as e_j = (x_j/2, p_j), with p_j a positional encoding.
A suitable key-query product selects the next positional entry, allowing the value projection to copy the next target into the current token.The resulting token has the form e_j = (x_j/2, 0, y_j/2+1).
For toy regression, the construction provides N/2 input-output pairs, while language-like settings may use nonadjacent token pairs.The paper notes that every two tokens can represent an input-output pair for an autoregressive inner-loop loss.

Rolling out experiment with different dampening strength

Repeated application of the trained self-attention update is evaluated under different dampening strengths, with performance compared against repeated gradient descent.

At λ = 0.75, gradient descent and the trained self-attention layer show matching performance for 50 rollout steps.The figure reports slightly better robustness for gradient descent across dampening strengths.

A.6. Dampening the self-attention layer

The experiments examine repeated updates, nonlinear regression, softmax and LayerNorm variants, curvature correction, and optimization stability. They show both similarities to gradient descent and architecture- and setting-dependent limitations.

Rollout and dampening: Repeated updates can diverge quickly, but scaling the update by λ yields similar out-of-distribution behavior for gradient descent and trained self-attention.At λ = 0.75, both methods steadily decrease loss within 50 steps.
Nonlinear regression: MLPs plus self-attention support nonlinear regression by learning a linear model on deep representations.The resulting dynamics are equivalent to one gradient-descent step on a kernelized squared-error loss.
Softmax attention: A single softmax self-attention layer does not match gradient-descent performance, whereas two heads substantially reduce the gap and improve alignment.The two-headed correction approximately cancels off-diagonal terms in weight space.
LayerNorm: LayerNorm experiments show qualitatively similar behavior, with lower performance and alignment than gradient descent, while possibly improving generalization through data normalization.The normalization interpretation is presented as speculation tied to larger in-context inputs and targets.
Curvature correction: Curvature correction improves conditioning by moving the eigenvalue distribution toward condition numbers near 1, but larger γ values quickly produce exploding condition numbers.The simple derivation does not fully capture task statistics or distribution shifts.
Optimization behavior: Two-headed self-attention mitigates random-seed-dependent training instabilities observed with a single head.The paper also reports grokking-like train and test loss transitions under selected training settings.
Optimization behavior: Without regularization, deeper Transformers with K > 2 can become unstable when approaching gradient-descent performance.The authors speculate that this reflects gradient descent already operating near divergence on the training tasks.
Optimization behavior: On a fixed batch of training tasks, meta-gradient descent still finds Transformer weights that align remarkably well with the proposed construction and gradient descent.