Which Attention Heads Matter for In-Context Learning?

Thank you for this thorough and constructive review! Your detailed feedback will definitely help improve our paper.

• We agree that we need to better contextualize our findings about ICL performance. In Appendix A.8, we have added the ablations broken down by task for each model, as well as a random baseline (red horizontal line) where we randomly sampled outputs from the training split to make predictions on the eval split. Most models, with the exception of Pythia 70M and 160M, outperform the random baseline. However, we also do note that by design, the random baseline is close to 0 for open-ended tasks instead of a finite set of labels. We have also computed majority class baselines (which by design, is similar to the random baseline we used) and zero-shot baselines using the same models we evaluate, the majority class baselines have similar performance to random baselines and zero-shot baselines are weaker. Since most of our models do show robust ICL abilities according to this evaluation, we believe it does not change the findings of the paper. Would you prefer we omit the smallest models with weak ICL abilities from our analysis, or retain the full spectrum while clarifying that the weakest models may not be able to do all the ICL tasks we evaluate on?

• Thank you for raising the important point about adaptive computation! Taking this into account, the most direct approach would be to measure the extent of the hydra effect in our ablations. However, this is an open research problem since we ablate specific heads that score highly on induction or FV scores from different layers, whereas McGrath et al.'s method applies to ablating all heads in a single layer. Another approach would be to apply logit lens to measure the direct contributions of different head types to the output, but this will not account for second order effects either. We agree that ablation on its own may not sufficient to show something is not important, although if the hydra effect is present in both ablations of induction and FV heads to the same magnitude (not sure if this is the case), our findings from this comparison would remain valid. We will present our findings more carefully taking this into account, add a discussion about adaptive computation, and provide the approaches mentioned above as suggestions for future research.

• We acknowledge Olsson et al.'s earlier conjecture that induction heads might be performing a more abstract version of the pattern matching-copying mechanism for ICL! However, we'd like to clarify that our C1 posits something different: FV heads are a distinct mechanism, not a generalization, of induction, but the existence of induction heads helps models learn the FV mechanism. This conjecture is based on results where some FV heads have high induction scores during training but get a drop in induction score while they increase in FV scores, while if FV is a generalization of induction, FV heads would maintain a high induction score.

• Thank you for your point on potentially misleading claims! We'd like to maintain clarity in our claims, and we have updated "ICL" to "few-shot ICL" where that is the case. Please let us know if any misleading claims remain! Thank you for your minor note as well, we will add error bars where needed.

Q1: We believe that the existence of FV heads that are also induction heads helps strengthen our claims and explain why previous literature attributed ICL to IH. Because of these heads, our work and prior works found that ablating IH (including FV heads that are also IH) leads to drops in ICL accuracy. However, once we only ablate non-FV IH, ICL accuracy is not affected as much. These heads also provide additional observations for C1, where FV heads "evolve" from IH, and the FV-IH heads are heads that did not fully lose induction mechanisms during training.

Q2: Individual heads are identifiable by how saturated their curve is in both top and bottom plots, so the FV head spiking the most in bottom right is the most saturated (it has the highest final FV score) but this head has a weak spike in induction score (spikes at around 0.8 score at 2^9). If this is difficult to parse, we can add markers on the plots to differentiate each individual heads across plots.

We apologize for our delayed response, and thank you again for helping us improve the clarity and rigor of our work. Please let us know if you have any remaining questions!

Thank you for raising this thoughtful point about novelty! We respectfully disagree, as our work provides several novel contributions beyond just comparison. Our joint analysis reveals previously unknown interactions between FV and induction heads, providing a more complete picture of ICL mechanisms. We also perform several novel analyses for the study of ICL mechanisms (such as performing ablations of one head type while excluding another, tracking the evolution of FV and induction heads over training in both their induction and FV scores). These findings help resolve apparent misunderstandings in prior work about induction heads' role in ICL. By bridging these previously separate lines of research, we advance our theoretical understanding of ICL and encourage future interpretability research to further study FV heads when studying ICL. We appreciate the opportunity to clarify these aspects of our work and thank you for helping us improve its presentation!

Thank you as well for noting the impact on future work! Our findings have direct implications for future research: in mechanistic understanding (clarifying how FV heads emerge during training, identifying which heads drive few-shot vs. token-loss ICL performance), and in recommended practices for interpretability research more generally (shift focus from induction heads to FV heads when studying ICL, question the universality hypothesis of interpretability, establish that different definitions of ICL in the literature measure very different things). These contributions provide concrete directions for improving model understanding. We'll clarify these research implications in the revision!

Thank you for your thorough and enthusiastic feedback!

Studying the effect of removing induction heads during training is a great suggestion! Although we do not have enough resources to re-train Pythia models with ablated heads, we agree that this experiment could help confirm C1, and we have added this suggestion for future work in our revision.

To address your questions:

1. This is an interesting question! We find it very curious that although Llama 2 7B has low induction and FV scores relative to other models (Figure 12), in our ablation experiments, ablating FV heads still lead to strong drops in ICL accuracy. One possible explanation is that these FV heads are polysemantic, so they exhibit weaker FV functionalities but implement a third functionality from induction/FV that are also important to ICL performance. Pythia 6.9B and Llama 2 7B are comparable in size, number of layers, number of attention heads, and size of attention heads, so this rules out these factors as the properties responsible for this finding. We believe the main difference might be the type of attention head used (Flash Attention in Pythia vs. Grouped Query Attention in Llama 2), further experiments comparing these two attention head types would help confirm this difference!

2. As for the effect of FV heads and induction heads with scale, we hypothesize that because FV heads are more complex (appear in deeper layers and later during training), FV heads that effectively perform ICL only emerge with scale, whereas induction heads are simpler to learn but perform ICL less accurately. Therefore, in smaller models, the contributions of the less accurate induction heads and the under-developed FV heads are comparable, whereas in larger models, the well-trained FV heads contribute more to ICL than the simpler induction heads.

We apologize for the delayed response, and thank you once again for your feedback!

Thank you for this detailed and constructive review! You raise several important points that will help strengthen our paper. To start, since you cited the layer location and overlap of heads as our key findings, we would like to clarify that our key findings are those illustrated in Figure 1: (1) ablating FV heads affects ICL score more than ablating induction heads, and (2) many FV heads evolve from induction heads during training. We present results on the layer location and head overlap primarily as a precursor to these findings, to first establish that FV and induction heads are indeed two distinct types of heads and therefore should be analyzed separately. We will clarify this in our paper!

W1: Thank you for raising this important point! We have added t-stat and p-values for the difference in head layer locations for each model in Figure 19 (Appendix A.9). As you have pointed out, the mean layer is often the same in many of our models, so we will clarify that this result does not hold in all models, but when the mean layer does differ, FV head mean layer is always deeper than IH mean layer. We will also emphasize the next result on head overlap more to establish the distinctness of the two types of heads.

W2: You raised that neighboring/other attention heads may perform similar functionalities to FV and induction heads which we might miss if we take exact overlap. Our FV score and induction score metrics precisely measure how much heads perform these functionalities, which is why we use these metrics to construct the set of heads to compare (e.g. these metrics would exclude heads in the same layer as FV heads but that does not perform FV functionalities). Also, please let us know if we misunderstood your point, since you cited McGrath et al. [1], however we understand this work as when we ablate all heads from a layer, a subsequent layer performs the functionalities of the ablated layer. Our section on head overlap does not perform any ablations yet so there would be no self-repair effects.

W3: Thank you for bringing the two related works to our attention, we have added them to our discussion! We'd also like to clarify that in our section on reconciling divergent findings, it stems from either because small models do rely on induction heads, or because earlier studies (on small or large models) do not account for the correlation of FV and induction heads. We believe the two works you have provided fall into the second category, where studying only induction head ablations can give an incomplete account of the ICL mechanism landscape. Our work reproduces these two works' findings as well, and we'd also like to clarify that we did not intend to claim to be the first to study ICL heads at scale. Please let us know if you found any of our claims misleading!

Q1: We calculated the induction and FV scores of attention heads from each checkpoint separately to compute the evolution of scores for specific heads.

Q2: C2 states that the FV heads used for ICL are polysemantic heads that implement a combination of induction and a different mechanism. Therefore, monosemantic FV heads that do not perform induction mechanisms are excluded from this category of ICL heads, which is why C2 would predict that ablating these heads would not hurt performance. We agree that since we do observe ablating non-induction FV heads drops ICL score, these monosemantic may be performing important operations beyond induction, hence why we state that this observation contradicts a prediction from C2.

We apologize for our delayed response, and thank you again for helping us improve the clarity and rigor of our work. Please let us know if you have any remaining questions!

Thank you for this detailed and thoughtful review. You raise several important points that will help improve our paper!

W1: Great question regarding our choice of ablation method! We chose mean ablation following recommendations from Hase et al., 2021 [1], Zhang and Nanda 2024 [2], Wang et al., 2022 [3] where the suggest using mean ablation over zero ablation to avoid the out-of-distribution problem with zero ablations. Since the Olsson et al., 2022 paper on induction heads use zero ablation in their ablation studies, we have run additional experiments using zero ablations (replacing the ablated head's output by a zero vector), as well as random ablations (replacing the ablated head's output by a random vector). We added these results in Appendix A.3 (Figure 9). Overall, these ablations give similar results to mean ablations and our findings are robust to the choice of ablation method.

W2: Great point, thank you! We performed ablations of random heads in the different layers that correspond to where induction and FV heads appear in Pythia 6.9B, and present results in Appendix A.4 (Figure 10) of the revision. While ablating random heads in certain layers, namely the later layers, gives more drop to ICL score than other (earlier) layers, overall, the effect of ablating random heads is much weaker than ablating induction / FV heads. Therefore, we are no longer concerned about the confounding factors of head location in this comparison.

W3: We provide a list of ICL tasks we use in Appendix A.7 (Table 3), and data for these tasks will also be provided with the code release. While some tasks would not reflect real-world use cases (e.g. "capitalize second letter", "next-item"), most of these tasks would represent real-world use cases (e.g. English-French lexical translation, sentiment analysis, currency of countries).

W4: Thank you for your feedback on our clarity! We have revised Table 1 to better distinguish findings backed by evidence in our paper and conjectures with no sufficient evidence.

W5: Thank you! We have added additional intuitive explanations of FV heads, please let us know if the background in Section 2 remains unclear.

Q1: We chose the top 2% (98th percentile) based on induction/FV scores to define our sets of induction and FV heads in Section 3.2, and calculated the overlap between these sets. While most induction heads do not lie in the 98th percentile of FV heads and vice versa, the middle and right plots in Fig 3 show that induction heads still have relatively high FV scores (90-95th percentile) and vice versa. This result is especially important for ablation experiments later in the paper, where naively ablating induction heads would also ablate FV heads with high induction scores and vice versa, hence why we perform ablations with exclusion to account for the overlap during ablation. We hope this clarifies!

Q2: For actionable implications, we believe that future work exploring how to inject or encourage smaller models to develop FV heads may be promising to improve ICL performance in small models. We would also be excited to see more work in the field of interpretability studying these FV heads, for example what exactly these function vectors are encoding or what mechanism FV heads are more concretely implementing, since a lot of works on ICL interpretability so far have focused on induction heads.

We apologize for the delayed response, and we appreciate your feedback helping strengthen the experiments and presentation of our work!

[1] Peter Hase, Harry Xie, and Mohit Bansal. The out-of-distribution problem in explainability and search methods for feature imporrtance explanations. Advances in neural information processing systems, 34:3650–3666, 2021. [2] Fred Zhang and Neel Nanda. Towards best practices of activation patching in language models: Metrics and methods. In The Twelfth International Conference on Learning Representations, 2024. [3] Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023.

To our knowledge, we have also responded to all reviewer questions, clarified confusions, and addressed all major concerns reviewers brought up:

1. Questions about robustness of results

• "I am not very convinced by the ablation method used in section 4.1, i.e., by replacing output vector by mean values. [...] How robust are the results, or is it specific only to the ablation method used here?" (1PLt)

We provided references explaining the advantage of mean ablation over other ablation methods (zero/random). We also reproduced experiments using zero ablation and random ablation instead, and show that our results are robust against ablation method.

• "Head "location" can be one confounding factor that contributes to the difference in ICL performance [...]. There should perhaps be a controlled baseline that ablates heads at different locations in the model." (1PLt)

We performed ablations of random heads at different layers as additional baselines. The effect of ablating random heads is minimal relative to ablating induction / FV heads, regardless of the layer chosen. We verified that head location is not a confounding factor.

• "Some forms of statistical tests might need to be done here to strengthen the argument [that FV heads appear in deeper layers than induction heads do]." (wGjS)

We added t-stat and p-values for the difference in head layer locations for each model in Figure 19 (Appendix A.9).

• "It is possible that neighboring attention heads could be performing similar functionalities [as FV and induction heads] in LLMs." (wGjS)

Our FV score and induction score metrics precisely measure whether each head performs these functionalities, so this is a non-issue.

• "Can the authors clarify if the models they're testing are consistently better than baselines for the tasks they use?" (Ft6Y)

We added a random baseline for each task we used. Most models, with the exception of Pythia 70M and 160M, outperform the random baseline, and clarified in our paper that these two small models can't do all the ICL tasks used.

2. Presentation

• "It is not clear how many tasks are evaluated, and which ICL tasks are used exactly? How well do the tasks represent real-world ICL/few-shot use cases?" (1PLt)

We provide a list of ICL tasks we use in Appendix A.7 (Table 3). Most tasks represent real-world use cases.

• "Paper can be made more sound to clarify conjectures from conclusions with substantiated results" (1PLt)

We revised Table 1 to better distinguish findings backed by evidence in our paper and conjectures with no sufficient evidence.

• "Works such as [2][3] have already extended induction heads to large models up to 20B/66B" (wGjS)

We added a discussion of [2][3] in our related work. These two works provide further examples of one of our paper's claims: studying only induction head ablations gives an incomplete account of the ICL mechanism landscape.

• "In some cases, ablation can cause adaptive computation [1]" (Ft6Y)

We added a discussion about adaptive computation and presented our conclusions more carefully taking this into account. We also suggested future research directions to take into account adaptive computation.

• "Some of the claims in the paper are a bit misleading [...] A way to make these kind of statements less misleading would be to qualify that this means "few-shot ICL" in the same sentence" (Ft6Y)

We updated "ICL" to "few-shot ICL" where that is the case to improve clarity.

We thank all reviewers for their helpful feedback! During the discussion phase, we significantly improved the robustness of our results with additional experiments, and further clarified our paper's claims.

We are excited by this work since it challenges the community's existing beliefs about the role of induction heads in ICL, and provides important lessons/directions for future interpretability research (disentangle "few-shot" vs. "token-loss difference" ICL, universality hypothesis in interpretability, conjectures on the evolution of FV heads from induction heads). We hope our work will also inspire more research studying FV heads in ICL!

---

[1] Thomas McGrath, Matthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg. The Hydra Effect: Emergent Self-repair in Language Model Computations. 2023 [2] Joy Crosbie and Ekaterina Shutova. Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning. 2024 [3] Hritik Bansal, Karthik Gopalakrishnan, Saket Dingliwal, Sravan Bodapati, Katrin Kirchhoff, and Dan Roth. Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale. 2023