Circuit for In-Context Learning in Language Models
Original title: A Circuit for In-Context Learning
The model implements a gated associative lookup using key-value memory over label-token positions and attention heads for few-shot sentence classification.
The paper traces a circuit for in-context learning in three 3-4B-parameter transformers performing classification with arbitrary labels. Middle-layer attention heads store class keys in label-token residual streams while upper-layer heads match and copy labels at the query. An early-layer signal gates the readout and separate heads suppress the most recently seen label. Activation patching and head ablations establish the causal computation which transfers across models.
- Key-value memory encodes label identity in residual stream.
- Middle-layer heads write class as key.
- Upper-layer heads match query class and copy label.
- Early-layer signal gates readout; absence leads to attending to prompt start.
- Computation conserved from 0.6B to 8B parameters and survives instruction tuning.
From the original article
1 min read at arxiv.orgIn-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.
By Andrew Gao · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.