Contextualize-then-Aggregate: Circuits for In-Context Learning in Gemma-2 2B
Authors: Aleksandra Bakalova
Yana Veitsman
Xinting Huang
Michael Hahn
Venue: COLM 2025 (2025)
Keywords: In-Context Learning
Mechanistic Interpretability
TL;DR: We use causal interventions to uncover a two-step strategy used to assemble task information from the fewshot examples in a prompt.

Abstract: In-Context Learning (ICL) is an intriguing ability of large language models (LLMs). Despite a substantial amount of work on its behavioral aspects and how it emerges in miniature setups, it remains unclear which mechanism assembles task information from the individual examples in a fewshot prompt. We use causal interventions to identify information flow in Gemma-2 2B for five naturalistic ICL tasks. We find that the model infers task information using a two-step strategy we call contextualize-then-aggregate: In the lower layers, the model builds up representations of individual fewshot examples, which are contextualized by preceding examples through connections between fewshot input and output tokens across the sequence. In the higher layers, these representations are aggregated to identify the task and prepare prediction of the next output. The importance of the contextualization step differs between tasks, and it may become more important in the presence of ambiguous examples. Overall, by providing rigorous causal analysis, our results shed light on the mechanisms through which ICL happens in language models.

Published as a conference paper at COLM 2025
Contextualize-then--Agggggge Circuits for In-Context Learn- ing in Gemma-2 2B
Aleksandra Bakalova, Yana Veitsman, Xinting Huang & Michael Hahn Saarland Informatics Campus, Saarland University (abakalov, yanav, xhuang, mhahn)@lst. uni -oiissaalaadddcm
Abstract
In-Context Learning (ICL) is an intriguing ability of large language mod- els (LLMs). Despite substantial amount of work on its behavioral as- pects and how it emerges in miniature setups, remains unclear which mechanism assembles task information from the individual examples in a few-shot prompt. We use causal interventions to identify information flow in Gemma-2 2B for five naturalistic ICL tasks. We find that the model infers task information using two-step strategy we call contextualize-ther aggregate: In the lower layers, the model builds up representations of indi- vidual few-shot examples, which are contextualized by preceding examples through connections between few-shot input and output tokens across the sequence. In the higher layers, these representations are aggregated to identify the task and prepare prediction of the next output. The importance of the contextualization step differs between tasks, and it may become more important in the presence of ambiguous examples. Overall, by providing rigorous causal analysis, our results shed light on the mechanisms through which ICL happens in language models.
1 Introduction
In-Context Learning (ICL) is an intriguing property of large language models and has spurred a substantial amount of interest into how transformers able to perform it (e.g. Brown et al., 2020; Min et al., 2022; Garg et al., 2022; Akyürek et al., 2023; Cho et al., 2025; Wang et al., 2023). Prior work in mechanistic interpretability has found function vectors (Todd et al., 2024; Hendel et al., 2023): attention heads whose output encodes task information, and which are causally responsible for the prediction of the response. However, the circuit by which these vectors are assembled remains only partly understood. Recent work has proposed that each few-shot example's 's output token computes task information, which is then aggregated by attention heads to predict the next output (Wang et al., 2023; Cho al., 2025; Kharlapenko et al., 2025), but this strategy has also been observed to leave part of the models' performance unexplained (Cho et al., 2025).
In this paper, we use causal interventions (e.g. Vig et al., 2020; Geiger et al., 2021; Meng et al., 2022) to identify circuit performing ICL on five naturalistic tasks formatted as fhw-shot prompts (Capitalization, Country-Capital, Present-Past, Person-Sport, Copying) in Gemma- 2 2B (Team et al., 2024). We replicate the relevance of the aggregation step suggested by prior work, but show that it explains only a portion of the model's full performance. We use causal interventions to identify information flow that recovers at least 90% of the model's performance. A key idea is to first identify a computation graph between the tokens in prompt (Figure 1). We find that the required graph differs between tasks: whereas the simplest circuit is largely sufficient on some tasks, others require additional computation paths. We find that the representations of few-shot examples later in the prompt are contextualized by information from prior few-shot examples, particularly, the last preceding few-shot example. Using causal interventions, we show that contextualization transports information about the input and output spaces, and the task itself. In particular,
httpss: //gns:://iita //t/ wwtp://aacoooollaallllc..ii.iiii

---

Published as a conference paper at COLM 2025
France tertis LibilgaiBBugggiiia tSofia InGermay
Parallel Contextualization Circuit
3
Circuit
ICL Task Input
Y1
2
X
y
X
V2
LM
Head
Contextualizatioon
Aggregation
Figure 1: An N-shot ICL prompt (here, N 2), at the example of the Country-Capital task. The green edges describe a AGGREGATION subcircuit that assembles information at each few-shot example in parallel (x; + yú) and then aggregates these at IN+1. The blue edges describe a CONTEXTUALIZATION subcircuit that contextualizes the representation of each few-shot. A circuit involving both components recovers most of Gemma2-2B's performance; aggregating without contextualization leads to breakdown in performance on this task. See Figure 24 for other tasks.
Aggregation edges generally move information about the task from individual examples to the last token. In the PeersonSSport task, they move information about input/ output spaces without regard to functional relationship.
Contextualization edges generally move information about input/ output spaces, and sometimes the task.
Overall, our main contributions are to:
1. Obtain causally faithful information flow from the few-shot examples to the next-token prediction explaining 90% of the model's performance.
A key idea here to focus on position-level circuits, where information flows only between a restricted subset of positions in prompt, which is key to making the resulting computation graph interpretable.
2. Establish the importance and function of a confextualization step that precedes aggregation. proposals(Che
This distinct from prior proposals (Cho et al., 2025; Kharlapenko al., 2025) focusing on the aggregation step. Contextualization becomes even more important in the presence of ambiguity (Section 3.3). Indeed, beneficial even in a synthetic setup (Appendix A).
2 Background: Aggregation and Function Vectors
Prior work on the mechanics of ICI. in LLMs provides a few robust insights. Multiple studies document the existence of task vectors or function vectors (Hendel et al., 2023; Todd et al., 2024; Song al., 2025; Kharlapenko al., 2025; Yin and Steinhardt, 2025): Several attention heads at the last prompt token (in our case, fN+1) output vectors that, across ICL tasks, encode task information. Erasing them prevents ICL; patching them elsewhere leads to execution of the task. Todd et al. (2024) specifically operationalize Function Vector Heads as heads whose activations, when patched from valid ICL prompts to shuffled uninformative prompts, can most increase the likelihood of the desired target. Yin and Steinhardt (2025) find that Function Vector Heads are more important to ICI. than Induction Heads, hypothesized to be key in prior work (Olsson et al., 2022; Crosbie and Shutova, 2024) (cf. also Bansal et al. (2023)). However, it remains largely open what circuit assembles this information, and how information is assembled from the different sho- examples to these function vectors.
Kharlapenko et al. (2025) discover a sparse feature circuit in Gemma-2 2B, using sparse autoencoder (SAE) features encoding task information (in few-shot examples) and features that execute them (at IN 1). However, the circuit does not distinguish between inputs and outputs for different shw- examples, leaving open whether task information is
2

---

Published as a conference paper at COLM 2025
created independently at each example or whether the representations are dependent across few-shot examples. Cho et al. (2025) suggest a different circuit: (1) assembling an encoding of each few-shot example's input and moving it to few-shot label position, (2) attending to labels whose input is similar to the query, and copying these labels to the query. Their experiments are restricted to semantic text classification (e-g., sentiment), in which few-shot inputs typically are either semantically aligned or opposed. In addition, the arguments are largely based on evaluating representational similarity and on supervised probing, which does not establish that this mechanism is causally sufficient. While causal intervention, through zeroing out the relevant components, shows this mechanism is relevant to ICL (Section 5.1 in Cho et al. (2025)), it leaves open if the mechanism alone can fully explain the only a part of the overall mechanism. Overall, while function vec-
behavior, or whether tors play an important role for many ICI. tasks, the information flow leading to aggregation of task information across few-shot examples remains unclear. To foreshadow our results, we show that, while parallel aggregation of per few-shot representations is important, the overall circuit is dask-dependent, and may require a variety of additional components which, as will show, contextualize the per few-shot representations on other few-shot examples.
3 Results
We focus on Gemma-2, especially its 2B version, but our methodology allows us to also investigate larger variants (9B and 27B). We evaluated the model on tasks from Todd et al. (2024), and determined five tasks on which the 2B model achieves high (between 88% and 100%) accuracy: Copying? 2, Present-Past, Capitalization, Country-Capital, Person-Sport (Figure 2).3 3 We decided to focus on tasks with high accuracy in order to obtain clear signal, as for these tasks there is a clear understanding of the input-output behavior exhibited by the full model. We focused experiments on shorter (3-shot) and longer (10-shot) prompts.
Each of the five tasks has an input space X' (e.g,, arbitrary lowercase words for Capitalization; countries for (ountry-Capital) and an output space y (e.g., arbitrary capitalized words for Capitalization; cities for Country-Capital, and defines a func- tion f : X' + y (e.g., mapping countries to their capitals) (Appendix, Figure 25). Each N-shot prompt then has the following structure: x1f1y1111 ...XNNXUNNNNNNNIIN...
V :: X; E X,y, E V,y, = f(x)), Vt,n; = \n.
3.1 Identifying Circuits
To localize the behavior of the model, we use patching, an approach widely adopted in the literature (Wang et al., 2022; Hanna et al., 2023). With this technique, model is viewed as a computation graph with activations in different layers as nodes, and computations between them as edges. We then ablate some of the edges in the graph, forcibly replacing computation along specific edge with the one computed on the counterfactual input. Counterfactual inputs are designed to erase relevant information from the input while leaving other information intact for localization of model behavior. For example, when substituting "Berlin" with counterfactual "Paris", we isolate computations specific the city's 's identity, not those activated by its category (city) or token-type (word). If ablating a set of edges does not lead to drop in the model's performance, then the information unique to the original input relative to the counterfactual input, which was transferred along these edges, is not causal for the model's prediction. This allows to discover circuits - subgraphs of the full model's computation graph explaining major part of model performance a specific task. In contrast to the standard approach, we focus on the information flow between positions in a prompt; consequently, we differentiate between activations of the same heads at different positions. We keep nodes for x; and y/ in different examples distinct, which allows us to analyze the information flow between few-shot examples."
2This is not included in Todd et al. (2024), but we added it due to its naturalness. We also evaluated six other models at the 2B/3B scale, finding that they overall underperformed Gemma-2 2B at 10 shots on these tasks (Table 7), further motivating our focus on Gemma-2. "This contrasts with Kharlapenko al. (2025), who collapse all x; and y; nodes into one node each.
3

---

Published as a conference paper at COLM 2025
Cap talication
Courtry
Present-Past
Person Sport
Copy
0.8
0.6
aogregetton
2221028587 camaptadllaap (l3
A2pregutoon on cantaduul pacen 14 Ni -0 -39
Chosen Croaft
Bo Frade
Figure 2: Accuracy of circuits (Gemma-2 2B, 9B, 27B). The ACCRECATION circuit under- performs the full model by large margin in some tasks. Contextualizing edges recover large part of the gap. Whereas the AGGREGATION already performs well in some tasks, we note that the contextualizing edges become crucial when evaluating on a difficult subset designed to be ambiguous between two tasks (Section 3.3).
We first discover a position-level circuit. Nodes are the positions in the prompt; edges are directed ares between them, representing the flow of information (Figure 1). Importantly, in position-level circuits we collapse all edges across layers into a single edge. Ablating such an edge means ablating all corresponding edges in every attention head and layer. For illustration, X2 in position-level circuit can only receive four incoming edges: from f,,11//1 or n1. 5 Position-level circuits provide useful and interpretable upper bounds on information flow, though they do not distinguish between computations happening in different layers. We treated all tokens within few-shot input x; output y; together, as the number of tokens varies between few-shot examples. As positions, we distinguish the few-shot inputs x; and outputs y; and the separators within (f;) and between (n;) examples; though n; play no role in any of the identified circuits.
Patching Methodology To identify circuits, we define a simple class of counterfactual inputs where all few-shot inputs and outputs (including the query xN+1) are replaced with random words (Appendix H). The corrupted prompts still contain the same separators, and maintain the overall number of tokens of each x/y;; (Appendix H, Figure 22). When we ablate an edge from position A to position B, the key (K) and value (V) activations of A when queried by B are replaced with activations computed on corrupted prompt. The K/V activations of A when queried by other positions, as well as Q activation of B remain clean (Appendix H. Figure 28). This patching applied simultaneously head. We define the output of a circuit on an input prompt by greedily decoding a response YN+1 until a separator is generated, while ablating all edges outside the circuit."
each layer and
AGCREGATION Circuit: Computing representations at each example and aggregate Based on prior work, we first hypothesized a minimal circuit consisting of edges between few-shot inputs and outputs (z, + yi), from few-shot outputs to the last separator (у; +
1), and also from the query to the last separator (IN + IN 11)?" We dub this the ACCREGATION circuit, as it allows each few-shot to assemble its own representation based on y;) in parallel, and then allows fN+1 to make its prediction based on both the query and the set of fho- examples (green edges in Figure 1). This is the primary mechanism suggested by prior work (Cho et al., 2025; Wang et al., 2023; Kharlapenko et al., 2025). We evaluated the degree to which this subcircuit is sufficient for task performance." For this,
"The edges from to bos/ eos tokens are always ablated (Appendix F)).
6We do, however, retain edges supporting the autoregressive prediction of the answer (XN+1 +
4N+1 and fN+1 + +1) 1) (Appendix E).
We also always include all edges from one position itself, and never ablate these Appendix E).
@While Cho et al. (2025) showed zeroing out these edges to hurt model performance, establishing that they do play role, their method does not show sufficiency of the subcircuit. Interestingly, Cho al. (2025) propose that x; and yi are connected via 4 whereas we find a direct connection is largely sufficient, simplifying the resulting circuit. We note that Cho et al. (2025) found that performance suffers when zeroing out x; + t or I; + Vi and hypothesized that these edges transport information about x,. However, they did not verify that these edges provide information about (rather than, e.g, about the presence of prompt template). In contrast, our patching methodology allows to
4

---

Published as a conference paper at COLM 2025
we ablated all edges other than those in this minimal circuit. With this ablation, the model achieves 2 85% of the full accuracy on three tasks, but shows a huge drop (< 60% of full model) on Capitalization and Country-Capital (Figure 2). This confirms that the mechanism is causally important, but insufficient for overall explaining model behavior.
CONTEXTUALIZATION Subcircuit: Edges between different few-shot examples We next investigated which edges are needed beyond the ACGRECATION subcircuit, aiming to find small circuit recovering 2 90% of the accuracy of the full model. We call all the edges in this circuit except those in the AGGREGATION subcircuit a CONTEXTUALIZATION Specifically, the CONTEXTUALIZATION subcircuit comprises edges that contextualize the representation of Mi) using prior few-shot examples (x V;), where j < i. These edges connect nodes A + B, with A E ززز < i} and B E (x for < N. Note that the exact edges in this subcircuit may vary slightly across tasks.
subcircuit.
We first ablated only edges involving non-final separators from the full model, and found this to do little harm across the five tasks compared to the full model, especially at 10 shots and for the 9B model (Tables 2, 5, 6), showing that relevant task information need not be routed through the separators in and between the few-shots, i.e., the separators may provide information about the presence of a prompt template, but do not play a nonredundant role in assembling information about the task. We thus continued with these edges ablated. We next considered the remaining logical possibilities, grouped by the types (input, output, separator) of positions: within-type edges (1) x; + xj (for all 1 <i<ja N + 1), (2) y: + Vj (for all 1 < jijjs S N), and across-fype edges: (3) x; yj (1 S < S N), (4) x; - IN+1 (1 i N), (5) yr + xx (1 S Sijjs N + 1). We investigated adding each of these groups individually to the aggregation circuit (Table 1). We found that adding (1) and (2) provided strong accuracy gains on some tasks, and focused on these two groups.
compaseparatonn
We next focused on different circuits involving these edges. We considered both general edge sets of types (1) and (2), and the local subset (x; + x/+1 and y; + //++)); the latter were sufficient in all but one case (y: - yj needed for Capitalization). For each task, we chose the simplest circuit achieving > 0.9 at N = 3, or (if there is none) the circuit achieving highest accuracy at N = 3, 10, all based on the 2B model's accuracies (Tables 2). We illustrate the result for Country-Capite Appendix, Figure 24. 2B parameters and 10 few-shots, the circuits explained 2 90% of the full model's performance, far exceeding the aggregation-only circuit on some tasks (Figure 2). We next evaluated the circuits at 9B and 27B, using the same methodology (Tables 5, 6). At 10 shots, the circuits that we had chosen based on the 2B model explained at least 90% of the full model's performance in these model. Importantly, the ACCREGATION circuit again underperformed the circuits involving contextualization on most tasks (Figure 2).
in Appendix, Figure 1; circuits for all five tasks are shown in
Projecting to Activation-Level Circuits In order to identify the heads responsible for the information flow we have identified, we next projected the position-level circuits to activation-level circuits. Here, nodes are activations of attention heads indexed by both their positions and layers; edges are directed ares between the input and output of indívidual attention heads. More precisely, heads can be defined as tuples of (layer, head, position), and edges can be described as (layer, head, start position, end position). We used gradient- based method (Michel et al., 2019; Syed et al., 2023) (see Appendix FF for technical details) to obtain a circuit on the level of activations with edges defined by attention heads, in the 3-shot setup (Figures 9-13). The circuit reveals that the contextualization step tends to precede the aggregation step: contextualization happens in the lower half layers, most aggregation edges are in the middle and upper layers. Prior work has established the existence of function vector heads, specific heads aggregating causally relevant task information at NN+1 (Section 2). We next identified function vector heads by using the method described in Todd et al. (2024) for each task independently and selecting the top-10-scoring heads for each
conclude that these edges, at least in Gemma-2B on our prompt template, need not causally provide information beyond the presence of prompt template (which is kept constant in our corrupted inputs). We caution that this difference need not invalidate the conclusion of Cho et al. (2025) as they used a different template, using the word "label" for f,, whereas we use the "\t" punctuation.
5

---

Published as a conference paper at COLM 2025
Yi fpe1
y/+ V
X1+
y with random word
within y
x within X
x with ra random word and y with random word and with other task example
0.00 0.25 0.50 0.75 1.00 0.03
0.25 0.50
0.75
1.00
0.00
0.25
0.50 0.75
1.00
Accuracy
Accuracy
Accuracy
5008 Not Signiicarr
Creul Performmece
Figure 3: Patching experiments manipulating information about the input available to each edge in the full circuit, the example of the Country-Capien task, 10-shot. For each ablation, we indicate whether the drop in performance is statistically significant compared to the full circuit performance in binomial test with 0.05. See Figure for results from all tasks.
task. These heads were largely part of the identified circuits and mostly contributed to the yi + IN+1 edges (Appendix, Figure 26). This confirmss that the circuits we have identified describe the buildup of task information up to the function vector heads.
3.2 Which Information is Routed?
So far, we have identified edges of information flow, finding that the aggregation step suggested by prior work is causally important to task success, but that it is preceded by contextualizing subcircuit. We next investigate which information is passed along the edges. We are interested not in all information, but only in information about the input that is causally implicated in the downstream prediction. We do this by performing targeted manipulations on the input to manipulate the information present in specific activations.
For each edge, we specifically contrast the following hypotheses: An edge from some few- shot example (x, yi) might transport causally relevant information about (1) the specific tokens, (2) the functional relationship between x; and y; (but not their token identities), (3) the type of the x; or y/ tokens without regard to their functional relationship.
Patching Methodology We vary the information about the input available to a particular edge by constructing contrastive input sample in which certain aspects have been altered (Appendix H, Figure 22), and patching the K and activations directly feeding into the edge with the corresponding activations computed Figure 28). For instance, when patching an edge y; - IN 1, we modify the key and value activations associated with the positions of y; across all attention heads in the model. These modifications occur specifically when the activations are queried from the positions corresponding fN+1. Patching done in parallel in each layer and head. All patching is done within the overall circuit identified for the task, i.e., edges not present in the circuit are ablated, and paths not present in the overall circuit play no role even in the patched version. For instance, to determine whether the edges y; + fN+1 transport contextualized information about preceding few-shots (xy, yr) ) (i' < i), we patch all inputs to this edge at position y; with activations computed on an input where the preceding few-shots have been corrupted; by varying the degree of corruption, we analyze which specific information is transported (Appendix H, Figure 22). We note that this methodology allows to focus causally relevant information used by downstream components. It is thus fundamentally different from supervised probing, which does not distinguish between information used or unused by downstream components. We focus on the 2B model for tractability.
that contrastive input (Appendix,
A AGGREGATION Subcircuit We first investigated the information transported by the edges in the AGGREGATION subcircuit, consisting of the x; - y;, XN+1 + IN and y; + IN+1 edges. Prior work (Cho et al., 2025; Kharlapenko et al., 2025; Wang et al., 2023) suggests that the y1 - IN-1 edges aggregate task information from the examples. Indeed, when we patched these with arbitrary x; or yi (which disrupts the functional relationship), functional- ity was largely destroyed (see Figures 3 and 8 for all patching results). In contrast, when we patched the 1/; + IN edges with other few-shots that are valid for the task, functionality was preserved. We also verified that the x; - yi edges transport xy: patching these edges with
6

---

Published as a conference paper at COLM 2025
corrupted x; leads to total failure. This is to be expected, as in the simple AGGREGATION subcircuit, these edges are the only way for x; to influence the output. These results confirm prior proposals about the aggregation step (Cho et al., 2025; Kharlapenko et al., 2025, Wang et al., 2023) with rigorous causal analysis.
Our methodology allowed us to obtain two further findings. First, the y; IN+1 edges causally transport only task information and no token information. Namely, even in the presence of contextualizing edges, i.e., even though there are multiple paths from few-shots to the +++,, performance did not deteriorate even when the y; - fN+1 and contextualizing edges transported information from different prompts when they were valid for the task.' 9
Second, we next asked whether the edges rely on the functional relationship between x; and yi or only their semantic types. We patched with pairs where the semantic types were preserved but the functional relationship had been disrupted (e.g., in Country-Capital "France \ &Loondon,, "Canada \ tBeijing", etc. Appendix H, Figure 22). In four tasks, this led to a failure, confirming that the functional relationship is transported by y: -> IN+1- The one exception is the Person-Sport task: here, functionality is largely preserved (> 90 % of the full accuracy) as long as the patched x E X, E y. This phenomenon provides an understanding of why ICL has sometimes been observed to be robust to perturbed labels (Min et al., 2022): for some tasks (in this case, the Person-Sport task), the task information is inferred on the basis of only the semantic types, without regard to their functional relationship. In the context of the overall circuit:
The y; + 1 edges generally move information about the function f from individual examples (z) to IN 1- In the Person-Sport task, they move information about X' and ) without regard to functional relationship.
In order to understand the role of these components further, we investigated the functional behavior of the model with ablations applied. Overall, we observed three fallback tasks which accounted for most interpretable errors: copying the query, producing from the correct output space, copying another token from the prompt. With full ablation of the y; IN edges, behavior was often accounted for by copying the query or copying another token from the corrupted input prompt (Appendix H, Figures 14, 15, 16, 17, 18). We note that as the corrupted prompts maintain the separators, the last separator position can still infer the presence of repeated structure, potentially explaining the presence of copying behavior. When patching the edge while preserving the output space, behavior often amounted to reproduction from the correct oulput space when there is well-delineated output space (e.g., Capitalization, Country-Capital, Person-Sport), or copying from the corrupted prompt when there is none (Copying). This is expected, as such prompts maintain repeated structure in terms of a specific output space, though with disrupted functional relationship.
Is the ICL aggregation circuit a type of induction circuit? An interesting question is whether attention in the yj - IN+1 edges depends on the specific inputs, or just on the few-shot template. Induction circuits - argued to be essential for ICL by Olsson et al. (2022); Crosbie and Shutova (2024) - involve attention specifically to preceding position preceded by the same as the current token, one might thus hypothesize that heads involved in aggregation specifically attend to few-shot examples similar to the query. Indeed, Cho et al. (2025) argued that, at least in the text classification tasks they studied, the ICL circuit is a type of induction circuit, and the strongest information flow came from shot examples where the input texts are semantically closest to the query text. However, this argument was based on representational similarity, not any direct causal evidence. Relatedly, theoretical constructions based on linear transformers performing ICL for linear functions (e.g. Von Oswald et al., 2023; Vladymyrov et al., 2024; Mahankali al., 2023) also crucially
'We emphasize that this is a claim about information that has a causal downstream effect in the circuit, not about information that can be obtained via probing, which is likely to be richer and include token information.
An indirect argument based on representational similarity is made by Cho et al. (2025); our analysis improves by causally localizing the invariance to perturbed input-output mappings to the y; + ty+1 edges, ruling out role for other pathways.
7

---

Published as a conference paper at COLM 2025
assume that attention is strongest to examples where x; is similar to the query XN+1, in line with the induction circuit. We thus investigated patching the key or query vectors involved in these attention edges with other valid prompts (preserving functional relationship but changing the specific tokens). Strikingly, across all five tasks, at both 3 or 10 examples, there was almost no discernible effect on accuracy (Table 4). This shows that, at least in the tasks and model considered here, the aggregating attention edges do not exhibit the functionality expected of the classical induction circuit: Attention not preferentialyy given to specific few-shots on the basis of similarity to the query." 11
CONTEXTUALIZATION Subcircuit We next moved to analyzing the information trans- ported by the contextualizing edges, starting with the edges connecting y's (y; - 1/+11 or general yi + vj). We found that these edges transport both information about f and about V. Patching the 1/i + Mi+1 edges for Country-Capital only decreases accuracy significantly when y leaves ), but not under other corruptions (Figure 3). In Capitalization, at 3 shots, replacing with y E y hurts much less than a general ablation of y (accuracy 81% vs 56 %; Appendix F, Figure 8), which is harmful as full corruption of these edges and leads to large increase in copying behavior; nonetheless, accuracy still drops compared to ablations keeping yx E f- At 10 shots, Capitalization only responds to ablations where y leaves V. In both cases, ablating the output space information leads to an increased number of copy-type mistakes (Appendix H, Figure 19). As noted before, edges do not provide specific token information used by downstream components at IN: we verified that a token mismatch created by intervening on (x y/-1) for all y/ + yj edges but not the y; + IN+1 edges has no discernible effect on accuracy if the functional relation between is preserved by the prompt used for intervention. Overall, our causal interventions warrant the conclusion:
The y; - yj y: >> Vi+1 edges contextualize the representation each few-shot example by transporting the type of y, and sometimes the task f, from (xx,yj) (j < i).
We nextstuudee the edges connecting inputs x's. Ablating these on Present-Past and Person- Sport turns out to not lead to statistically significant drop, but the drop is substantive in Country-Capital (Figure 3). Here, ablations only hurt when x left X', showing that x; -> Xj transports the input space, but not (causally) the token x,. Such ablations again lead to substantial fraction of copying responses, i.e., the functional input-output behavior is impacted (Appendix H, Figure 20). Patching prior x;'s for these edges, we found that they affect the final prediction both via x; + x1+1 + Vi+1 -> IN+1 and via XN - XN+1 - fN+1 edges (Appendix, Table 8). Overall:
The x; + X1+1 edges contextualize the representations of xN+1 and of individual few-shot outputs y; with information about X.
3.3 A CONTEXTUALIZATION subcircuit is important in the presence of ambiguity
We conducted controlled experiment aiming stress-test the aggregation step. We designed the following hard version of the Present- task. We take advantage of the fact that some English verbs have identical present and past tense forms (e.g., "put", "spread"). In prompt for the Present-Past task, examples using such verbs are thus ambiguous between the Present-Past task and simpler Copy- ing task (Appendix H, Figure 23.1). In this exper- iment, most individual examples do not give full task information, and may in fact provide mislead- ing information. We hypothesize that direct aggregation of uncontextualized per few-shot representations as hypothesized by prior work may be less useful in this case, because
H1
cut It cut
give 1t gave
make V.
ambiguous examples
unambiguous examples
final separator
Figure 4: Two hypotheses for the causal graph on ambiguous inputs.
NNonethelsss, ty+1, collecting information from fhw-shot outputs that followed the same separator token \t, might still rely on induction-head-like behavior, with the single preceding token the key.
8

---

Published as a conference paper at COLM 2025
individual examples might, on their own, provide conflicting task information. The contex- tualization step might help resolve this ambiguity before the aggregation step. Even when out of 10 few-shot examples are ambiguous, the full model continues to perform the task ACGGREGATION)-NN
at very high accuracy (95%). Whereas the aggregation performed at 93% on the standard Present-Past task (and similarly on the Copying task), its performance on the ambiguous prompts dropped to 56% (Appendix H, Table 9).
circuit had
Wrong outputs largely consist of copying, expected as the ambiguous examples are compat- ible with the Copying task. Note that, in the EEGRRGATTIO--DN circuit, the activations at y; feeding into !/ + fN are a mixture of activations for Copying or Present-Past; they each largely succeed in inducing the correct prediction on their own without contextualization when the task is consistent across the prompt, but (N i1 fails at resolving conflicting task information when the task information is mixed across the prompt.
We found that contextualization is needed to explain the model's accuracy. Adding x; + Xj41, y/i - Vi-1 edges recovered part of the drop, but still underperformed the full model (68%). We next investigated which further edges help close the gap to the full model; indeed, in contrast to the standard tasks, we found separators now to play a role, as the full model with only eeparator-involiing edges ablated performs at 85% (Appendix H, Table 9). Overall, adding contextualization subcircuit with edges y; - y/, y/ - f x, fi, fj Vi + f++1 X; Xi+1 recovered almost all of the full model's performance (89.9%).
We investigated two hypotheses as to how contextualization resolves ambiguity (Figure 4):
1. (H1) Ambiguous examples obtain information from unambiguous examples, con- textualizing their own representations so that they encode the Present-Past task instead of the Copying task.
2. (H2) Unambiguous examples obtain information from ambiguous examples, con- textualizing their own representations so they can override misleading functional information from ambiguous examples.
We note that these hypotheses are not exclusive and might be true at the same time. To distin- guish between them, we patched the K and V activations at either all ambiguous or all un- ambiguous examples feeding into yi + IN edges, with activations computed from either all-ambiguous or allllul all-unambiguous prompts, in both cases erasing the presence of the other type of example (see Appendix H, Figure 23.2). Patch- ing the yi 's in ambiguous examples had little ef- fect the accuracy (Table 3, Figure 5). This sug- gests that, even if these examples may receive information about the presence of Present-Past examples, this is not causally important to task disambiguation.
unambiguous 1
ambiguous y; 1
0.0
0.2
0.4
0.6
0.8 1.0
Accuracy
Figure 5: Patching!! yi + fN+1 connections only ambiguous examples (bottom) using activations taken on all-ambiguous (red) or all-unambiguous (green) barely hurts. In contrast, performing these patches on only unambiguous exam- ples (top) leads to a drastic drop, even when patching with an all-unambiguous prompt (green).
In contrast, patching the unambiguous examples had great effect: Accuracy dropped to almost zero when patching with ambiguous prompts (in which case all information would be routed through the ambiguous examples, but we saw that they do not causally provide disambiguation signal via y/1 - IN 1). Strikingly, accuracy dropped drastically even when patching with fully unambiguous prompts, showing that contextualizing the unambiguous examples with awareness of the presence of ambiguous examples is key to downstream performance. Overall, our results allow us to reject H1. They are as expected under H2:
In the presence of ambiguous examples, contextualization changes the representations of unambiguous few-shot examples. These altered representations of unambiguous exam- ples are key to transporting task information to N++1, overriding potentially misleading information from ambiguous examples.

---

Published as a conference paper at COLM 2025
4
Discussion
On five tasks, we have shown that Gemma-2 2B performs ICL using cascade contextual- izing the representations of individual examples, and aggregating task information from them. Using rigorous causal interventions, we extend prior work on mechanistically un- derstanding the mechanisms of ICL in transformers (e.g. Wang et al., 2023; Cho al., 2025; Kharlapenko et al., 2025; Hendel et al., 2023; Todd et al., 2024; Bansal et al., 2022; Sia et al., 2024; Song et al., 2025; Yin and Steinhardt, 2025; Crosbie and Shutova, 2024). We extend on work function or task vectors by showing that they are largely part of the aggregation circuit hypothesized by Cho et al. (2025); Kharlapenko et al. (2025). Most importantly, we establish the importance of contextualization step, especially in the presence of ambiguous examples in a prompt.
Our work is distinct from Cho et al. (2025) in providing computation graph supported by causal effects rather than representation similarity. A major difference between causal experiments in our work and Cho et al. (2025) is that we construct circuit sufficient for recovering model performance, whereas Cho et al. (2025) zeroed out edges to confirm their importance. major difference to Kharlapenko et al. (2025) is that we keep different few-shot examples distinct in our circuits, allowing us to identify the contextualization step.
Theoretical work has considered ICL both from the perspective of learning simple parametric functions (e.g. Akyürek et al., 2023; Von Oswald et al., 2023; Mahankali et al., 2023), and as general Bayesian inference (e.g. Xie et al., 2022; Wies et al., 2023; Hahn and Goyal, 2023). Theoretical constructions for transformers performing ICL often correspond to the ACCREGATION subcircuit (Akyürek et al., 2023; Mahankali al., 2023), but contextualization of the few-shot inputs (Von Oswald et al., 2023; Vladymyrov et al., 2024) outputs (Chen et al., 2024) does play a role in some constructions (see Appendix A for more), in agreement with findings in an LLM. Mechanistic understanding of ICL also has the potential to explain some of its behavioral aspects. For instance, we found that, in the Person-Sport task, the circuit only assembles input and output spaces without regard to functional relationships, explaining the robustness to label noise sometimes observed (Min et al., 2022).
Our work is grounded in tradition of identifying important components in neural networks (e.g. Michel et al., 2019). Modern circuit discovery in LMs was pioneered by Wang et al. (2022); Hanna et al. (2023); unifying framework was provided by Conmy al. (2023). Scaling this to larger models (e.g. Syed et al., 2023; Hanna et al., 2024; Ferrando and Voita, 2024) is focus of recent research.
A limitation of the present study is that, to keep computational cost manageable, it focuses on one model family and five tasks. second limitation that it focuses on causal under- standing of information flow between positions, and does not interpret ML.s or individual heads. Future work might combine our results with a further scaled-up version of the sparse feature circuits of Kharlapenko et al. (2025) to understand contextualizatton at the level of individual representations It is likely that further advances in interpretability are needed to comprehensively interpret all individual components in models of the scale studied here.
Conclusion
We have conducted extensive causal interventions to understand the flow of information from individual few-shot examples to the next-token prediction in ICL, at the example of five tasks in Gemma-2 2B. Our most important result is to establish a two-stage procedure, whereby the model contextualizes apprsennattins of individual few-shot examples, which are then aggregated to prepare next-token prediction.
Author Contributions
AB led the project and conducted the experiments. YV contributed to scaling experiments. XH contributed to discussions throughout the project. MH supervised the project, provided input the experimental design, and drafted the paper text. All authors refined the paper.
10

---

Published as a conference paper at COLM 2025
Acknowledgments
Funded by the Deutsche Fosscungsgemeeisshafft (DFG, German Research Foundation) Project-ID 232722074 SFB 1102. We thank Entang Wang for comments on the manuscript.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. 2023. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614 446500.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. 2023. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations.
Hritik Bansal, Karthik Gopalakrishnan Saket Dingliwal, Sravan Bodapati, Katrin Kirch- hoff, and Dan Roth. 2022. Rethinking the role of scale for in-context learning: An interpretability-aaseee case study at 66 billion scale. arXiv preprint XrXo:2212..9095.
Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirch- hoff, and Dan Roth. 2023. Rethinking the role of scale for in-context learning: An interpretability-baeee case study at 66 billion scale.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901.
Xingwu Chen, Lei Zhao, and Difan Zou. 2024. How transformers utilize multi-head attention in in-context learning? case study on sparse linear regression.
Hakaze Cho, Mariko Kato, Yoshihiro Sakai, and Naoya Inoue. 2025. Revisiting in-context learning inference circuit in large language models. In The Thirteenth International Confer- ence on Learning Representations.
Arthur Conmy, Augustine Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga-Alonso 2023. Towards automated circuit discovery for mechanistic interpretabil- ity. Advances in Neural Information Processing Systems, 36:16318-1633
Joy Crosbie and Ekaterina Shutova. 2024. Induction heads as an essential mechanism for pattern matching in in-context learning, ar preprint ar xrXi:22407000111
Javier Ferrando and Elena Voita. 2024. Information flow routes: Automatically interpreting language models at scale. Xrv preprint ar arXio:2403.00824.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. 2022. What can trans- formers learn in-context? case study of simple function classes. Adoances in Neural Information Processing Systems, 35:30583-30599
Atticus Geiger, Hanson Lu, Thomas Icard, and Christopher Potts. 2021. Causal abstractions of neural networks. Adoances in Neural Information Processing Systems, 34:9574-9586.
Michael Hahn and Navin Goyal. 2023. A theory of emergent in-context learning implicit structure induction. Xiv preprint arXiv:2303.07977
Tal Haklay, Hadas Orgad, David Bau, Aaron Mueller, and Yonatan Belinkov. 2025. Position-
aware automatic circuit discovery. arvio preprint ar XiX::2250.04457.
Michael Hanna, Ollie Liu, and Alexandre Variengien. 2023. How does gpt-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model. arv preprint ar XrXio:2305.0056.
11

---

Published as a conference paper at COLM 2025
Michael Hanna, Sandro Pezzelle, and Yonatan Belinkov. 2024. Have faith in faithful- ness: Going beyond circuit overlap when finding model mechanisms. arXio preprint (rXiv:2403.17806
Roee Hendel, Mor Geva, and Amir Globerson. 2023. In-context learning creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 9318-9333.
Dmitrii Kharlapenko, Stepan Shabalin, Arthur Conmy, and Neel Nanda. 2025. Scaling sparse feature circuits for studying in-context learning. In Forty-second International Conference on Machine Learning.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. 2023. One step of gradient descent provably the optimal in-context learner with one layer of linear self-attention. Xiv preprint ar Xtt::2307.03576.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. 2022. Locating and editing factual associations in gpt. Advances in neural information processing systems, 35:17359-
17372.
Paul Michel, Omer Levy, and Graham Neubig. 2019. Are sixteen heads really better than one? Advances in neural information processing systems, 32.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2022. Rethinking the role of demonstrations: What makes in-context learning work?
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, al. 2022. In-context learning and induction heads. arxiv preprint arXtt:2209.11895.
Suzanna Sia, David Mueller, and Kevin Duh. 2024. Where does in-context learning happen in large language models? Advances in Neural Information Processing Systems, 37:32761-32788
Jinyeop Song, Seungwook Han, Jeff Gore, and Pulkit Agrawal. 2025. From context to concept: Concept encoding in in-context learning.
Aaquib Syed, Can Rager, and Arthur Conmy. 2023. Attribution patching outperforms automated circuit discovery. arviv preprint drXio:2310.10345
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. 2024. Gemma 2: Improving open language models at a practical size. arXiv preprint ar axXiv:2408.00118
Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. 2024. Function vectors in large language models. In The Twelfth International Conference on Learning Representations.
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber. 2020. Investigating gender bias in language models using causal mediation analysis. Advances in neural information processing systems, 33:12388--2201.
Max Vladymyrov, Johannes Von Oswald, Mark Sandler, and Rong Ge. 2024. Linear trans- formers are versatile in-context learners. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. 2023. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pages 35151-35174. PMLR.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. 2022. Interpretability in the wild: A circuit for indirect object identification in gpt-2 small. Xiv preprint ar Xit:2211.0059.
12

---

Published as a conference paper at COLM 2025
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. 2023. Label words are anchors: An information flow perspective for understanding in-context learning. In Proceedings of the 2023 Conference on Empirical Methods in Nat- ural Language Processing, pages 9840-9855, Singapore. Association for Computational Linguistics.
Noam Wies, Yoav Levine, and Amnon Shashua. 2023. The learnability of in-context learning. Advances in Neural Information Processing Systems, 36:36637-36651.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. 2022. An explanation of in-context learning implicit bayesian inference. Xrv preprint ar arXiv:2111.0000
Kayo Yin and Jacob Steinhardt. 2025. Which attention heads matter for in-context learning?
A ICL on Linear Regression Benefits from Contextualization Step
A substantial amount of work has studied ICL in small transformers trained to solve simple synthetic tasks, particularly linear regression. However, the link between these mechanisms and the ICL circuits inside LLMs largely remain unknown. Theoretical constructions for ICL on linear regression based on gradient descent (GD) and similar algorithms often involve just multiple rounds of parallel aggregation from each few-shot; in particular, a single aggregation step, and is the optimal
single round of GD can be implemented ICL strategy for a one-layer (linear) transformer (Mahankali et al., 2023). However, there theoretical reason believe that a contextualization step is beneficial in preprocessing the representations of individual few-shots before aggregation. First, multilayer linear transformers implement the GD++ algorithm (Von Oswald et al., 2023), which iteratively contextualizes x, 'ss to approximate the input data covariance (Von Oswald et al., 2023), and which allows faster convergence than GD (Vladymyrov et al., 2024); such preconditioning also appears in constructions for ICL on non-isotropic data (Ahn et al., 2023). Second, Chen et al. (2024) showed that preprocessing of individual few-shots is beneficial in sparse linear regression. While these works show that contextualization helps outperform gradient descent, they do not establish that the step is needed for optimal performance, as a priori there might be undiscovered algorithms implementable by transformers outperforming GD with only the aggregation step.
We empirically investigated which attention edges needed when transformers learn linear regression in context. Following Garg et al. (2022), we trained transformers to perform linear regression, i.e., each prompt has latent task E Rd and each few-shot consists of vector E ad and the output y oTx E R. We chose d 20.
We investigated three subcircuits:
1. AGGREGATION: x; + V/, 1/1 + XN+1
2. Between-Xs: x; + x 1
N +))
3. Between-Yst Ys: y; - y (all (al]
We trained the model with the aggregation-only circuit, and three combinations of this and the other two subcircuits (Figure 6). As we trained from scratch, we ablated edges simply by zeroing out their attention weights throughout training.
Garg et al. (2022) trained the model on the aggregation of the losses of each yi (i S 40) conditioned on the prior examples. This was not feasible in our setup, as the set of ablated edges varies with the target. We instead randomized the length N of the prompt in every training step (between 10 and 40) and trained only on the last output (1/++) on each prompt. As there are no separators, the position at which this output is predicted is the query, XN+1-
We varied the depth (4, 6, 8, 12 layers). The model of Garg et al. (2022) had used 12 layers; reasoned that contextualizaation might be more important in shallow models because these provide less opportunity for deep computations at the last position We used the original
13

---

Published as a conference paper at COLM 2025
1.5
Looe
1.0-
Perald
Porale + Between_ys
03
Parad Betweer Between_Ys
Parallel: x; + 1/2 1/: + XN+1
Between-Xs: X; (all <i<js N Between-Ys: y/ + yj (all si<jsM
Figure 6: Loss (MSE) on ICL on linear functions when trained with attention links ablated during training, with transformers with / 4, 6, 8, 12 layers. The AGGREGATION circuit performs much better than random baseline (loss 20.0), but underperforms the full model by large margin shallow transformers. Adding links within x's or y's improves perfor- mance, particularly the latter. Adding both performs at par with the full model, in line with our findings on real-world tasks in Gemma-2.
code from Garg et al. (2022), but increased the batch size to 1024 and decreased the number of steps to 150K for tractability. The model was run on one NVIDIA H100 card.
Results are in Figure 6. The AGGREGATION subcircuit performs much better than random baseline (MSE loss of 20), but suffers substantially for shallow models. X; -> x, edges improve, but cannot close the gap to the full model. y₁ + y edges nearly or fully close the gap. Overall, while parallel aggregation plays an important role, contextualization step is needed to close the gap to the full transformer.
B Computational Resources
For the 2B model, a single patching run takes roughly 40 minutes on 1 NVIDIA A100 40GB GPU on one 3-shot task, and roughly 4 hours on a 10-shot task. Costs are higher on the 9B and 27B models.
C Datasets
We took Person- -poor, Country-Capital and Present-Past datasets from Todd et al. (2024), and selected words for Capitalization dataset randomly from 5000 English words. For the Copying dataset we used the same source as for Capitalization.
All datasets were randomly divided into train and test parts. To construct an ICL prompt, select examples from the train split of the dataset without replacement to as few-shot examples, and one example from the test split to use as final query and target. We used test size of 100 examples in each dataset." 12 Train sizes of datasets range from 39 to 4000.
For the ambiguity dataset, we selected 20 English verbs with identical present and past tense forms, and combined them with the Present-Past dataset. For each prompt, we randomly selected N positions the Present-Past prompt to replace with ambiguous examples.
D Patching methodology
We use path patching (Wang et al., 2022; Hanna et al., 2023) to identify position-level circuits and analyze information flow along edges. This method involves: (1) representing the model as a computation graph with token embeddings as inputs, loss value as output, and
The high computational cost of the full set of patching experiments (Figure 8) prevented us from running larger test sets. Importantly, our key conclusions are drawn only from patching effects with statistically significant changes in accuracy (Figure. 8
14

---

Published as a conference paper at COLM 2025
specific activations nodes; (2) ablating graph components to isolate those responsible for specific behaviors. This section details our computation graph definition (nodes, edges), loss function and ablation implementation.
Computation graph Nodes consist of activations at the input/output of each attention head across positions. Each node is represented as tuple (input/output, layer, head, position), where "position" indicates a token's 's semantic role (e.g., X1, V/,, f2, n4). Note that x; and yi positions typically span multiple tokens but are treated as single nodes. Each node is thus tensor of size (number of tokens in position, hidden dimension).
Edges represent computational connections between nodes. In our position-level circuits, edges exist exclusively within attention heads, as only these transfer information between positions (Appendix É details edge inclusion criteria). Removing an edge between (input, L, H, pos=x)) and (output, L, H, pos=y/1 disconnects these nodes, implemented by computing:
ott.ge/...t.ossssseeeeeeeeorrrep (htph:/==s)...kppppppptpLLptp.): (inpl
instead of Oatpop/-EHpppsss---teeerrrrrrrpp (input,,,,pps==t))) http://phee_upppppooopuououooood (input,,,H.oos-yy))
where subscript C denotes counterfactual activations.
When discovering position-level circuits, we abstract all edges between the same pair of positions into a single edge. Specifically, we treat all edges of the form (liputt, L, H, pos=A) + (output, L, H, pos=B) for all L, H) as a single edge: (input, pos=A) - (output, pos=B). Thus, removing such an edge corresponds to removing all edges between positions A and B across all attention heads and layers.
We use cross-entropy loss over all tokens in the output position (yN 1). Performance drop is measured as accuracy decrease between ablated and normal model runs.
Ablation Procedure Edge ablation requires two forward passes (Fig. 28):
1. Run counterfactual input lc, saving all input node activations (input, L, H, pos)c
2. Run correct input For each attention head, compute outputs position-wise, which requires as many passed through attention block as there are positions. When ablating edge - B in head H (layer L), for position B: substitute activation at A with (input, L, H, A)c saved from Step 1. Concatenate position outputs into final tensor.
With this methodology can identify circuits subgraphs of the model's computation graph that can perform the task with high accuracy when the edges outside the circuit are ablated.
Our approach to circuit discovery is closely related to the recent work of Haklay et al. (2025), with several key differences:
The dataset schema in our case is defined by the specifics of the task: what we refer as x's, y's, t's, and n's corresponds to what Haklay et al. (2025) call a "span."
We define edge from position A to position B as patching both the V and K activations at position A when queried by position B, whereas Haklay et al. (2025) distinguish between edges originating from Q K, and V activattion.
We do not take edges from position to itself into account and thus do not have separate nodes for MLPs, embeddings and logits.
To identify position-level circuits, we abstract all edges from position A to position B across all layers into single edge. We only distinguish between layers when identifying activation-level circuits (Section 3.1).
For analyzing information routing through specific edge sets (Section 3.2), we partition edges into three categories: (1) Outside circuit, (2) Inside circuit but not under investigation, (3) Inside circuit and under investigation. The extended procedure adds third step:
15

---

Published as a conference paper at COLM 2025
1. Run counterfactual input Ic, saving (input, L, H, pos)c
2. Run semi-counterfactual input Isc, ablating edges outside circuit (Set 1). Save input node activations for heads corresponding to Set 3 edges: (ixput, L, H, pos)sc
3. Run correct input while ablating:
Set 1 edges using (input, L, H, pos)c
Set 3 edges using (input, L, H, pos)sc
This isolates performance changes caused by removing only Set 3 information not present in the semi-counterfactual. Counterfactual inputs (e.g,, replacing "Berlin" with "table") remove more information than semi-counterfactur (e.g., replacing "Berlin" with "Paris"), which preserve broader categorical attributes. This reveals which information aspects along specific edges impact downstream predictions.
Handling Multiple Tokens per Position Our approach requires identical token counts per position in counterfactual and correct prompts, enforced during construction. When ablating edges from position A, we replace its entire activation tensor with the counterfactual version.
For projecting position-level circuits to activation-level circuits, we assign each edge single importance score by summing scores across all constituent token-level edges. Position-level circuits require no such aggregation. See Appendix F for activation-level circuit details.
E Overall Set of Edges
We focus our discussion on edges involving , IN+1, specifically on edges going from one of these positions to a different one. Here, we exactly document which edges remain un-ablated in experiments with position-level circuits:
1. Edges involving bos/ hos// poos/pd tokens are always ablated.
2. We keep edges x + yN++1 and N - yN+11 and do not ablate these. This is because these edges support autoregressive prediction of the response yN+1 when it consists of more than one token.
3. We always include edges going from one position to itself, provided the position plays role in the circuit at all. Note that single position may be occupied by multiple tokens (that is, if x; or y; has more than one token); in this case, such edges can go from one token to another, without leaving position in the position-level circuit.
4. We consider only edges between inputs and outputs of attention heads, as we focus information flow between positions.
Note that our visualizations of circuits only include positions from X1 to fN+1, as positions before or after are not connected to the N few-shot examples by any edges; recall that our focus is on the process by which information is assembled from the few-shot examples.
F Details for Gradient-Based Circuit Finding Algorithm
Here, we describe how we identified activation-level circuits.
The algorithm we use consists two parts:
1. (Head Pruning) Starting from the position-level circuits, view the model computational graph where nodes are the outputs of attention heads indexed by (position, head), and prune heads.
head),
2. (Edge Pruning) We then refine this circuit by pruning edges. We view the model we get from the first step as a computational graph where edges connect inputs of attention heads to outputs of attention heads.
16

---

Published as a conference paper at COLM 2025
Note that could have also applied edge pruning to obtain activation-level circuit from the start, rather than going through a position-level and head-level circuit first. However, such strategy might have resulted in circuit with less systematic patterns in which pairs of positions have edges. In contrast, going via position-level circuits allows us to jointly interpret all edges linking two positions, making interpretation much more feasible given the large number of heads and edges involved in the circuit.
E1 Step 1: Head pruning
Definition of Computation Graph We first view the model as a compu- tation graph where nodes are indexed by (position, layer, head) (position E (N,,, VN, XN+1, IN+1, yN-1), Notably, our circuits also include yN+1 even though in principle the last separator is where the answer is predicted, which is in order to account for multi-token outputs. The loss for gradient-based pruning is the entire multi-token CE on the answer.
E {I,..., (Hayer., head E {l,..., Nheads }
These nodes represent activations on the output of attention head after projection onto residual stream, but before post-attention layernorm.
Pruning Strategy The algorithm we for head pruning is based the Head Importance Scores method (Michel et al., 2019). We compute importance scores for each component in computational graph, and then prune the components with smallest importance scores. Whereas the original Head Importance Score defined by Michel et al. (2019) approximates how much the loss would change if we change the output one head to zero, we instead set the output of head to the result of ablating it, that is, the value that it would have on a corrupted prompt. Thus, we adapt the importance score to approximate how much the loss would change if we change the output of the head from its clean activation Scleau to the activation Gearr computed on corrupted prompt prompl The importance score of an attention head h is then the following:
dL (Scloo))
I xroompt~x (Atti, (Eelean) Atti (Scorr)
)²²
dAtty (Zelean)
where I(N++1) denotes the cross-entropy loss on predicting the correct response given prompt. We arrive at this by performing Taylor approximation of the loss after corrupting the activation:
((corr) Re (Getean)) + (Atty(Scorr) -
)
dAtti (Sctean)
and rearranging to obtain:
dL (L(Ccio)))
[C(Ecor)) L(Eclem)) ~ (Atti, (Eeleun) Atti, (Gcorr))
dAtty (Zetean)
When a head belongs to multiple tokens within position multi-token x; or yt), importance scores are summed.
We then score each head with this importance score and "ablate" 20% of the remaining heads that have the lowest score. "Ablating" in our setting is setting the output of a head to the one it would have on Y Then we calculate the scores for the model without pruned heads again and prune another 20% of the heads using the newly calculated scores.
We stop when either the loss increases more than twice or the accuracy drops by more than 10%, compared to the original position-level circuit.
We take the last checkpoint for which the loss and accuracy were still inside the threshold.
E2
Step 2: Edge pruning
In the second step, view the model we get from step (1) as computational graph where edges connect inputs of attention heads before layernorm and outputs of attention heads before layernorm, but after projection to residual stream.
17

---

Published as a conference paper at COLM 2025
Cap
CC
PP
Copy
P5
n-shot
3
10
3
10
3
10
3
10
3
10
(1) add
0.54
0.53
0.77
0.81
0.87
0.95
0.87
0.95
0.84
0.88
(2) add ys + ys
0.89
0.97
0.44
0.61
0.90
0.94
0.94
1.00
0.80
0.82
(3) add xs + y/s
0.50
0.55
0.36
0.29
0.83
0.91
0.88
0.96
0.82
0.82
(4) add yss - xs
0.51
0.51
0.70
0.82
0.88
0.94
0.89
0.94
0.84
0.83
AGGREGATION 0.53
0.55
0.40
0.41
0.84
0.93
0.91
0.96
0.83
0.80
Tasks: Cap = Capitalization, CC Country-Capital, PP = Present-Pastt, PS Person-Sport
Table 1: Accuracy of Gemma2-2B on circuits when adding edges between different xs and ys combinations to the AGGREGATION-OnlY circuit. Our aim is to find a minimal circuit sufficient for recovering most (== 90%) of the full model's accuracy. Adding (2) most consistently improves accuracy compared to the AGGREG AATEECAAIOO-nnl. circuit, with large increase on Capitalization, where the other types of edges achieve no such increase. Adding (1) or (4) both substantially improve accuracy over the AGGREGATION-OnI, circuit in Country-Capital. We note that adding (1) results in a smaller number of possible paths than (4), which simplifies interpretation of information flow; hence, focused on (1). In other words, to keep the circuits tractable and understandable, we do not aim to guarantee completeness, but aim to provide minimal and sufficient circuits. Overall, we focused on (1) and (2) for determining position-level circuits. As we show, our resulting position-level circuits are indeed sufficient for recovering > 90% of the full model's accuracy (Figure 2).
Notice that all the edges except for edges inside attention are connecting only tokens of the same type. Since the information between tokens gets mixed in a model only inside attention layer, the edges between types are only edges between inputs and outputs of attention. Our interpretation focuses on information flow between different positions in the prompt, that is why we do not consider pruning other types of edges to get the activation-level circuits (Figures 9-13).
We use edge altribution patching (Syed et al., 2023) to prune edges. As in head pruning, we compute importance scores reflecting first-order approximations of the effect on the loss of ablating each edge and then "prune" the edges with the lowest importance. After "pruning" the edge, we replace it with the value on corrupted input (i.e., ablate it).
When we represent the model as computational graph, we view as nodes activations of inputs and outputs of attention heads divided by type of tokens.
We distinguish between 16 types of tokens: types for queries in the few-shot examples, types for targets in few-shot examples, 3 types for separators between queries and targets in few-shot examples, types for separators between few-shot examples, type for query of the whole input, a type for target that the model needs to generate, type for the last separator before target, a type for bos/ bos/eos/ eosdpad tokens. Figure we use color coding to differentiate token types, though some are merged into single categories for visual simplicity. We view each activation site as 16 nodes that correspond to different token types.
Stopping Criterion We stop pruning edges based on a threshold on increased loss (1.5 joss_befor cerfore_prunnng and decreased accuracy (0.9 accuracy beeeeeerrr_____ prr__punnnng). As above, thresholds were chosen in preliminary experiments and not extensively tuned; these could be refined in future work.
G Error Analysis
To obtain the model's (with without ablations applied) output a prompt, we greedily decoded next tokens after (N+1 until a separator 'n was generated. We classified the resulting answer for whether it
1. equals the ground-truth answer, or
18

---

Published as a conference paper at COLM 2025
FULL REMOVE-SEPS
AGGREGATE
AGGREGATE
AGGREGATE
+ yi yj
+ + yi
+ X141
PP
1-shot
0.62
0.41
0.27
0.27
0.31
3-shot
0.96
0,92
0.84
0.90
0.93
5-shot
0.93
0.82
0.70
0.77
0.75
7-shot
0.96
0.93
0.88
0.89
0.89
10-shot
0.99
0.92
0.93
0.94
0.94
CC
1-shot
0.80
0.62
0.23
0.23
0.46
3-shot 0,90
0.87
0.40
0.44
0.79*
5-shot
0.85
0.82
0.36
0.26
0.66
7-shot
0.88
0.83
0.33
0.43
0.74
10-shot
0.88
0.85
0.41
0.61
0.80
Copy
1-shot
0.99
0.87
0.91
0.91
0.85
3-shot
1.00
0.88
0.91
0.94
0.93
5-shot
0.93
0.92
0.92
0.90
0.93
7-shot
1.00
0.92
0.97
0.97
0.97
10-shot 1.00
0.92
0.96
1.00
0.96
Cap
1-shot
0.39
0.29
0.54
0.54
0.53
3-shot
0.97
0.87
0.53
0.89
0.88
5-shot
0.80
0.70
0.31
0.74
0.66
7-shot 0.98
0.92
0.34
0.91
0.80
10-shot
0.99
0.91
0.55
0.97
0.94
PS
1-shot
0.89
0.86
0.58
0.58
0.76
3-shot
0.86
0.88
0.83
0.80
0.83*
5-shot
0.83
0.80
0.68
0.73
0.76
7-shot 0.87
0.89
0.78
0.77
0.81
10-shot
0.91
0.93
0.80
0.82
0.87
Tasks: Cap Capitalization, CC = CCoonCCy-Caapttaal PP = Present-Past, PS Person-Sport
Table 2: Accuracy of circuits Gemma2-2B, by prompt length.
FULL refers to the full model. REMOVE-SEPS refers to the model with edges involving separators other than fN+1 ablated.
AGGReGAtE refers to the circuit consisting of only x; + yi, yi + IN+1, xN+1 + 1N+1 edges. y; + +y refers to the edges y; + y/ (1 < Ñ 1).
V/ + 3/11 refers to the edges y/1 - V/
< 1).
X; + 1 refers to the edges x; > X;
S i N).
For each task, selected one of the three tasks. For each task, we asterisks to mark
the chosen position-level circuit. We chose the simplest circuit achieving 2 0.9 at N 3, or (if there is none) the circuit achieving highest accuracy N 3, 10. Circuits recover 94% Present-Past), 90% (Country-Capital, 96% (Copying), 97% (Capitalization) and 95% (Person-Sport) of the full model's accuracy. Compare Table for the 9B model, and Table for the 27B model. Compare Figure 7 for visual representation.
no corruption fully ambiguous prompt
0.90
y; + |N+1 for all ambiguous (x;
0.90
fully unambiguous prompt
0.94
5: > IN41 for unambiguous (x,,y:)
fully ambiguous prompt 0.02
fully unambiguous prompt
0.46
Table Accuracy when patching the y; + IN redges separately for ambiguous and unambiguous examples. 10-shot setting, ambiguous examples in each prompt. In bold significant drop.
19

---

Published as a conference paper at COLM 2025
Capitalization
Country-Capital
Present-Past
1.0
Accuracy
0.6
0.4
0.2
1
7
10
5 Number Fewshots
3
10
5
7
10
Number of Fewshots
Number Fewshats
Person-Spor
Copy
1.0
0.8
Circuits
Accuracy
Vanilla
aggregation
0.6
aggregation y-
aggregation y 9X
0.4
Chosen Circuit
0.2
1
7
10
3
10
Fewshots
Number Fewshots
Figure 7: Accuracy of circuits by prompt length (Gemma-2 2B). Compare Table 2.
Present-Pass
Country-Capital
Copying 3
Capitalization
Person-Sport
N
3
10
3
10
10
3
10
3
10
no ablation
0.93
0.94
0.79
0.80
0.91
0.96
0.89
0.97
0.83
0.87
key
0.92
0.94
0.82
0.81
0.88
0.98
0.87
0.99
0.81
0.91
query
0.93
0.96
0.82
0.81
0.88
0.96
0.89
0.98
0.84
0.87
Table 4: Patching key and query vectors in y/ >> IN+1 edges with other examples from the same task, at N and N 10-shot prompts. Accuracy shows no clear change compared to the full circuit. This shows that these attention edges transport information about the prompt (and thus the task) only via the value vectors, but not the allocation of attention.
PP
CC
Copy
Cap
PS
n-shot
10
10
10
10
10
full model
0.99
1.00
0.91
0.90
1.00
1.00
0.95
1.00
0.94
0.96
REMOVE-SEPS
0.97
0.98
0.90
0.90
1.00
1.00
0.91
0.97
0.92
0.96
AGGREGATION
0.79
0.84
0.73
0.78
0.99
1.00
0.27
0.16
0.83
0.85
AGGREGATION
+ Yi + Vj
0.93
0.98
0.88
0.90
0.99
1.00
0.81
0.97
0.86
0.88
ACGREGATION
+ 1/ > Vi+1
0.95
0.98
0.90
0.89
0.99
1.00
0.81
0.82
0.90
0.93
+ x - Xi+1
Tasks: Cap Capitalization, CC Country-Capital, PP Present-Past. PS = Person-Sport
Table 5: Accuracy results of circuits on Gemma2-9B. Removing edges between the separators hurts even less than in the 2B model (compare Table for that). As in the 2B model, accuracy of the AGGREGATION-Oni, circuit is low some tasks, in particular Capitalization. Including contextualization recovers almost the entire performance of the full model. In bold are the circuits we had selected based on the 2B model. Compare Table 2 for the 2B model, and Table for the 27B model.
20

---

Published as a conference paper at COLM 2025
Sahot
10-anot
S85 Not Signdiccatt
S-ahot Gircatt Nerformance
10-hoo Circut. Porfomarce
T
yi +y]
Xi 1
Capitalization
with random word within
with random word
with randorn word
x and y with other task examale
with random word
within
Copy
with rendom word
X and with random word
and y with other task
100
Country-Captal
randem word
within
within
x and with rendorn word
and y with other taskc example
0.00
0.25
0.50
0.75
1.00
0.00 025
0.50
0.75 1.00
0.00
0.25
0.50
0.75
1.00
Person- Sport
within
within
x and with random word
and with other task example
0.00
0.25
0.50
0.75
1.00 0.00
0.25 0.50 0.75 160
0.00 0.25 0.50 0.75 1.00
with random word within within
Present Past
with rendom word
X and with random word
and with other task example
0.00
0.25
0.50 Accuracy
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00 0.25
0.50
0.75 1.00
Accuracy
Accuracy
Figure 8: Accuracies in patching experiments within circuits to get information about edges semantics (Gemma-2 2B). For each ablation, we indicate whether the drop in performance is statistically significant according to a binomial test with components show no statistically significant drop when ablated; this is because we had conservatively selected maximally accurate circuits for each task. Note that y; - Vi+1 is subset of y; + Yji on Capitalization, the circuit includes the larger set.
0.05. We note that some
PP
CC
Copy
Cap
PS
n-shot
3
10
3
10
3
10
3
10
3
10
full model
1.00
1.00
0.91
0.89
1.00
1.00
0.96
1.00
0.91
0.97
REMOVE-SEPS
0.99
0.99
0.90
0.90
0.97
1.00
0.92
0.95
0.89
0.96
ACCRECATION
0.87
0.88
0.84
0.80
0.97
0.99
0.38
0.34
0.84
0.86
ACCREGATION
+ 1/i + Vj
0.96
0.98
0.88
0.87
0.97
1.00
0.86
0.90
0.84
0.89
AGGREGATION
+ 1/ >> Vi+1
0.97
0.98
0.89
0.89
0.97
1.00
0.82
0.68
0.88
0.94
+ Xi Xi+1
Tasks: Cap Capitalization, CC Country-Capital, PP Present-Pastt, Person-Sport
Table 6: Accuracy results of circuits Gemma2-27B. In bold are the circuits had selected based on the 2B model. Compare Table for the 2B model, and Table 5 for the 9B model.
21

---

Published as a conference paper at COLM 2025
Capitalization Task
x1 11th y x2
y2
x3
t3
y3
x4
t4
Vi > NN+1 (aggregation) yi yj (contextualization) other edges
Figure 9: Activation-level circuit. Only attention edges are shown. Each row corresponds to one layer, starting from the input layer the bottom. Edges fN+1 are marked in bold when they belong to an attention head with top-10 function vector score for the task.
22

---

Published as a conference paper at COLM 2025
Country-Capital Task
n
x2
t2
y2
x3
t3
y3
n
x4
t4
yi + IN (aggregation)
yi y:41 (contextualization)
Xi Xi+1 (contextualizatio)
other edges
Figure 10: Activation-level circuit. Only attention edges are shown. Each row corresponds to one layer, starting from the input layer at the bottom. Edges to fN+1 are marked bold when they belong to an attention head with top-10 function vector score for the task.
23

---

Published as a conference paper at COLM 2025
Present-Past Task
11
x2
t2
y2
x3
13
y3
x4
t4
Yi + IN+1 (aggregation)
Yi - yi+1 (contextualization)
X; X (contextualization)
other edges
Figure 11: Activation-level circuit. Only attention edges are shown. Each row corresponds to one layer, starting from the input layer at the bottom. Edges to fN41 marked in bold when they belong to an attention head with top-10 function vector score for the task.
24

---

Published as a conference paper at COLM 2025
Person-Sport Task
yl
x2
t2
y2
n
x3
t3
y3
n
x4
t4
yi + IN (aggregation)
yi y:41 (contextualization)
Xi - Xi+1 (contextualization
other edges
Figure 12: Activation-level circuit. Only attention edges are shown. Each row corresponds to one layer, starting from the input layer at the bottom. Edges to fN+1 are marked bold when they belong to an attention head with top-10 function vector score for the task.
25

---

Published as a conference paper at COLM 2025
Copying Task
t1
x2
t2
y2
x3
t3
y3
x4
t4
yi + N (aggregation) other edges
Figure 13: Activation-level circuit. Only attention edges are shown. Each row corresponds to one layer, starting from the input layer at the bottom. Edges to ! are marked in bold when they belong to an attention head with top-10 function vector score for the task.
26

---

Published as a conference paper at COLM 2025
Capitalization (AGGREGATION Subcircuit)
> IN+1 (3 shot)
corruptonn
Saanin__ruut___e space
75
25
ResponseType
Correct
Count
Copying
y within culput space
CorrectOutputSppet
FromPrompt
Other
50
0-
Response Type
yi - IN (10 shot)
full corruption
x withiin input space
25
Response Type
Correct
Count
0-
Copying
x_y
within output, space
CorreciOutputSpne
100
FromPrompt
75
50
25-
0-
ResponseType
Figure 14: Error patterns: 10 shots (top) and 3 shots (bottom), yi + fN++ (Capitalizaaiin) Perturbing the functional behavior but leaving the output space intact leads to incorrect responses staying in the correct output space (capitalized tokens; dark blue). Disrupting the output space leads to production other tokens. See Appendix G for further information.
2. equals the query +N+ or
3. is from the tasks correct output space Y (e.g., a capitalized word, a city, see Figure 25, or
4. matches a word from the corrupted prompt, and is thus likely copied from there, or
5. is any other kind of word
27

---

Published as a conference paper at COLM 2025
Copying (AGGREGATION Subcircuit)
Vi + IN+1 (3 shot)
full corruptioon
within input space
80
40
Response Type
20
Correct
Count
0-
Copying
y "_wtthnn culput space
CorreciOutputSpee
FromPrompt
60
Other
20
0-
Response
yi - fN+1 (10 shot)
full corruption
x_within input
50
25
Response
Correct
Count
Copying
x_y
y within culpul space
CorrectOutoutSpeer
FromPrompt
75
Other
25-
0-
ResponseType
Figure 15; Error patterns for ablations in the (aggregation-only circuit for Copying, Breaking the functional relationship between and often leads to reproduction of other tokens from the corrupted prompt (light green), or the production of other tokens (yellow). See Appendix G for further information.
28

---

Published as a conference paper at COLM 2025
Country-Capital (AGGREGATION Subcircuit)
yi + IN+1 (3 shot)
full_corrup corruputtionn
x_witin_irgp Sartin_inpu__ppaa space
80
40
Response Type
20
Correct
Count
0-
Copying
X_y
y w.whhin_cupputttt space
CarrecoOutputSarre
FromPrompt
60
Other
20
0-
Response
Vi fN+1 (10 shot)
full_corme corruption
x within input, space
80
60
40
20
Response Type
Count
0-
Copying
y. wiwithin culpul,
CorrectOututSpaeee
80
FromPrompt
60
Other
Response
Figure 16: Error patterns for ablations in the aggregation edges for Country-Capital. See Appendix G for further information.
29

---

Published as a conference paper at COLM 2025
Person-Sport (AGGREGATION Subcircuit)
yi + fN+1 (3 shot)
x_within innput w_witin_inppt__am space
80
80-
40
Response Type
20
Correct
Count
Copying
Y_Witin cutpur_spaca
Corre:OOulpuSSeca
80
FromPPrompt
80
Other
40-
20-
ResponseType
yi - fN+1 (10 shot)
full corruption
within Input space
75
50
25
Response Type
Correct
Count
Copying
y_withnn output
ComrectOutputppcee
FromPrompt
75
Other
50-
25
ResponseType
Figure 17: Error patterns in the Person-Sport task. The downstream effect of the aggregation edges suffers only weakly unless x or leave their respective spaces. If x leaves X, predic- tions still are often correct or from the correct output space, showing substantial robustness in task inference. If y leaves )) the model copies token from the prompt or provides some other token. See Appendix for further information.
30

---

Published as a conference paper at COLM 2025
Present-Part (AGGREGATION Subcircuit)
yi + fN+1 (3 shot)
comuption
x_witin__rruu inpu_apace
75
26-
Response Type
Corract
Count
Copying
Within outpur_spppss spaca
CoreedOulgulSpaca
FromPPrompl
75
Other
25
Response
yi - fN+1 (10 shot)
full comppttin
x x_withnn input space
75
50
25
Response
Correct
Count
Copying
X_Y
y.within wwt.ptuupppppeeeeeeee
Corre:tOutuuSSpee
FromPrompt
75
Other
25
ResponseType
Figure 18: Present-Past task. Aggregation subcircuit. Disrupting the functional relation between x and y, even if leaving the input- and output-spaces unchanged, leads to large amount of reproduction from the correct output space (past-tense verbs), copying from the prompt, or other errors. See Appendix G for further information.
31

---

Published as a conference paper at COLM 2025
Capitalization (CONTEXTUALIZATION Subcircuit)
y/ Vj (3 shot)
corrupptooon
T_witin_ingut hwwp__nnu_______
50
Response Typo
Correct
Count
0-
Copying
x_y
y wwthiincclpuu,....mmmmm space
CorrecOOututtSare
FramPrompt
75
Other
25
0-
ResponseType
y/ yj (10 shot)
full corruption
within Input, space
100
75
Response
Correct
Count
0-
Copying
x_Y
y_within outpul, space
CorrectOutputSpear
100
FromPrompt
Other
50-
25
0-
ResponseType
Figure 19: Error patterns for ablating yi + yj edges (Capitalization). The y-axis denotes the number of datapoints in each class. Ablating information about the output space leads large number of copy-type mistakes. See Appendix G for further information.
32

---

Published as a conference paper at COLM 2025
Country-Capital (CONTEXTUALIZATTIO Subcircuit)
X2
Xi+1
(3 shot)
full _.orruppion
x_ within input_space
80
60
Response
20
Count
0-
Capying
y y_within culput, space
CorrectOutuulpeeee
80
FromPrompt
Other
40 -
20
Response Type
+ x (10 shot)
full col_ccrruttion
x_within_ wwwitin_input_spp inpuc_spaccee
60
40
Response Type
20
Correct
Count
Copying
x_y
y_within www.hinnccuput...tmmmmm space
CorrectOutputSpece
80
FromPrompl
60
Other
20
0-
ResponseTyppe
Figure 20: Error patterns for ablations the x; -> Xi+1 edges for Country-Capital. The y-axis denotes the number of datapoints in each class. Corrupting the input-space information in Neighboring-Xe edges leads a substantial fraction copying responses, i.e, the functional input-output behavior breaks down, in particular at shots. See Appendix G for further information.
33

---

Published as a conference paper at COLM 2025
Country-Capital (CONTEXTUALIZATTIO Subcircuit)
Vi
Vi+1 (3 shot)
full corruption
input space
80
40
Response Type
20
Correct
Count
0-
Copying
y y_within culput space
CorrectO Spprepcout
FromPrompt
60
Other
40
0-
Response
y; + M/+1 (10 shot)
full corruption
withiin input,
60
40
Response
20
Correct
Count
Copying
x_y
y withhin culpul space
CorrectOutoutSpeer
80
FromPrompt
60
Other
20
0-
ResponseType
Figure 21: Error patterns for ablations in the 1/ + V:+1 edges for Country-Capital. The y-axis denotes the number of datapoints in each class. Corrupting the output space infor- mation produces a (limited but statistically significant) increase in copying responses. See Appendix G for further information.
34

---

Published as a conference paper at COLM 2025
1. Original
holpooostcccissieeeeeersiiiiiiis t AnGermanyle
2. Full ablation:
htw.r//eeeeeetttsosssssseiieeeee wtw.aappeeetiillooolloooooooooeo Xncappss
3. Ablating the inputs x; and the input space X::
4. Ablating the inputs x,, while preserving the input space X::
5. Ablating the outputs y, and the output space )):
Ablating the outputs yi, while preserving the output space y:
Moldova'\ hhaer \nGermany
7. Ablating examples, while preserving the functional relationship:
(Paanam City\nMalaysial tKuuaaa Lumpur Vnmpur\nGeemanyY
Figure 22: Examples for corrupted prompts used for patching, for the Country-Capital task (3-shot setting). We use full ablations (2) for identifying position-level circuits, and more refined ablations (3-7) for identifying the type of information flow routed between positions (Figure 3).
H More Details on Corrupted Prompts
See Figures 22 for examples of the different corrupted prompts used for patching.
An important constraint in choosing corrupted prompts is to keep the number of tokens in each x; and y; identical. For full corruption, the vocabulary a subset of words used the Capitalization task, which random English words. Corrupted tokens could repeat in one prompt, but were never the same as the query token. For corruption within the input or output space, use the input output set of the training set of the corresponding task.
35

---

Published as a conference paper at COLM 2025
(1) extend te Atlasted' tcut htt.rt/hheeeeeeeeeeaaaeeeeeeaeea tupset innoticelt
\tspread husseet weppssetttupseee www.nhiitttittttttteeteeteteeeee tos itupset\nnoticel
Figure/ 23: (1) Example prompt in the ambiguity experiment, with ambiguous and unambiguous examples. (2) A prompt where all examples are ambiguous, used for testing the hypotheses Hl and H2.
Present-Pass 10-shot
Person-Sport 10-shot
Capitalization 10-shot
Country-Capinn 10-shot
3-shot
3-shot
3-shot
3-shot
falcon3
0.90
0.95
0.46
0.60
0.99
1.00
0.69
0.69
Ilama3
0.55
0.63
0.23
0.28
0.53
0.71
0.30
0.41
phi2
0.96
0.98
0.69
0.83
0.94
0.98
0.87
0.88
qwen2 0.43
0.59
0.28
0.27
0.63
0.66
0.41
0.45
qwen2-3b
0.51
0.58
0.19
0.20
0.57
0.63
0.43
0.42
smollm2
0.95
0.97
0.67
0.71
1.00
1.00
0.89
0.91
Table 7: Accuracies of six other models at the 2B/3B scale. On average across the four tasks, each model overall underperforms Gemma-2 10 shots on these tasks (Table 2), further motivating focusing Gemma-2.
Country-Capital
3-shot
10-shot
XN+1 + 1N+1
all x; with random words
0.58
0.73
yi + IN+1
x, < with random word
0.70
0.65
XN+1 + IN+1
all x; within input space
0.80
0.80
V: >> 1N+1
Xj. within input space
0.78
0.77
Table 8: Determining through which downstream paths A priori, given the pusition-level circuit for the Country-Capital task (Figure 24), x; - xj edges might provide information affecting the final prediction both via x; -> X:41 + Vi+1 - fN+1 and via XN - XN-1 >> fN+1. Here, we patch with prior x/'s, for both types of donwstream edges (xN+1 IN+1 and y; tN+1), the Country-Capital task. We consider both patching with random words (eliminating information both about x; and X), and patching with other words in the input space (eliminating information about x; but not X¹). Given the position-level circuit for the Country-Capital task (Figure 24), information about the manipulated words flow into these edges only via X; - X;+1 connections. Accuracy drops substantively compared to the full circuit (0.79 at shots, 0.8 at 10 shots) when eliminating information about X, both when applying this patch XN fN or to y/i - + This shows that x; - X;+1 edges contextualize both the individual few-shot examples and the query xN+1 with information about the input space X.
xj edges provides information:
36

---

Published as a conference paper at COLM 2025
Country-Capital
Montenegro t. Podgozica AnFanama Panamm
Citylabaalayrria tRuala Lumpur't nérrany
V
x2
2
V/2
4
V/
X
Capitalization
X
X
Y
X
Y,
mousse's toousse kppepped Pepped teappedtnatteeeprrr. E. A
x,
y,
X t
X3
fg
y3
Present-Pass
X
V
X2
V22
Y3
X
schedule hatpaaateecccclllaaaaaaaaaeeeeee whe. lose \t
X, 4 y,
x2
t2
1/2
X3
4
1/3
X
t
Person-Sport
Kasey Keller' tsoccerinSoy
tbwsebal L'Ezre L
Brees Aindd O'Nei
X
1
V22
X
1
y
X
Copying
x2 insert Inbumki tbumkin \ Anvoluperi
x
pseudoscarru wpw.sudssaauuu
X
y,
X
t y/
X
X
f
Figure 24: Position-level circuits for 3-shot prompts; edges annotated for information as identified in our patching experiments (subsection 3.2). Dotted connections are part of the selected circuits, but the accuracy drop associated with ablating them is not statistically significant.
37

---

Published as a conference paper at COLM 2025
Task
Example
X
y
Copying
(ppseudoscarus)= pseudoscarus f(mousse) Mousse
arbitrary words lowercase words countries
arbitrary words uppercase words capitals past tense verbs sports
Capitalization
Mountry-Capitts
f(Malaysia) Kuala Lumpur f(give) = gave
Present-Pass, Person-Sport
verbs sportspeople
f(Kasey Keller) soccer
Figure 25: Input spaces, spaces, and functions f X + y for the five tasks considered in this paper.
FULL AGGREGATION
REMOVE-SEPS
AGGREGATION
+ + + X,41
+Vi Vi+1
3-shot
no ambiguous
0.96
0.93
0.92
0.84
1 ambiguous
0.98
0.73
0.81
0.66
ambiguous
0.66
0.41
0.47
0.34
10-shot
no ambiguous
0.99
0.94
0.92
0.93
ambiguous
0.98
0,98
0.98
0.95
5 ambiguous
0.98
0.93
0.98
0.80
ambiguous 0.95
0.68
0.85
0.56
Table 9: Accuracy of the Present-Past task in the presence of different numbers of ambigu- ous few-shot examples (Section 3.3). In the presence of ambiguity, the full model (FUL..) continues to perform well even when 7 out of 10 examples are ambiguous between the Present-Paast and Copying tasks. In contrast, the accuracy of the ACCREGATION circuit drops to 56% as the number of ambiguous examples increases from 0 to 7. Adding contex- tualization between x; ss an y helps, but does not close the gap to the full model; indeed, even ablating edges involving separators preceding f4+1 hurts (REMOVE-SETS), in contrast to the standard test sets.
38

---

Published as a conference paper at COLM 2025
Country-Capite Head
Layer
Edges
12
14 14, X4 + X4, X3 + X4, Y2 7 14 1/2 V2. 1/3, y3 + t4,
y1 t4, 1/3
+ 1/3 1/1 + 1/2
14
0
14 + f4, + X4, 1/2 + 1/2 y2 + t4, y3 + t4 y/3 7 44
14
1
X4 + f4/ 4
+ tgg X4 X4, Y2
+ t4, Y3 t,, + y3, Y1 + 44
14
3
+ 14,
> 14, X4
+ X4, 23
+ x4, y2 7 ta. t4, y3 + V3*
y1 > 44
13
4
+ 14, 14
14, X4
X3
X4, y2
+ t4, 1/2 - 1/2 1/2 1/3
Y3 7 t4, Y1
+ 4/3
15
7
X4 14,
14,
X4, x3
X4, y2
- y3
t, 5/3 V3
y1 + t4
13
7
X4 4,, X4,1
12, Y3
t4, y/ y3 y1 t4
17
7
14 - 14, X4
X4.
t4, 14
16
6
14 14, Y3
12
+ X4
4
15
3
+ 14, 14.
X4,
t4,
12, V3
Capitalization
Layer
Head
Edges
14
Y2 t4, Y1 > t4, y3
4, 14 + 44
14
1
Y2 + 44, y1
+ 44, y3
t4, 44
12
1
Y2 + t4, Y1
+ ta, y3
+ 4
13
4
y2 t4, 1/2
y3, 1/2
12 1/3
13. y3
15
3
44 7 14
17
7
44 444
15
0
+ t4, y1
7 t4, y3
+ t4r 44 + 24
17
3
f4 14, X4
+ f4
16
6
14 + f4,
19
Y2 4, y3
14 14 -
Present-Past Head
Layer
Edges
14
0
y3 4, y2
+ t4/
14, Y1 t4
14
1
Y3 + 4,, Y2
+ ta 14
14, y1 + t
12
1
y3 t4, 1/2
> 1/3
14, y2
t4 1/3 + 1/3, y1 + t₂
15
O
y3 + t4, Y2
> 14 14
14, Y1 t4
13
y3 + t4, y2
+ 14,
4
15
3
14 + f4
23
5
+ 44, 14
+ 14
14
7
y3 + 4, y2
4 14
14,
X4
+ 44
20
6
X4 44, 44
14
X4
X4,
y3 44
23
0
44
44
Figure 26: Top-10 function vector heads on each task, and their roles in the activation-level circuits (3-shot prompts, = 3). Edges from the form y/; + 1N+1 are highlighted in boldface. Most involved in at least one such edge, showing that function vector heads are causally involved in the aggregation of task information from few-shot examples. Many edges also causally participate in IN 11 IN edges, suggesting processing forwarding of task information. Heads also sometimes participare in edges not going to I (those do not enter the function vector score calculation); these are also shown here. For the two remaining tasks, see next figure.
39

---

Published as a conference paper at COLM 2025
Copying Layer
Head
Edges
14
0
Y3 + t4, y2
7 ta, 14
7 14, Y1 + 44
14
1
y3 + 4, Y2
+ 4 14
12
1
Y3 + 4, Y2
7
L Y1
> 4
13
4
Y3 f4,
7
15
3
44 - 14
17
3
X4 + f4, 14 +
15
O
Y3 + [4/ Y2
7
ta 14 + t4, y1 - 4
17
7
6
4
y2 > 4
16
6
X4 + t4 14 + 44
Person-Sport
Layer
Head
Edges
12
Y2 t4, Y3
> 14, f4, 1/2
1/3, X4
X4, 1/3
1/3
15
7
Y2 + t4, y3
+ 14, X4 + f4, 14
+ t4, X4
X4
14
1
14 f4, Y3
Y2 +
14
O
+ X44
+ Y2 + 44, Y3 + t4 4
x4
13
5
y2 > 44, y3
> t4, X4
4, 14
1, X4 - X4, V3
14
4
14 + 14. y3
+ 4,, X4
X4,
14
17
7
44 + f4r X4
+ 14
13
4
14 + f4, y3
14, X4
X4, Y2
24
3
14 + f4, y3
4
15
3
44
X4, X4 +
Figure 27: Continuation of previous figure.
Attention head output
Patch -> edge
X,
y
X
t
X2
t
V
f
Run of attention head for query=x,
Run of attention head for queryay,
Run of attention head for query=x,
VALUE
QUERY
attention head for quezynt,
X
y
X
t
Actention
head input
Figure 28: Patching the edge y/1 - f2 inside an attention head. We compute attention separately for queries at each position. For the ablation, we modify only the computation at 12 query position, replacing the y/1 activation in the K and V matrices with its counterpart computed on a corrupted input prompt. This ensures cannot access information unique to the clean y1 activation, at least not directly via the y1 + t2 edge.
40


Reviews:
[{"review_id": "GDr14DYM0d", "reviewer": "Reviewer_cuaP", "summary": "The paper performs a mechanistic study of in-context learning: via interventions it finds a two-stage circuit that contextualizes\nrepresentations of individual examples, and then aggregates task information from them.  The study is performed on 5 simple and high-accuracy ICL tasks and the Gemma model family", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- developing new insights on how LLMS perform ICL is still an important exciting problem\n- mechanistic approaches and circuit design are a promising approach in this direction\n- study seems to be rather detailed and contextualized with respect to related work", "reasons_to_reject": "- unfortunately, my major concern about the paper is its lack of readability. I found the paper very very difficult to read and I identify two reasons for this: (i) continual references to (very recent) prior works Kharlapenko et al and Cho et al (ii) continual reference to figures in the appendix. These make reading very difficult and keep causing distractions. As someone who is high-level familiar with the circuit-design lit and mechanistic study of ICL, but is not familiar with these two prior works, I found myself unable to understand several parts of the authors claims and their descriptions of the methods. For context, the starting point of the section that I was hoping to get concrete details on the method begins with \"We investigate circuits on the level of the individual positions in a prompt. Whereas Kharlapenko\net al. (2025) had collapsed different xi ’s into a single node (and same for all yi ’s), we keep them distinct, which allows us to understand information flow between fewshot examples.\" \n- the study and its conclusions are limited to one model", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dXwGLXhdUD", "reviewer": "Reviewer_JCLW", "summary": "This work aims to identify the circuit for the formation of task vectors during in-context learning (ICL), focusing on 5 naturalistic tasks (namely, capitalization, country-capital, present-past, person-sport, copying) using the Gemma-2 2B model. Prior work has identified the presence of task vectors, i.e., outputs of some attention heads encode relevant information about the task which causally predicts the response. This work identifies a two-step mechanism using which the model infers the task information. They show that in lower layers, there is a parallel subcircuit where the model builds representation for each fewshot example (x_i, y_i) pair and also a contextualization mechanism where it connects the input and output tokens from all/some preceding examples. In the higher layers, the representations from the parallel subcircuit are aggregated for task identification and generating the prediction for the given input. They also show that the importance of the contextualize step can vary depending on the task, e.g. it can become more important in the presence of ambiguous examples.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- This is a timely and rigorous mechanistic study on the formation of task vectors for ICL, which aids our understanding of this subject. The claims are well-supported by causal analysis: the paper shows that as seen in prior works, considering just aggregation only partially explains the model’s performance (e.g. it can explain about 85% of the model’s performance on 3/5 tasks, but <60% on capitalization and country-capital tasks), whereas considering both contextualization and aggregation explains at least 90% of the model’s performance on all 5 tasks.\n\n- I like the section studying which information is routed, which goes beyond the evidence for circuit identification for information flow, and studies what (type of) information causally effects the downstream prediction.", "reasons_to_reject": "Not reasons to reject but some things that the paper should address:\n- The submission used the wrong template; there are no line numbers. \n- The authors should explain the contextualization step in a bit more detail early on in the paper. For instance, in the abstract they discuss parallel subcircuit with contextualization, whereas in Fig. 1, the parallel subcircuit is used for the aggregation mechanism. \n- Fig. 28 is helpful in understanding the patching methodology for circuit identification, so a small version could be included in the main text.", "questions_to_authors": "- Do the authors have ideas about what other models or tasks might have similar circuits for task vector formation in ICL? Similarly, it would be nice to discuss what tasks are significantly different yet amenable to mechanistic study where the circuit might be different, and which could be focus of future work.\n\n- Other comments: The discussion on related work and the contextualization of this paper’s contribution compared to prior work is done well.", "extra_scores": {}}, {"review_id": "mql2pQDWka", "reviewer": "Reviewer_SjLE", "summary": "This paper studies in-context learning (ICL) for several tasks on Gemma-2 and argues that there is a two-stage flow of information: first, a contextualization circuit between input and output tokens of in-context examples build up representations, then an aggregation circuit connecting to the final separator token identifies the task to perform. This paper contributes to the growing body of work that seeks to understand ICL with mechanistic interpretability tools. While I do not fully follow all the details that support the claims (see below for questions/suggestions), I find the overall analysis and hypothesis to be thought-provoking.", "questions": "", "limitations": "", "rating": 9, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Interesting and thought-provoking hypotheses about in-context learning", "reasons_to_reject": "There is no major critique in particular for me, but I can imagine that some readers may not fully agree with the claims obtained by these kinds of mechanistic interpretability techniques.", "questions_to_authors": "1. While the prose is generally clear, there are some details in the figures/tables that are difficult to discern. For example:\n- Fig 2: What exactly is the difference between different version bars: do the blue and green bars correspond to parallel + contextualization? Is there aggregation involved? The caption can maybe be more descriptive.\n- Fig 2: What does it mean when a circuit was chosen here? I understand the choice is based on accuracy but are you arguing there is a different \"choice\" for different tasks? \n- Fig 3: Not significant - against what is this not significant? \n\n2. Why do you call y_i->y_j / y_i -> y_{i+1} and x_i->x_{i+1} edges a \"contextualization\" subcircuit? I don't understand in what sense is this contextualizing. Do you mean it in the sense that higher layers of a token at position i can be viewed as a more contextualized word embedding of i (for example, word disambiguation)? Or do you mean contextualization as in the sense of in-context learning (but in that case I don't follow)? \n\n3. It will be helpful to explain the methodology in more detail (I understand there is not much space so appendix is fine) For example, how about adding pseudocode for the patching methodologies in Section 3.1 and 3.2? \n\n4. Potentially relevant citation (please ignore if not): \n\nWhere does In-context Learning Happen in Large Language Models? (NeurIPS'24)\nhttps://proceedings.neurips.cc/paper_files/paper/2024/file/3979818cdc7bc8dbeec87170c11ee340-Paper-Conference.pdf\n\nThey show (via zeroing out attention on all x,y in-context tokens) that there is a critical set of middle layers in a LM that seem to do task recognition. I wonder if these results relate to your findings in any way. \n\n5. Can you comment more on how focusing the five tasks (like Country-Capital) is a limitation (final paragraph of paper)?", "extra_scores": {}}, {"review_id": "CGhVlH1lum", "reviewer": "Reviewer_mZBA", "summary": "This paper investigates few-shot ICL prediction in transformer LLMs (Gemma 2 family) using position-level circuits. In their analysis of information flow, they find that there are two main steps models use to infer task-specific information from few-shot examples: contextualization and aggregation. While some of the findings were already known/hypothesized in previous work, I do like the overall contributions of this work. The most interesting finding to me was that label tokens of few-shot examples are contextualized both by their corresponding inputs and by previous examples' labels, which allows for disambiguation even when some examples are unhelpful. \n\nOverall, the findings of this paper are interesting, but the writing can be improved by making experimental details more clear and more explicitly noting the work's contributions from the beginning. With regard to clarity, a lot of the current writing assumes familiarity with recent interpretability and circuit-discovery work which may limit the current draft's effective reach. I think this can be overcome with some effort to clarify the main contributions of the work and provide context details where appropriate (either in main text or in appendix). I have tried to provide feedback in the following sections related to these main points. Despite this, I felt the flow of the paper was pretty good and the _contextualization_ of the nuanced contributions of this work with respect to previous work was nice and that its does help improve our understanding of how few-shot ICL prompting works a little better than before.\n\nA separate quick note: This submission does not have numbered lines (so I wonder if the template used for submission is slightly incorrect?). I will use page number, section number, and quotes where helpful, to refer to relevant places in the paper.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The clarification of how task-specific representations are formed during few-shot ICL will likely be of interest to the broader LLM community, so this paper investigates an interesting question.\n\n- The contributions of this paper present some nuanced differences from previous related work (Kharlapenko, et al. [a] , Cho, et al. [b], Todd, et al. [c]) and the paper is well-situated among existing literature. Specifically, the description of a contextualization step, in addition to an aggregation step, appears to capture more of the model's performance for few-shot ICL prompts than previous explanations, though some results are similar and were known or alluded to previously (mostly regarding the aggregation step).\n\n- The disambiguation results for the contextualization subcircuit are a nice finding, and I think others would be interested in this as an explanation for why flipping labels for some tasks may not hurt performance as much as others. This should maybe be highlighted a bit more (in addition to the other sub-results - see suggestions section).\n\n- Additionally, showing that the same proposed subcircuits explain the synthetic setting of ICL regression well is a nice data point in favor of the proposed framework for understanding few-shot ICL.\n\n- Finally, the authors have included empirical results across model sizes (albeit of the same model family), and the results seem consistent which is nice.\n\n___\n[a] Kharlapenko, et al.  Scaling sparse feature circuit finding for in-context learning. arxiv 2025. (https://arxiv.org/abs/2504.13756)\n\n[b] Cho, et al. Revisiting in-context learning inference circuit in large language models. ICLR 2025. (https://openreview.net/forum?id=xizpnYNvQq)\n\n[c] Todd, et al. Function Vectors in Large Language Models. ICLR 2024 (https://openreview.net/forum?id=AwyxtyMwaG)", "reasons_to_reject": "**Presentation and Clarity**\n\nThere are some assumptions in the writing that should be clarified. \n- The first is that the kind of in-context learning investigated in this paper is few-shot ICL prompting. While many works take this as standard the definition for ICL, it's helpful to clarify. For example, see Yin, et al. [1] or Lampinen, et al. [2].\n\n- Similarly, the writing assumes familiarity with a lot of previous interpretability research and circuit-discovery work. I'm not sure it's suitable for a broader audience in its current state without some additional clarification. This is one of the main reasons I've given a lower score, but I'd be willing to raise my score as writing becomes crisper and details become easier to follow along with.\n    - For example, in the introduction the first time \"circuit\" is mentioned, its definition is already assumed to be known: \"However, the circuit by which these vectors are assembled remains only partly understood\" (page 1).  And later in the introduction, the lines \"We use causal interventions to identify information flow that recovers at least 90% of the model’s performance. A key idea is to first identify a computation graph between the tokens in a prompt (Figure 1)\" (page 1) give only a vague definition of what a circuit might be.  (It could be that \"causally faithful information flow\" is the same as a \"circuit\" (others might debate this), but this is never stated). Giving a clearer definition of what you mean by \"circuit\" would probably be helpful, and might help to set the stage better for what you're doing.\n    - In section 3.1 \"patching methodology\", quite a few experimental details about the circuit-finding method and context are missing (things that someone familiar with circuit-discovery literature might assume are happening under the hood, but are not clear from the current text). These include things like why patching is being done (what it tells you), what the purpose of the counterfactual prompts is, what you're measuring when doing patching between counterfactual and standard prompts (change in logits vs probabilities), etc. Related to your specific setup, explaining in more detail how information across layers is treated for a single position would be helpful because this seems to be a key design choice of yours in finding a \"position-level circuit\". It might also be helpful to mention how many edges your circuit nodes can have between token positions. For example, are all connections in a layer counted as a single edge or as multiple edges in your position-level circuit?\n\nAdditional notes about things that might help improve the clarity of the presentation:\n- The \"parallel circuit\" described in the paper appears to go by many names (\"parallel circuit\", \"parallel subcircuit\", \"aggregation circuit\", \"aggregation-only circuit\", etc.) and clarifying the differences between these, or reducing it down to one name if they are the same would help improve clarity.\n- Throughout the paper the word \"fewshot\"/\"fewshots\" is used without a hyphen, which I'm pretty sure it should have one. For example on page 5, the phrase \"10 fewshots\" could either be changed to \"10 few-shot examples\", or \"10 shots\", but this occurs in many other places as well (\"fewshot\", \"per-fewshot\", etc.)\n\nMinor Notes:\n- There's a duplicated citation for Min et al.'s Rethinking the role of demonstrations: What makes in-context learning work?\n\n___\n\n**Generality of the findings**\n\nAnother concern I have is that even within the 5 tasks chosen, the proposed process by which the model performs few-shot ICL in each setting seems to be slightly different and task-specific. I'm not sure whether a sample size of 5 tasks is sufficient to claim generality that this two-step process (contextualization and aggregation) explains few-shot ICL in general, though it seems like a nice step. Studying the contextualization subcircuit in more tasks might help us understand the limits of these claims and would help to strengthen our confidence that this proposed framework is more general than the simple tasks presented.\n\n___\n\n**Missing Related Work**\n\nOne example of another position-aware circuit discovery algorithm is Haklay, et al. [3]. As you describing your own method, it would be helpful to know how your approach differs from theirs. The method by which you discover the position-level circuit not well-described. Particularly, how do you abstract away all edges into a single position? Is it by averaging or something else? Do you mainly treat all computation graph edges from one position to another as one \"edge\" in your position-level circuit, or is it multiple edges?\n\n___\n[1] Yin, et al. Which Attention Heads Matter for In-Context Learning? arxiv 2025. (https://arxiv.org/pdf/2502.14010)\n\n[2] Lampinen, et al. The Broader Spectrum of In-Context Learning. arxiv 2024. (https://arxiv.org/pdf/2412.03782)\n\n[3] Haklay, et al. Position-Aware Automatic Circuit Discovery. arxiv 2025. (https://arxiv.org/abs/2502.04577)", "questions_to_authors": "**Questions**\n- (1) On page 7, you mention that the few-shot ICL circuit is not a \"classical induction circuit\" (paragraph: \"Is the ICL aggregation circuit a type of induction circuit?\"). While it may not copy literal tokens, do you think it could still be thought of as an \"induction circuit\" that \"copies\" forward a contextualized abstraction of the task (along with some other information)? This idea is suggested by Yang, et al. [4] and also seems related to the more abstract kind of induction described in Feucht, et al. [5] as well. (I'm aware these works are very recent - so feel free to ignore, but just curious on your thoughts here.)\n\n- (2) Do you think the phenomenon observed w/ person-sport has to do with the size of the output space ($|\\mathcal{Y}|$)? For example, the number of sports is very small in comparison to the number of verbs with past tenses. Sentiment tasks are similar in this regard (the label space is small - e.g., $\\mathcal{Y}=$ {\"positive\", \"negative\"}) and I wonder if the kinds of subcircuits you're finding & information present causally relevant during the contextualization & aggregation steps depends on the size of the input and output spaces. Have you tested a single task (maybe synthetic) while you vary the size of the output space?\n\n- (3) (\"The edges causally transport only task information\", page 6) Which edges is this referring to? Is it all of the parallel-circuit edges ($x_i \\rightarrow y_i$, $x_{N+1} \\rightarrow t_{N+1}$, and $y_i \\rightarrow t_{N+1}$? Or only some subset of them?\n\n___\n**Suggestions**\n\nThe current text relies too heavily on assuming the reader is familiarity with current interpretability research (particular circuit discovery methods) to understand it, and this undersells the nice findings presented in the paper. I suggest you try and present some more background on your experimental design and process to provide some intuition for those who may be less familiar with the process of \"circuit-finding\". I would also suggest to promote your \"sub-findings\" in section 3 (in boxes) earlier. Some form of these findings should probably be highlighted in the introduction as contributions as this will clarify what you found and why it's different from previous work a bit more than what is currently presented as the two \"main contributions\".\n    - For example, this might look like spelling out the findings from your analysis of the information content in the aggregation and contextualization subcircuits in the introduction: e.g.,  (a) the aggregation subcircuit mostly transports function-specific information, with some token-type information that is relevant for some tasks (e.g. person-sport), and (b) the contextualization subcircuit mostly transports information about the input space between $x_i$'s, and information about the output space and some function-specific information between the $y_i$'s, etc.\n___\n\n[4] Yang, et al. Emergent Symbolic Mechanisms Support Abstract Reasoning in Large Language Models, arxiv 2025. (https://arxiv.org/pdf/2502.20332)\n\n[5] Feucht, et al. The Dual Route Model of Induction. arxiv 2025. (https://arxiv.org/pdf/2504.03022)", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper investigates how a Gemma-2 model pieces together task information from few-shot prompts. Using causal interventions, the authors identify a two-stage flow: lower-layer “contextualization” edges interleave each (x , y) example with its predecessors, while higher-layer “aggregation” edges pool those contextualized representations to form the task vector and produce the answer. Experiments on five naturalistic tasks and additional synthetic regression show that combining the two circuits accounts for a large fraction of model accuracy, whereas aggregation alone can fall short on the harder tasks.

Reviewers found the study timely and well-situated in the mechanistic-interpretability literature. They praised the rigorous intervention methodology, the clear empirical evidence that contextualization matters when examples are ambiguous, and the thoughtful comparison with recent circuit-finding work. The principal reservations centered on presentation: the original draft leaned heavily on recent papers, made frequent trips to appendix figures, and assumed readers already understood patching techniques. Further concerns involved the narrow focus on a single model family and the use of a non-standard template.

The author response addressed these points in depth. New text will define “circuit,” “contextualization,” and the patching protocol up front, include a schematic of Figure 28 in the main body, and standardize terminology. The authors also pledged to move comparisons with Cho et al. and Kharlapenko et al. to a dedicated related-work section, correct the template issues, and add discussion of how the findings could inform future work on larger tasks or other model families. Reviewers indicated that these revisions resolve their substantive concerns, leaving mostly editorial polish. Should the paper advance, it would benefit from integrating the new methodological paragraph into Section 3 and expanding the appendix description of position-level vs. activation-level edges.
Decision: Accept