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.
Figures & tables
Figure 1: One twelve-demonstration prompt, stage by stage. Columns are tokens (sentences si , arrows, labels and the query q ) and rows are layers of Qwen3-4B. Each coloured block marks the tokens and layers where a stage acts: the duplicate detectors and the encoder early on, the writers into the label tokens in the middle layers, and then, at the query arrow, the query code, the gate, the readers, the brakes and the readout band. Later sections give the intervention behind each block.
Term
Meaning
slot
the residual stream at a label token, which holds a key and a value
key
the class of the sentence the label followed; written into the slot by writer heads, read through the readers’ WK
value
the identity of the label; read through the readers’ WV
writer heads
about twenty middle-layer heads that copy a sentence’s class into its label token
reader heads
about ninety upper-layer heads that attend from the query arrow to slots whose key matches the query’s class, and copy their values
query code
the representation of the query’s class that the readers match against the keys
Table 1: Terms used throughout the paper.
Figure 2: The circuit for few-shot in-context learning in Qwen3-4B-Base. Prompt tokens run down the left (sentences Si , labels Li , the query Q and its arrow END). Boxes are classes of attention heads, with layer.head numbers beneath. Writer heads give each label token a key (the class of its sentence) alongside its value (the label). At END, reader heads match the query’s class against those keys and copy the matching labels, copier heads write the logits, and brake heads write against the most recent label. The gate g , assembled by the pooling heads, opens the read by lowering the readers’ attention-sink logit (dashed). Sections 3 – 8 give the intervention behind each arrow.
Figure 3: The same circuit laid out along the layers of Qwen3-4B-Base. The top lane shows what happens at the label tokens: duplicate-token heads mark a repeated label, and writers deposit the sentence’s class into the label’s residual stream, which makes it a slot with a value and a key. The bottom lane shows what happens at the query arrow: the encoder computes the query’s class, the pooled repeat signal g lowers the readers’ sink logit, the readers match the query’s class against the slot keys and copy the matching values, and the readout band and late MLPs write the logits while the brakes write against the last label.
Figure 4: The four main interventions, each drawn on the end of a 12-shot prompt. The red component is replaced or steered, and the output is read at END. (a) Sentence swap: the keys at a label token come from the twin in which that demonstration’s sentence belongs to the other class, so the slot’s key no longer matches its value. (b) Label swap: the values at the label token come from the twin with the labels exchanged. (c) Gate steering: the model’s own gate direction g is added at END at its natural magnitude, before the reader layers, and we measure the readers’ attention to the sink and to the matching slots. (d) Query-code handoff: the outputs of the two query heads at END come from the twin whose query is of the other class; later attention outputs can be held fixed to test whether the effect passes through them.
Figure 5: Anatomy of a slot. The residual stream at a label token holds two things, which are read through two different matrices. The label’s identity comes from the token’s own embedding and the early layers, and is read by the readers’ WV . The class of the preceding sentence is written in by middle-layer heads that attend back to the sentence, and is read by the readers’ WK as the key that the query’s class is matched against.
Figure 6: The slot holds a value and a key on different sides of attention. (a) Exchanging the labels: in all three models the effect runs through the values at the label tokens. (b) Changing a sentence while keeping its label: the effect runs through the keys at the label tokens, and only in the middle layers. (c) Handoffs from a same-class sentence of the same length move the answer by at most 0.06 nats. Bars are paired effects on the margin with 95% bootstrap intervals; grey bars are size-matched random controls.
Figure 7: Three key–value groups carry half of the label’s value. Patching the values of all 72 KV groups of layers 21–29 at the label tokens, from a prompt with the labels exchanged, carries 86% of the exchange’s effect on the margin (left). By layer, the effect sits in layers 22, 24 and 29 (right, grey), and within each of these layers one KV group carries most of it (red). Group 24.5 alone carries a quarter, and groups 24.5, 29.2 and 22.5 together carry about half. Random head sets of the same size move the margin by 0.03 or less (dotted). Qwen3-4B, 12-shot invented-label prompts, 72 prompts.
Figure 8: The key is written in the middle layers, and the write has to be spread out. (a) In Llama-3.2-3B, patching the label-token keys and values from the sentence-exchanged prompt one layer band at a time: the two middle bands carry the effect, and the first and last carry almost none. (b) In SmolLM3, adding a fixed total dose of the class direction u to the label tokens at one layer helps less than spreading it over four layers, and falls back at double dose while the spread write keeps climbing (95% intervals; 144 prompts).
Figure 9: Within a regime the read follows a softmax law, and the law transfers to fewer slots per class but not to longer prompts. Each point is one head at one slot count, showing the measured matching-slot mass against the prediction of Equation 1 from twelve-demonstration constants. (a) Qwen, in the fitted regime. (b) Qwen and Llama, one slot per class (held out). (c) Qwen and Llama, eighteen demonstrations (held out; the prediction fails). Sharp readers sit on the line and diffuse heads fall below it.
Figure 10: The gate lowers the null logit and moves matching far less. (a) Projection of the arrow residual onto g for prompts with repeated and with unique labels, intact and with the sixteen duplicate-token heads ablated, with random and induction-scored sets as controls. (b) Change in each reader’s sink-minus-match logit when g is added at natural and double dose; same-norm random directions in grey. (c) Change in the same readers’ match-minus-other logit.
Figure 11: What a reader looks at. Attention from the query arrow to the first position (the sink) and to the four label slots, for three Qwen3-4B readers, averaged over 30 four-demonstration prompts. Top row: two labels, each used twice; the matching slot (slot 3 or 4, depending on the prompt) draws the attention. Middle row: the same prompts with the sixteen repeat-detector heads mean-ablated. L29h11 moves back toward the sink (0.21 to 0.40), while L24h21 and L22h20 do not, so the detectors’ signal reaches some readers and not others. Bottom row: four distinct labels, none repeated; L29h11 puts 0.91 of its attention on the sink and reads nothing.
Figure 12: Two heads write a query code that the reader population consumes. (a) Share of the query-flip effect carried by the pair L22h10+L23h26 and by all attention in layers 21–29, and what remains when downstream attention queries are held at their intact values. (b) Energy of the pair’s write in the span of the readers’ match functionals, against random subspaces of the same dimension.
Figure 13: No single head carries the answer. Each cell is the effect on the 12-shot class margin of knocking out one head at the query position (mean over 144 prompts from six label families; blue lowers the margin and red raises it; colour clipped at ±0.15 nats), for every head in the reader layers of Qwen3-4B-Base (layers 21–35) and Llama-3.2-3B (layers 16–27). Outlined cells are the five most negative and five most positive heads. The largest single-head effects are −0.14 nats (Qwen L23h26, a query head) and +0.13 (Qwen L31h30, a brake), and −0.09 and +0.12 in Llama (L18h10, L19h1), against a 1.3-nat margin. The sets built from these rankings move the margin by −0.74 and +1.39 in Qwen and by −0.56 and +1.90 in Llama when knocked out together (Table B ).
Figure 14: Readers and brakes act as populations with opposite signs. (a) Joint knockout of 20 readers, of 24 brakes, and of both, on a fresh draw, against size-matched random sets (grey). (b) Order effect: the margin against the position of the last matching demonstration, intact and with the brakes knocked out.
Figure 15: Every nat of the answer assigned to a component. (a) Direct contributions to the margin by layer band, intact and under the class flip, in Qwen and SmolLM3. (b) Cumulative share of the flip against the number of heads, sorted by contribution.
Figure 16: The answer as a sum of bands. Left: the intact twelve-demonstration margin in Qwen, built up band by band; the copier band of layers 30–35 writes almost all of it. Right: the change under the exact class flip. The copier band carries most of the swing and the readers a quarter, and the late MLPs push the other way.
Intervention
Battery
SmolLM3
Llama
Class-flipped twin, everything
concept
−2.05 / −1.30
−1.92 / −0.86
twenty readers’ outputs from twin
concept
−2.16 / −1.52
−0.86 / −0.46
+ hold late attention inputs
concept
−1.38 / −1.24
−1.02 / −0.61
+ hold late attention and MLP inputs
concept
−0.79 / −0.53
−1.09 / −0.52
random sixteen heads from twin
concept
−0.11 / −0.07
−0.16 / +0.05
Twenty readers knocked / random twenty
slotgate
−1.32 / −0.05
−0.13 / +0.23
Table 2: The connection cuts in SmolLM3-3B and Llama-3.2-3B, on fresh draws. Change in answer margin (nats) from intact, mean over prompts; concept battery rare / symbol ( n=71 each), slotgate ( n=60 ). The hold rows fix the inputs of the late attention and MLP layers (SmolLM3 L30–35, Llama L22–27) at their intact values.
Intervention
Battery
Δ margin
Class-flipped twin, everything
concept
−1.80 / −0.95
twenty readers’ outputs from twin
concept
−1.08 / −0.54
+ hold attention inputs L30–35
concept
−0.81 / −0.35
+ hold attention and MLP inputs L30–35
concept
−0.39 / −0.19
Duplicate heads knocked
slotgate
+0.15 / −0.28
+ g at layer 23 ( ×1 )
slotgate
+0.58 / +0.04
Table 3: Four interventions on connections in Qwen3-4B-Base, on fresh draws. Change in answer margin (nats) from intact, mean over prompts. Concept battery: rare-word / symbol labels ( n=72 each); slotgate: unique / recurring label prompts ( n=30 each).
Link
Kind
Qwen3-4B
SmolLM3-3B
Llama-3.2-3B
Slot value (label) at label-token V
forced
yes
yes
yes
Slot key (class) at label-token K, mid layers
forced
yes
yes
yes
Writers, dose spread across depth
implementation
population
population
–
Read law within regime / down-transfer
forced
127/140, 0.07
7/8
124/125, 0.04
Read law up-transfer (18 demos)
tested
failed
–
failed
Duplicate detector feeds gate
implementation
named
yes
weak
Table 4: Each link of the account in three models. Forced : follows from the architecture and is expected in every model. Implementation : set by training and reported per model. Named : a small set of heads is necessary and sufficient. Population : no single head is necessary, and the link is characterized by a direction with a size-matched random control. –: not run.
Figure 17: The null action in two families. (a) Llama-3.2-3B emits a label on 67% of unique-label prompts and Qwen3-4B on 10%: the same fallback, gated much more weakly in Llama. (b) In Qwen, adding the gate signal g at its natural dose lowers the sink logit relative to the matching slots at the three strongest readers (red), and moves matching relative to the other slots much less (grey). (c) In Llama, knocking out the sixteen heads that a generic induction score selects moves the middle-layer readers’ attention out of the sink and onto the slots, so in Llama the induction-scored heads hold the read closed.
Figure 18: Closing the circuit on natural labels. For each model, the part of the margin that follows the demonstrations ( R , blue) and the part that follows what the label words mean ( P , orange), intact and with only the named heads kept at the query (final sets of 189, 122 and 163 heads, with the late MLPs). The reference is the mean over same-kind prompts or a resampled same-kind prompt. The controls are a random set of the same size and layers, and no heads. Bars are means over 96 natural/flipped prompt pairs; whiskers are 95% bootstrap intervals.
Figure 19: One census at four sizes (Qwen3 base, 0.6B, 1.7B, 4B and 8B). (a) Intact margins for invented labels, natural word labels and flipped word labels. (b) Reader and brake populations: the summed effect of the heads whose single removal at the query lowers (blue) or raises (red) the margin with an interval that excludes zero; the number is the head count. (c) The brakes’ summed hold on word meaning (bars, left axis) against the flipped-label margin (line, right axis). The two rise together.
Figure 20: Qwen3-4B base (a) and instruct (b) on the word-label prompts, with the base model’s 189 heads in both. Blue is the part of the margin that follows the demonstrations ( R ), and orange the part that follows the words’ meaning ( P ); the band is the intact R interval. Post-training multiplies both parts by about 2.5 and keeps their ratio. The 189 heads alone, with a resampled reference, keep 79% of R in the instruct model, and a random set of the same size keeps none.
0.6B
1.7B
4B
8B
Intact margin, symbol / rare labels (nats)
0.32 / 0.62
0.87 / 1.12
1.62 / 1.57
2.43 / 1.88
Natural / flipped word labels
2.18 / −0.03
3.30 / +0.14
5.40 / +0.35
6.74 / +1.05
R / P
1.07 / 1.10
1.72 / 1.58
2.87 / 2.52
3.90 / 2.85
Readers: count, summed effect
70, −2.08
48, −2.60
89, −4.69
103, −5.94
Largest reader, % of intact
22
18
18
15
Brakes: count, summed effect
13, +0.82
13, +1.21
56, +3.76
32, +3.27
Table 5: One census at four sizes (Qwen3 base, concept12n , 384 prompts each, upper-half heads removed one at a time at the query). Readers: removal lowers the margin, with an interval that excludes zero; brakes: removal raises it. R and P are the parts of the word-label margin that follow the demonstrations and the meaning. Gate dose: change in the unique-label margin when the size’s own fitted direction is added at the query arrow at the model’s own dose and at twice that dose, and change in the recurring-label margin when it is taken away. 0.6B and 1.7B screen layers 10–27 (288 heads); 4B and 8B screen layers 16–35 (640 heads).
In-context learning (ICL) enables language models to perform new tasks from demonstrations without weight updates. However, every ICL inference requires processing the full set of examples, resulting in inefficient deployments, and how ICL works mechanistically is not fully understood. Prior work compresses ICL into fixed activation vectors extracted from specific layers or positions, but these input-independent interventions fail on complex tasks where the output depends on fine-grained interactions with the input. By analyzing the ICL forward pass, we show that each attention head's output is an affine transformation of its context-masked counterpart, and that the parameters of this transformation are empirically stable across samples for a given task. Building on this, we introduce Task Operator (TO), which replays this transformation as an analytically derived update to the attention output projection. Across lexical, algorithmic, and reasoning tasks, TO achieves the best overall performance among prior methods and substantially narrows the gap between zero-shot inference and ICL. We further show that the extracted knowledge concentrates in a task-specific sparse circuit across layers and positions, and that averaging operators from disjoint demonstration batches enables effective many-shot scaling without expanding the context window. Our code is available at https://github.com/gzxiong/task_operator.
In-context learning (ICL) allows language models to adapt their behavior from instructions and demonstrations in a prompt, but accuracy alone does not reveal how contextual information controls their predictions. We study the acquisition and trainability of ICL with a controlled diagnostic that separates copying a nearby demonstration from integrating information across the prompt. Across 17 models, about a third prefer the copied answer, a few are undecided, and the rest prefer the rule-consistent answer, which goes with higher accuracy and with control by the requested value rather than the closest conflicting example. Checkpoint analyses show that models begin prioritizing task-relevant prompt tokens over conflicting examples long before their accuracy improves beyond chance level. In matched training experiments, synthetic documents that define endings and then query them shift models toward compositional responses on the controlled diagnostic, whereas fine-tuning on definition exercises built from rewritten real documents improves accuracy when prompts reassign label meanings. The results suggest that ICL is not acquired as a single, uniform capability. Instead, prompt control, answer preference, construction without candidates, and transfer can develop along distinct trajectories.
In-context learning (ICL) excels at new tasks from minimal examples, yet we still lack a mechanistic explanation of how few-shot prompts shape a model's function vector (FV)--a causal activation direction that drives task behavior on the ICL query. Across tasks and models, an n-shot FV is well-approximated by a linear combination of example-level sub-FVs, suggesting additive and composable contributions from individual demonstrations. Beyond additivity, we show that models contextualize individual examples' representations based on prior examples to adaptively reweight which demonstrations dominate the FV: attention shifts toward examples that are more informative and less ambiguous under the context. Finally, a causal decomposition separates Query-Key routing from Value updates, finding that contextualization's most consistent contributions to FV quality arise from Query-Key alignment--particularly in ambiguous settings--while Value-mediated effects are more heterogeneous. Together, these results unify additive superposition with context-dependent attention reweighting into a mechanistic, testable account of how few-shot prompts implement tasks.