Transformer models use memory circuits to learn from prompts
A Circuit for In-Context Learning
Computation and LanguageMachine Learning
Summary
Some language models can learn to do new tasks just by seeing examples in their prompt, without changing their internal settings. This paper explains how certain parts inside these models work together to remember labels and match new sentences to the right category in just one pass. The authors studied several models and found a small circuit inside them that acts like a memory, storing information from examples and using it to decide what label fits best. This discovery helps us understand how such models learn on the fly from the examples they see.
What this means in practice
- •For machine learning engineers: Design language models that can rapidly adapt to new classification tasks without retraining by using memory-based circuits identified in this work.
- •For natural language processing developers: Create efficient prompt-based classifiers that leverage the internal associative lookup mechanisms described for improved few-shot learning performance.
Authors
Andrew Gao, William Gao
Abstract
In-context learning (ICL) lets a pretrained language model perform a task specified by demonstrations in its prompt, with no parameter updates. With arbitrary labels the model must acquire a mapping it cannot have memorized, so the learning happens within a single forward pass. Prior work found candidate parts but not how they compose. We give the composition for few-shot sentence classification with arbitrary labels in three 3-4B-parameter transformers. The model implements a key-value memory over label-token positions: each label token's residual stream encodes the label's identity, and middle-layer attention heads write into it the class of the preceding demonstration as its key. At the query, upper-layer heads match the query's class against the stored keys and copy out the matching labels. An early-layer signal encoding whether any label has repeated gates this readout; in its absence the readers attend to the start of the prompt and no label is favored. A separate set of heads suppresses the most recently seen label. Each step is established causally by activation patching, head ablations against random-head baselines, and steering with the model's own signals. A softmax model of the readout with two constants per head, fit on one prompt distribution, predicts another. The computation is conserved across the three models, the named heads alone reproduce the behaviour, and a transformer trained from scratch grows the same parts. The circuit transfers to semantic concepts and to labels with meaning, and the same parts, with a causal gate, appear from 0.6B to 8B in one family and survive instruction tuning. ICL in this setting is a gated associative lookup built from a few identifiable component types, traced causally from demonstrations to prediction.