A Circuit for In-Context Learning: Transformers Implement Gated Associative Lookup

Forum topic · 小凯 · 2026-11-04

Summary

This arXiv paper (2611.05344) by Andrew Gao and William Gao reverse-engineers the mechanism behind in-context learning (ICL) in three 3-4B-parameter transformers performing few-shot sentence classification with arbitrary labels. The authors show the model implements a key-value memory over label-token positions: each label token's residual stream encodes label identity, middle-layer attention heads write the preceding demonstration's class as a key, and upper-layer heads match the query's class against stored keys and copy out the matching label. An early-layer signal detecting whether any label repeats gates this readout; when absent, readers attend to the prompt start and no label is favored, while separate heads suppress the most recently seen label. Each step is established causally via activation patching, head ablations against random-head baselines, and steering with the model's own signals. A softmax readout model with two constants per head transfers across prompt distributions. The circuit is conserved across the three models, emerges in transformers trained from scratch, transfers to semantic labels, persists from 0.6B to 8B in one model family, and survives instruction tuning. ICL in this setting is a gated associative lookup built from a few identifiable component types.

Paper Overview

Field: NLP Authors: Andrew Gao, William Gao Published: 2026-11-03 arXiv: 2611.05344

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.

Key Findings

Key-value memory over labels: each label token's residual stream encodes label identity; middle-layer heads write the preceding demonstration's class as the key.
Query-side readout: upper-layer heads match the query's class against stored keys and copy the corresponding label.
Gating signal: an early-layer "has any label repeated?" signal gates the readout; without repeats, readers attend to the prompt start and no label is favored.
Recency suppression: a separate set of heads suppresses the most recently seen label.
Causal evidence: activation patching, head ablations against random-head baselines, and steering with the model's own signals.
Transferability: a softmax readout model with two constants per head, fit on one prompt distribution, predicts another; the circuit transfers to semantic concepts and meaningful labels.
Robustness: the circuit is conserved across the three models, appears in transformers trained from scratch, persists from 0.6B to 8B in one family, and survives instruction tuning.

Conclusion

ICL in this setting is a gated associative lookup built from a few identifiable component types, traced causally from demonstrations to prediction.

---

*Auto-collected on 2026-11-05.*