Contextualize-then-Aggregate: Circuits for In-Context Learning in Gemma-2 2B

Another 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.

We agree that this is a limitation of the paper, and that scaling the results to more difficult tasks is a nontrivial direction of future work. On the other hand, we also consider it a very nontrivial finding that some parts of the mechanism are shared between the tasks we used, particularly: the information that edges transport in different tasks, shared contextualization circuits between some of the tasks, cross-target-source edges ($x_i \rightarrow y_j$ and $y_i \rightarrow x_j$) being less important than other edges. Also, we believe it is expected that the exact mechanism is different for each task, and that putting them under one umbrella requires this umbrella to be quite abstract: which in our case is the concept of contextualization. Finding distinct smaller circuits for each task in our study goes somewhat beyond the abstraction, and this necessarily means that we identify more task-specific elements.

Missing Related Work: One 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.

Thank you for pointing out this concurrent work, which indeed proposes a method quite similar to the one we're using. The differences can be summarized as follows: (1) the schema of the dataset in our case is well-defined by the tasks specifics: what we call $x$'s, $y$'s, $t$'s, $n$'s is called a span in Haklay, et al. For automated circuit discovery, we sum scores for the same edge between different tokens, and average over dataset examples, same as Haklay, et al. (2) we define an edge from position A to position B as patching both V and K activations in position A when queried by position B, while Haklay, et al. differentiate between edges from Q, K and V activations. (3) We do not take edges from a position to itself into account and thus do not have separate nodes for MLPs, embeddings and logits in our graph. (4) In most of our experiments, we also abstract all edges between position A to position B in all layers into one edge to get a more interpretable mechanism. We differentiate between layers when finding the Activation-Level Circuits (Section 3.1).

We will add detailed discussion in Appendix E.

To specifically answer this question:

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?

We do not perform averaging within a position. Rather, as the reviewer suggests, we treat (in the position-level circuits) the set of all computation-graph edges from one position to another as a single edge of the position-level circuit. For position-level circuits, ablating a position-level edge from one position to another means ablating all computation-graph edges between this pair of positions; including the position-level edge means ablating none of these computation-graph edges. If any computation-graph edge plays an important role, we fully keep the overall position-level edge containing it. We will ensure that this is made explicit not just in the Appendix, but also in the main paper.

On page 7, you mention that the few-shot ICL circuit is not a "classical 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?

An induction-like mechanism might indeed still be happening, but our results suggest that it is not very selective to individual examples. Cho et al. argue that among the in-context examples the model chooses to copy only those that are most similar to the final query. We argue against this, as shuffling the K inputs in few-shot examples in these heads (which amounts to neutralizing attention weight differences between examples more or less similar to the query) does not affect model performance much. It may still be the case that $y_i\rightarrow t_{N+1}$ edges are a part of an induction-like circuit, but the sensitivity of this circuit to the specific query is not as strong as suggested by Cho et al. and various theoretical studies of ICL, as least in the model and tasks studied by us.

Do you think the phenomenon observed w/ person-sport has to do with the size of the output space?

This is a very reasonable idea, and we hypothesize this phenomenon could indeed be attributed to the size of the output space and semantic meaningfulness of the labels. However, we did not experiment with this. Designing such an experiment is not so trivial, because one needs some care in finding tasks that have high accuracy in 2B model, and finding a set of suitable tasks might be tricky. However, this is an exciting direction for future work.

("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->y_i$, $x_{N+1}->t_{N+1}$, and $y_i->t_{N+1}$? Or only some subset of them?

This refers to $y_i\rightarrow t_{N+1}$ only, we will clarify that in the text.

The 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.

Thank you for this suggestion. We will follow it, and spell out these findings already in the Introduction. We will also provide more background on the experimental design, especially in Section 3.1, to make the paper self-contained to a broader audience.

We are glad the reviewer found our work interesting and we thank the reviewer for all the constructive suggestions. We hope the following clarifications can address the reviewer's concerns.

The first is that the kind of in-context learning investigated in this paper is few-shot ICL prompting.

Thank you for pointing this out, we will clarify this in introduction.

(1) The writing assumes familiarity with a lot of previous interpretability research. (2) Giving a clearer definition of what you mean by "circuit" would probably be helpful. (3) In section 3.1 "patching methodology", quite a few experimental details about the circuit-finding method and context are missing.

We will add a paragraph in section 3.1 with information about patching setup and definitions, which we adopted from prior work. We provide the draft of the paragraph below:

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, a model is viewed as a computation graph with activations in different layers as nodes, and computations between them as edges. Then, it is possible to ablate some of the edges in the graph, forcibly replacing the computation along a specific edge with the one computed on the counterfactual input. Counterfactual inputs are designed to erase relevant information from the prompt while leaving other information intact for even better localization of model behavior. For example, when substituting "Berlin" with counterfactual "Paris", we isolate computations specific to the city's identity, not those activated by its category (city) or token-type (word). If ablating a set of edges does not lead to a 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 the circuits - subgraphs of the full model's computation graph that can explain a major part of model's performance on 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.

For readers unfamiliar with mechanistic interpretability, we hope this section offers an accessible introduction to patching. We will provide an overview of patching approach, a detailed description of our experimental setup and its distinctions from standard methodologies in the Appendix. This will include basics of patching implementation, a thorough description of Figure 28, the metric we're measuring when doing patching between counterfactual and standard prompts, and the details of how we collapse multiple edges into one.

Related to your specific setup, explaining in more detail how information across layers is treated for a single position would be helpful

We will adjust the second paragraph in section 3.1 to include this information. Please see the draft below:

We first discover a position-level circuit. Nodes are the positions in the prompt; edges are directed arcs 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, $x_2$ in a position-level circuit can only receive four incoming edges: from $x_1, t_1, y_1$, or $n_1$. Position-level circuits provide useful and interpretable upper bounds on information flow, though they do not distinguish between computations happening in different layers.

(1) This submission does not have numbered lines (2) The "parallel circuit" described in the paper appears to go by many names ("parallel circuit", "parallel subcircuit", "aggregation circuit", "aggregation-only circuit", etc.) (3) Throughout the paper the word "fewshot"/"fewshots" is used without a hyphen (4) There's a duplicated citation for Min et al.'s Rethinking the role of demonstrations: What makes in-context learning work?

Thank you for pointing this out. Indeed, we'll reduce references to "parallel circuit" to a single name, change the template, and adjust the text of the paper to include all of these changes.

We are glad the reviewer found our work thought-provoking and we thank the reviewer for all the constructive suggestions. We hope the following clarifications can address the questions the reviewer raised.

some readers may not fully agree with the claims obtained by these kinds of mechanistic interpretability techniques

While the methods we use are not perfect (we do not rule out the potential existence of multiple faithful, but distinct circuits, and our circuits recover close to, but not 100%, of the full model's performance), they are widely adopted in the literature and have become almost standard for this kind of work. Based on our experiments, we argue that the circuits we found are causal to the model's performance on the tasks, and we show causality via intervention in the model's behavior. Importantly, we avoid using correlational evidence, in order to make our analysis as rigorous as possible.

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?

These bars indeed correspond to parallel + contextualization; different colors correspond to different versions of the contextualization subcircuit. Aggregation in this context means the same as parallel circuit. We will update the legend to clarify that.

Fig 2: What does it mean when a circuit was chosen here?

We test two versions of the contextualization subcircuit, which involve different edges between few-shot examples. In one version, all edges between ys are involved ($y_i \rightarrow y_j$); the other includes only edges between adjacent ys ($y_i \rightarrow y_{i+1}$) and also edges between adjacent xs ($x_i \rightarrow x_{i+1}$). In principle, we could have defined the contextualization subcircuit as the union of all those edges ($y_i \rightarrow y_j$ and $x_i \rightarrow x_{i+1}$), but we decided to use smaller subcircuits per task to make further analysis more tractable.

As described in the paper, 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, based on 2B accuracies.

Fig 3: Not significant - against what is this not significant?

This means that the drop in accuracy when doing a specific ablation is not significant according to a binomial test with α = 0.05 compared to the full circuit performance. We will clarify it in the caption.

Why do you call $y_i\rightarrow y_j$ / $y_i\rightarrow y_{i+1}$ and $x_i\rightarrow x_{i+1}$ edges a "contextualization" subcircuit?

We mean that the mid and higher layer representations of the i-th few-shot example become contextualized, so that they reflect certain information about the other (preceding) few-shot examples.

For example, the higher layers of "Berlin" in a prompt "France\tParis\nBulgaria\tSofia\nGermany\tBerlin\nChina" in a contextualization circuit that includes edges $y_i\rightarrow y_{i+1}$, $x_i\rightarrow x_{i+1}$ can be contextualied with information from "Paris" or "Sofia", but not "France" or "Bulgaria". In a parallel circuit without contextualization, when the only edges we leave in a model are $y_i\rightarrow t_N$ and $x_i\rightarrow y_i$, higher layers of "Berlin" can not have any access to the information in previous few-shot examples (but it does have access to the information in embeddings of "Germany"), and thus its representaion is not contextualized with previous few-shot examples.

It will be helpful to explain the methodology in more detail.

We will extend section 3.1 to better describe the methodology we used, and also add a section in Appendix with more details about our setup and how it is different from standard patching. Apart from that, we will add a link to github with the implementation.

Currently, the details can be found in Appendix D (definition of computation graph), Appendix G (counterfactual inputs) and Figure 28 (patching edges between positions).

Potentially relevant citation: Where does In-context Learning Happen in Large Language Models? (NeurIPS'24)

Thanks for this pointer. We will add a citation and discussion. We believe that this finding is in good conceptual agreement with results about the role of function/task vector heads, which are in middle layers. Our work expands by clarifying the information flow leading up to the point where the task is recognized.

how focusing the five tasks is a limitation?

Our analysis focuses on five tasks, and we cannot claim that our findings will directly generalize to more advanced tasks (math, reasoning, etc.) used in real applications. We limit the study to tasks on which the model has high accuracy and therefore do not discuss tasks with lower accuracy. Further modifications of the circuits might be needed to extend our analysis to other tasks.

We thank the reviewer for all the suggestions and hope that additional clarifications make our findings more informative.

We thank the reviewer for all the constructive suggestions and appreciate the acknowledgement of thoroughness and novelty of this work. We hope the following clarifications can address the questions the reviewer raised.

The submission used the wrong template

Thank you for pointing at that, we will adjust the template accordingly.

The authors should explain the contextualization step in a bit more detail early on in the paper.

Thank you for your suggestion. We will update the "Edges between different fewshot examples" paragraph on page 4 to include the definition of a contextualization subcircuit earlier.

For reference, we include the draft of the change below:

Contextualization Subcircuit: Edges between different few-shot examples. We next investigated which edges are needed beyond the Parallel subcircuit, aiming to find a small circuit recovering ≥ 90% of the accuracy of the full model. We call all the edges in this circuit except those in the Parallel subcircuit a Contextualization subcircuit. Note that this subcircuit can include slightly different edges for different tasks.

Fig. 28 is helpful in understanding the patching methodology for circuit identification, so a small version could be included in the main text.

We will do our best to do this, the page limit permitting. We will extend section 3.1 with more details about the methodology we used, so that readers can refer to that early in the paper.

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.

The question of what task properties determine which circuit best explains the model's behavior is very interesting, and we currently do not have a rigorous answer for that. In preliminary experiments we found our contextualization circuits explain most of the model's behavior on tasks where the vanilla model achieves >90% accuracy. On the other hand, on tasks with lower accuracy (<75%) the circuits generally do not generalize well. We think the exact answer to this question is an exciting direction for future work.

We thank the reviewer for all the questions and suggestions and hope that adding additional clarifications make our findings more informative.

We are glad that the reviewer found our study to address an exciting problem, and to be detailed and contextualized with respect to related work.

The reviewer notes as their major concern questions about readability. We are glad to address this as follows.

(i) continual references to (very recent) prior works

We acknowledge that the paper has many references to this recent prior work. We did this in order to contextualize our work, and make the contributions transparent. We also are encouraged by the fact that the reviewer positively noted the contextualization of our work with respect to the literature.

That said, in order to improve the flow for the reader, we will largely move comparison to Cho et al and Kharlapenko et al to a dedicated section in the Discussion.

(ii) continual reference to figures in the appendix

We acknowledge that the paper has a good number of references to Appendix figures. We respectfully point out that the paper includes a large number of results, all of them quantitatively backed up with intervention results. We link our claims with tables and figures to show they are well-supported, but understanding the linked Appendix tables and figures is not necessary for following the main text. We note that it is infeasible to report the tables and figures backing up all of these in the main paper.

We will add explicit notes, such as "(Appendix, Figure 19)" instead of "(Figure 19)". This will improve the flow for the reader, while maintaining easy access to the relevant figures and tables. We believe that this effectively addresses the reviewer's concern.

unable to understand several parts of the authors claims and their descriptions of the methods.

To improve clarity in our description of the method, we will add a paragraph in section 3.1 with information about patching setup adopted from prior work and definitions. We will add a detailed description of our setup, and the way it is different from standard methodologies in Appendix. We provide the draft of the paragraph we will add to the main paper below:

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, a model is viewed as a computation graph with activations in different layers as nodes, and computations between them as edges. Then, it is possible to ablate some of the edges in the graph, forcibly replacing the computation along a specific edge with the one computed on the counterfactual input. Counterfactual inputs are designed to erase relevant information from the prompt while leaving other information intact for even better localization of model behavior. For example, when substituting "Berlin" with counterfactual "Paris", we isolate computations specific to the city's identity, not those activated by its category (city) or token-type (word). If ablating a set of edges does not lead to a 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 the circuits - subgraphs of the full model's computation graph that can explain a major part of model's performance on 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.

conclusions are limited to one model

We would like to point out that we do have results on different models within the Gemma family. As described in the paper, our motivation for the Gemma model family is that it performs relatively better on relevant ICL tasks than other families even at the 2B scale. Thus, including other model families would necessitate studying them on different tasks, which makes the findings across models less comparable. That is why we decided to stick to one model family and leave others for future work. We also note that this approach is common in circuit discovery papers (Lindsey et al. [1], Wang et al. [2]), as it allows controlled in-depth study with substantial insights.

[1] Lindsey, et al., "On the Biology of a Large Language Model", Transformer Circuits, 2025. (https://transformer-circuits.pub/2025/attribution-graphs/biology.html) [2] Wang, Kevin Ro, et al. "Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 Small." The Eleventh International Conference on Learning Representations. (https://arxiv.org/abs/2211.00593)

Overall, we appreciate the reviewer's feedback on writing, and noting two specific directions. As the reviewer described readability as their major concern, we believe that the changes described above effectively address this concern, and hope that the reviewer reconsiders their score. If there is any further feedback on readability, we would be happy to address that, too.