Which Attention Heads Matter for In-Context Learning?

Thank you for the detailed feedback, and we especially appreciate the clear questions to resolve reviewer concerns!

We first address reviewer concerns on the claims:

1. We used Pythia models mostly to facilitate the analysis over training dynamics, which other open-source models do not enable since they do not release intermediate checkpoints. We will add Llama 3 8B and Qwen 2.5 7B to our paper - as these models were not yet released during the development of this paper.
2. In many cases, we had to decide on a threshold to differentiate meaningful FV/induction heads from the long tail of other heads. The 2% threshold was carefully chosen following previous work by Todd et al. (2024), and after we verified this threshold does seem to meaningfully separate important heads (Figure 7 in appendix shows the 2% threshold does select heads outside the cluster of other low-scoring heads - taking a higher percentage would select heads with FV scores close to 0). If we take the top 5% of heads, there would be between a 20-40% overlap (except for our smallest model with a 67% overlap) - which is still a small enough overlap to motivate our subsequent experiments in comparing the two sets of heads. The main goal of the claim on overlap is to establish that the two sets of heads we're studying are not identical, and furthermore, our ablation and training experiments later that shows different trends for FV vs. induction heads reinforces that the two are distinct types. Beyond a 5% threshold, we would be selecting too many heads with very low FV and induction scores that it would not be meaningful to consider them FV/induction heads.
3. We agree that the observation on layer depths is not statistically backed up - we will clarify that this observation is speculative.
4. Great suggestion! We have computed and added the evolution of ICL accuracy to Figure 5 of section 5: we observe that in all models, few-shot ICL accuracy begins to improve around the same time as when induction heads appear, and continues to gradually increase throughout training until the end. Since ICL accuracy continues to improve even after the formation of induction heads, we speculate that this suggests the sharp emergence of induction heads contributes to an initial rise in ICL performance, but the emergence of the FV mechanisms contributes to further improvements in ICL (which reinforces our conjecture 1).

We also address the questions at the end:

1. Thank you very much for pointing this out and for the thorough, actionable feedback under "Other comments or suggestions"! We have incorporated your feedback to remove repetitions; reorganize the structure of related work / background / methods; added suggested references to related work; edited Table 2 in the main body to show the model, parameters and |L| only (while putting the rest in the Appendix); moved figures in page 7 to be closer to their text; added a limitation section (where we discuss limitations on the complexity of ICL tasks studied and the scale of the models in our analysis); added additional markers, grids, subcaptions to figures; removed unnecessary borders and empty images.
2. We are currently repeating our experiments on Llama 3 8B and Qwen 2.5 7B. We will also try to include models in the 12-13B parameter range that might fit in A100s. However, we cannot perform experiments on larger models since, to compute the FV score, we need to store several caches of all the attention head activations - we will suggest experiments on larger models for future work.
3. Our conclusion mainly serves to clarify misunderstandings in the current interpretability literature, that attributes few-shot ICL to induction heads, as well as providing general lessons for interpretability such as how interpretability conclusions change with model scale, and the definition of the metric used to measure ICL can also affect conclusions. Since few-shot ICL is still an active field of study in interpretability, we believe our conclusions are very important to share with the field, and we hope that our general lessons for interpretability will also apply to future interpretabiltiy research on other methods!

Thank you again for your feedback, please let us know if any questions or concerns remain unsolved!

Thank you for your feedback! We will first address the concerns around claims and evidence:

• On the layer distribution, we will clarify that the claim on the layer depth difference is more of a speculative observation than a claim sufficiently backed up statistically.
• On the selection of 2%: the 2% threshold was carefully chosen following previous work by Todd et al. (2024), and after we verified this threshold does seem to meaningfully separate important heads. Figure 7 in appendix shows the 2% threshold does select heads outside the cluster of other low-scoring heads - taking a higher percentage would select heads with FV and induction scores close to 0, which would prevent meaningful analysis. We also repeated ablation experiments taking 5% and 10% as thresholds, and find that our claims are robust to the thresholds.
• On head overlap: the main goal of the claim on overlap is to establish that the two sets of heads we're studying are not identical, to motivate subsequent experiments comparing the two sets. Our ablation and training experiments later that shows different trends for FV vs. induction heads further reinforces that the two are distinct types. If we take the top 5% or top 10% of heads, there would be between a 20-40% overlap (except for our smallest model with a 67% overlap for top 5% and 80% overlap for top 10%) - while taking 5% and 10% as the threshold does not make as much sense since it includes heads with FV and induction scores close to 0, it still shows that the overlap is low enough to justify treating the two sets of heads as distinct heads for study.
• On the rigor of conjectures: we would like to clarify that these are mere conjectures, inspired from observations in our study, but that remain to be rigorously proven and we are not positing them as empirically verified claims. We leave these conjectures as discussions and interpretations of our empirical findings, and to inspire future analyses.

We also address weaknesses raised by the reviewer:

• On ICL tasks: our study combines ICL tasks that are synthetic but require reasoning that is achievable by the models we study, with realistic tasks (such as translation, commonsense question answering, news topic classification, etc.) so we believe that the selection of ICL tasks is sufficient for our analysis. We do appreciate the suggestion for realistic tasks, and will consider adding further tasks that is feasible by our models (such as sentiment classification).
• On token loss: the difference between TL difference and few-shot ICL accuracy was first described in Olsson et al. (2022), but we will include a discussion in our revision as well. In short, TL difference approximates how much the model is using context by measuring how much the model's loss decreases on tokens further in context than earlier. When we observe high ICL accuracy with low TL difference (which is what we observe when we ablate random heads and induction heads for example), this means the model preserves its ICL abilities but loses other abilities associated with learning important signals from context. These other signals could capture non-ICL signals such as remembering entities, tracking style from context, etc. It would be great for future work to analyze what is being lost when we see these drops in TL difference!
• On newer models: we are currently reproducing all experiments except section 5 on training dynamics on Llama 3 8b and Qwen 2.5 7B. Since our conclusions have applied to three families (GPT-2, Pythia, Llama 2) that were released in different years (2019, 2022, 2023), we would expect our conclusions to generalize to the models released in 2024 as well. If surprisingly they don't, we would still like to share the interesting negative results!
• On suggestions for future development: while papers in the interpretability literature are often not expected to provide this (e.g. the Olsson et al. (2022) and Todd et al. (2024) papers that we compare the mechanisms for do not) as the interpretability contributions are standalone from model development, we will include a discussion on how evidence of the higher influence of FV heads may suggest that this is a more effective mechanism than induction heads. This suggests that ICL can be optimized in small models by using training methods that promote the formation of FV heads.

Thank you again for your feedback, please let us know if any questions or concerns remain unsolved!

Thank you for your feedback! We will first respond to the questions about claims and evidence:

• Choice of percentile: In many cases, we had to decide on a threshold to differentiate meaningful FV/induction heads from the long tail of other heads. The 2% threshold was carefully chosen following previous work by Todd et al. (2024), and after we verified this threshold does seem to meaningfully separate important heads (Figure 7 in appendix shows the 2% threshold does select heads outside the cluster of other low-scoring heads - taking a higher percentage would select heads with FV scores close to 0).
• Ablations with exclusion: we repeated our ablation with exclusion experiments on Pythia-6.9B taking 5% and 10% as the threshold. Overall, we observe the same trends as with 2%, where ablating induction heads does not hurt ICL much more than random, while ablating FV heads quickly significantly degrades ICL performance. Our ablations are robust to the threshold chosen for exclusion. We will repeat these experiments with other thresholds for our other models, and add them to the Appendix.

We will also address weaknesses:

• Newer models: we are currently reproducing all experiments except section 5 on training dynamics on Llama 3 8b and Qwen 2.5 7B. Since our conclusions have applied to three families (GPT-2, Pythia, Llama 2) that were released in different years (2019, 2022, 2023), we would expect our conclusions to generalize to the models released in 2024 as well. If surprisingly they don't, we would still like to share the interesting negative results!
• Threshold selection: we will elaborate in our revision the rationale behind choosing 2% - that in addition to following prior work, we also empirically verified in Figure 7 that selecting a higher threshold such as top 5% or 10% would end up selecting attention heads with FV and induction scores that are close to 0, which would prevent meaningful analysis. We have also previously computed the correlation between the two scores, but the correlations are not informative due to the long tail of heads that are neither high on FV or induction. As for the question about contradiction, we would like to clarify that we first established a lack of significant overlap to motivate our subsequent experiments that study FV and induction as two separate phenomena (if the two heads significantly overlap and correspond to the same set of heads, we no longer have a reason to compare these two heads). However, while the two sets of heads are mostly distinct, there are some heads that have both induction and FV properties which may have confounded earlier studies (e.g. heads in purple on Figure 7). Thank you for raising this confusion, we will clarify this in our revision!
• For our ablations, we have also computed and plotted the average scores and percentils of the heads preserved. We have added these plots to our appendix.

As for your last suggestion on separating the performance, we have plotted the ablations for each task in Figures 15-18 of the Appendix. Following your suggestion, we have added markers for the 8 new tasks to better differentiate them from the other 11 tasks.

Thank you again for your feedback, please let us know if any questions or concerns remain unsolved!

Thank you for the positive assessment! We also address minor concerns:

• Larger models: we will consider experiments on models in the 13B range for our camera-ready. However, computational constraints on our end prevent testing on larger models due to memory limits, since the FV score computation and ablation experiments require several caching of model activations. We would love to see future work extend our studies to bigger models! The models we study are also much bigger than earlier mechanistic analyses of ICL (e.g. Olsson et al. 2022 on toy models with 1-6 attention-only layers or MLPs) which shows promise of scaling our methods to bigger models given enough compute.
• Task difficulty: while small models limit complex reasoning tasks, our findings hold on tasks such as translation, binding, named entity recognition, commonsense QA... It would be a great idea to also include the more challenging tasks you suggest for future work studying bigger models that have a reasonable base accuracy for those tasks!
• Thank you for the suggestion! We have repeated the ablation experiments and training dynamics analysis using 5% as the threshold, and results confirm the robustness of the observed trends.
• Practical applications of findings: we will include a discussion on how evidence of the higher influence of FV heads may suggest that this is a more effective mechanism than induction heads. This suggests that ICL can be optimized in small models by using training methods that promote the formation of FV heads.

We thank the reviewer for their thoughtful feedback! We address concerns below:

1. Our task suite does include dictionary-style tasks as you mentioned, and we plotted the task-specific ablations in Figure 18 of the appendix - for the dictionary tasks (e.g. national_parks, park-country, person-occupation, person-instrument, english-german, french-english) we do still consistently find that FV heads contribute more than induction heads significantly.
2. In many cases, we had to decide on a threshold to differentiate meaningful FV/induction heads from the long tail of other heads. The 2% threshold was carefully chosen following previous work by Todd et al. (2024), and after we verified this threshold does seem to meaningfully separate important heads (Figure 7 in appendix shows the 2% threshold does select heads outside the cluster of other low-scoring heads - taking a higher percentage would select heads with FV scores close to 0). If we take the top 5% of heads, there would be between a 20-40% overlap (except for our smallest model with a 67% overlap).
3. While we used a held-out set of tasks to identify FV heads, the FV heads are task-agnostic (as shown in our ablations where they contribute to tasks that were not used to identify them) - so regardless of the choice of the ICL task, with enough tasks the same set of heads should be identified as FV heads. In addition, the tasks we used to compute FV heads are ones that were previously often associated with induction heads, but demonstrated in our paper as being less influenced by induction heads than FV heads. Therefore, the lack of overlap is a general phenomenon. Also, the main goal of the claim on overlap is to establish that the two sets of heads we're studying are not identical, to motivate subsequent experiments comparing the two sets. Our ablation and training experiments later that shows different trends for FV vs. induction heads further reinforces that the two are distinct types.
4. We agree token-loss difference trends are noisy and we will remove the claim about scale-dependence. The key finding we'd like to emphasize here is that token-loss difference and ICL accuracy measure distinct phenomena.
5. We agree that the observation on layer depths is not statistically backed up - we will clarify that this observation is speculative.
6. Great point! We used a different set of tasks to compute FV scores than the tasks used in our ablation experiments, to specifically ensure that our main claim is not tautological. What's interesting about FV heads is that their contribution to ICL seems to generalize across different ICL tasks.
7. Interesting question! The difference between the two metrics was first described in Olsson et al. (2022), but we will include a discussion in our revision as well. In the prompt you described, it would be a normalized version of few-shot ICL, but in practice, TL difference is not computed using such prompts (in previous work and ours, it is taken over natural sentences). And yes, high ICL accuracy with low TL difference (which is what we observe when we ablate random heads and induction heads for example) means the model preserves its ICL abilities but loses other abilities associated with learning important signals from context. These other signals could capture non-ICL signals such as remembering entities, tracking style from context, etc. (it would be great for future work to analyze what is being lost when we see these drops in TL difference!)

Thank you for providing the additional references, suggesting better clarity in 2.2, and pointing out the typos - we have incorporated these in our revision!

“However, C2 would predict that ablating monosemantic FV heads would not hurt ICL performance, […]” Why is that ?

C2 posits that FV heads are polysemantic heads that implement both induction and the FV mechanism. Then, if we ablate monosemantic FV heads but preserve the polysemantic FV heads, models should still be able to perform ICL - we will clarify this in our revision!

Thank you again for your feedback, please let us know if any questions or concerns remain unsolved!

We thank the reviewer for their feedback!

To address your weaknesses:

• On the shift from induction to FV: we will elaborate on why this shift might be necessary for larger models. We hypothesize that induction is a more lightweight mechanism (based on just attending to the subsequent token of a previous copy of the current token) that offers a crude way for models to retrieve information from context, however this method does not sufficiently solve more complex ICL problems. FV heads, on the other hand, offer a compact representation of the ICL task given the examples - the exact mechanism of FV heads is still unknown, and plausibly more complex, but allows larger models to achieve higher training accuracy. We also discussed the possibility of other mechanisms that may play a role in ICL, we leave this investigation for future work.
• On the conjecture between induction and FV heads relationship: we performed an additional experiment measuring the evolution of ICL accuracy during training. We observe that in all models, few-shot ICL accuracy begins to improve around the same time as when induction heads appear, and continues to gradually increase throughout training until the end. Since ICL accuracy continues to improve even after the formation of induction heads, we speculate that this suggests the sharp emergence of induction heads contributes to an initial rise in ICL performance, but the emergence of the FV mechanisms contributes to further improvements in ICL. This reinforces our conjecture that induction heads may serve as a stepping stone for FV, and we leave further exploration to rigorously verify this conjecture to future work.
• We will also add a further discussion on the application of our findings: our evidence of the higher influence of FV heads suggests that this is a more effective mechanism for ICL than induction heads. This suggests that ICL can be optimized in small models by using training methods that promote the formation of FV heads.

Thank you for suggesting the relevant work by Wu et al.! We also extensively discussed and compared with the Anthropic work suggested (which we cite as Olsson et al. 2022).

Thank you again for your feedback, please let us know if any questions or concerns remain unsolved!