II · THE IDEA · ARTIFICIAL INTELLIGENCE
Mechanistic Interpretability of In-Context Learning
▶ Listen · narrated
In-context learning looks like training without training. Mechanistic work on simplified models asks whether the network is literally implementing gradient steps between the prompt examples and the query.
At a glance
- Core question
- Which learning algorithm, if any, does in-context learning implement?
- Testbed
- Transformers on linear-regression prompts
- Candidates
- Gradient descent and closed-form linear estimators
- Method
- Hand-built constructions plus trained-model probes
- Limit
- Strongest claims stay inside linear problems
Think of a calculator that cannot change its stored programs, yet can still solve a new mini-problem if you type the worked examples first and the question last. Nothing inside the calculator is permanently retrained; it only shuffles the numbers you just typed through fixed circuitry until an answer appears.
In-context learning in a transformer is like that. The model’s weights stay frozen. The examples in the prompt act as a temporary training set. Mechanistic work asks what fixed circuitry is doing with those examples. For the simplified case of straight-line (linear) rules, ordinary machine learning already knows good procedures: take a few gradient steps toward better line-fit parameters, or solve for the best line in one closed-form shot. Researchers both hand-build small transformers that perform those procedures layer by layer and train transformers to see whether they discover something similar. When the match is good, “the model learns from the prompt” is no longer only a figure of speech — it is a description of arithmetic happening in the activations. When the match is poor, the model must be doing something else. The linear case is a clean test bench; language is messier and still less settled.
The object of study is a transformer that, given a prompt of pairs (x_1, y_1), …, (x_n, y_n) and a query x_q, predicts y_q under the assumption that y ≈ w⊤x for some w drawn per episode. Parameters θ of the transformer are fixed at inference; any adaptation to the episode must live in activations.
Two candidate algorithms dominate the comparison. Iterative: initialise a predictor ŵ and apply gradient updates on the empirical squared loss (optionally with ridge regularisation), using the in-context pairs as the batch. Closed-form: compute a regularised normal-equation solution from the same pairs. Constructive results exhibit transformer weights whose residual stream maintains ŵ (or the sufficient statistics X⊤X and X⊤y) and whose blocks implement one update step or one fragment of the closed form, subject to width and depth scaling with input dimension, step count, and numerical precision.
Trained models are then measured against these references on held-out covariances and noise levels: output MSE curves versus step index, implicit regularisation strength, and — where probes allow — alignment of intermediate activations with ŵ_t or with the running Gram matrix. Agreement supports the claim that next-token training on a distribution of linear episodes can induce an in-forward-pass optimiser. Disagreement outside the trained regime, or when n exceeds what width can store, bounds how literally the identification should be read.
Transfer remains the main limitation. These identifications are for linear regression under known data-generating processes. They do not by themselves establish that few-shot language behaviour, or gains from instruction-data selection at tuning time, are the same algorithm. They do supply a proof template: name the candidate procedure, construct it, then test whether gradient descent on the pretraining objective finds a neighbouring implementation.
Look closer
The prompt is the training set
In the linear setup, each in-context example is a pair the model must treat as data, not as text to copy. The query arrives last. Whatever algorithm the network runs has to accumulate evidence from those pairs inside the residual stream and attention patterns, then apply the resulting predictor to the query — all without changing any weights.
Construction before observation
One strand of the work builds transformers by hand whose layers are wired to perform the arithmetic of gradient descent or of a closed-form linear estimator. If such constructions exist and match the depth and width the algorithm needs, the architecture is at least capable of hosting those procedures. Capability is not proof of what trained models do, but it narrows the search.
Trained models under the same probe
Separately, transformers trained on distributions of linear-regression prompts are compared with the hand-built algorithms. Behavioural match on held-out tasks, and intermediate-state match where it can be measured, are used to argue that training can discover something close to gradient-based updates rather than an unrelated heuristic. Agreement is strongest in restricted regimes; outside them the picture is less settled.
The story
In-context learning is the observation that a model, given a few examples in its prompt, can continue the pattern on a new input without any gradient update to its parameters. From the outside it resembles rapid learning. From the inside it is only a forward pass. Mechanistic interpretability asks what computation that pass is actually performing.
A productive way to make the question sharp is to strip the problem down until the ordinary answer is already known. Linear regression is such a strip-down. The model sees a sequence of input–output pairs drawn from a linear rule, then a query input, and must predict the corresponding output. Ordinary machine learning already supplies candidate algorithms: iterate gradient descent on a squared loss, or solve a ridge-style normal equation in closed form. If a transformer’s internals line up with one of those, the metaphor of “learning in context” becomes a claim about implemented arithmetic rather than a loose analogy.
The investigative pattern splits into two moves. First, constructive: show that there exist transformer weights whose layer-by-layer computation realises gradient-descent steps, or realises a closed-form linear estimator, on the in-context data. Attention and residual updates have to maintain running quantities — estimates of weights, accumulated gradients, sufficient statistics — and each block must implement one algorithmic step or one piece of the closed form. Where the construction succeeds, the architecture is expressively adequate. Depth and width must scale with the algorithm’s needs; a shallow net cannot host many descent steps.
Second, empirical: train transformers on random linear problems and test whether their behaviour, and where measurable their intermediate states, track the same algorithms. A model that matches gradient descent on novel covariances, step-size choices, or regularisation strengths is harder to dismiss as a superficial mimic. A model that instead matches the closed-form estimator suggests a different internal route to the same predictions. Disagreement between behaviour at the output and the story told by the internals is itself informative: the network may approximate the algorithm only on the distribution it was trained on.
Why linear models at all, when the public interest is in language? Because in free-form text there is no agreed target algorithm against which to score the forward pass. In the linear case the loss, the optimum, and the iterative updates are explicit, so mismatch can be measured rather than narrated. The cost of that clarity is scope. Results that hold for linear regression do not automatically transfer to in-context learning over natural language, instruction following, or the kinds of prompt collections studied when selecting instruction-tuning data. They supply a lower bound on what the architecture can do and a template for how one might reverse-engineer richer cases later.
Read this way, mechanistic work on in-context learning is less a single discovery than a research style: pick a task whose classical algorithm is known, prove the transformer can host it, then ask whether training finds something close. The linear investigations are an early, comparatively clean instance of that style. They do not close the question for large language models; they show what a satisfying answer would have to look like.
Why it mattered then
When in-context learning became hard to ignore in large language models, the dominant accounts were behavioural: the model “learns from the prompt,” or retrieves a latent task, or pattern-matches. Those descriptions organise experiments but do not say what the layers compute. Investigations with linear models mattered in that moment because they replaced the metaphor with candidates that could be written as equations — gradient descent, closed-form ridge-style solutions — and then asked whether transformers could implement them in a forward pass. Constructive proofs answered a prior objection: that attention blocks might be the wrong substrate for iterative learning algorithms. Empirical matches in trained miniature models answered a second objection: that only hand-built curiosities would ever do so. Together they shifted part of the field from describing in-context learning as a surprising capability toward treating it as an object for algorithmic reverse-engineering.
Why it matters now
The same questions still sit under much everyday use of large models. Few-shot prompts, tool demonstrations, and chain-of-thought exemplars all assume that examples in context change the model’s effective behaviour without weight updates. If that change is sometimes a real optimisation procedure compressed into activations, then prompt design is closer to choosing a training set and a step budget than to writing instructions for a person. If it is only sometimes that procedure, and only on distributions that resemble the model’s pretraining, then failures on shifted formats or thin languages are less mysterious. Mechanistic accounts also affect evaluation: accuracy on a few-shot benchmark does not reveal whether the network ran ten gradient steps, inverted a small matrix of statistics, or executed a brittle heuristic that collapses outside the test suite. For anyone inspecting open-weight models, the linear case remains a sandbox where probes can be checked against ground truth before they are trusted on messier tasks.
The surprising detail
The striking move is not that transformers can predict well on linear prompts — many architectures could. It is that a depth-limited stack of attention and multilayer-perceptron blocks can be wired, and in some regimes trained, to mirror an algorithm that machine learning usually runs as an explicit loop over parameters. Learning, in that picture, is relocated from the optimiser into the residual stream. The “training set” is tokens; the “parameter update” is a transient activation pattern that disappears when the prompt ends.
What is disputed
Evidence that trained transformers implement gradient descent or closed-form linear estimators is strongest inside controlled linear-regression settings. How far those algorithms explain in-context learning in natural language, or in instruction-tuned models selected for data quality, remains open; behavioural similarity alone does not identify the internal procedure.
Remember this
In-context learning can be studied as an algorithm in the forward pass; linear models make that claim testable, not merely metaphorical.
Test yourself
A transformer matches a closed-form ridge estimator on the linear problems it was trained on, but diverges when the number of in-context points exceeds the width that construction would need to store the sufficient statistics. What does that pattern suggest about reading its success as “it implements ridge regression”?
It suggests the match is partial and resource-bound rather than a full, general implementation. Behavioural agreement on the training regime is compatible with an approximation that reuses the same statistics only up to the network’s finite width and depth. Once the prompt demands more memory or more algorithmic steps than the construction allows, divergence is expected. The honest claim is therefore narrower: the model implements something ridge-like within a capacity envelope, not ridge regression as an unrestricted procedure. That is still mechanistic progress, but it blocks the leap from “matches on these prompts” to “is the algorithm.”
Go deeper
- [2211.15661] What learning algorithm is in-context learning? Investigations with linear models · arxiv.org
- [2307.06290] Instruction Mining: Instruction Data Selection for Tuning Large Language Models · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.