Source-linked AI summary
Which Attention Heads Matter for In-Context Learning?
Kayo Yin, Jacob Steinhardt
TL;DR
The paper asks which of two proposed mechanisms primarily drives few-shot ICL and compares induction and FV heads across 12 language models using ablations and training-dynamics analysis. It finds that FV heads are more causally important, especially in larger models, while many FV heads develop from earlier induction-head behavior.
Problem
The paper addresses uncertainty over whether induction heads or FV heads primarily drive few-shot in-context learning.
Method
The authors compare induction and FV heads across 12 decoder-only models and 45 natural-language ICL tasks using ablations and training-dynamics analysis.
Results
FV heads are the primary drivers of few-shot ICL performance, with their importance relative to induction heads increasing as model scale increases.
Takeaways & Limitations
Many FV heads begin with strong induction behavior before transitioning to FV functionality, supporting induction as a possible precursor to the more effective FV mechanism.
Takeaways & Limitations
The interpretation remains partly conjectural because the study observes transitions during training but does not establish that removing induction heads during training changes FV-head development.
Abstract
from arXiv · showhide
Large language models (LLMs) exhibit impressive in-context learning (ICL) capability, enabling them to perform new tasks using only a few demonstrations in the prompt. Two different mechanisms have been proposed to explain ICL: induction heads that find and copy relevant tokens, and function vector (FV) heads whose activations compute a latent encoding of the ICL task. To better understand which of the two distinct mechanisms drives ICL, we study and compare induction heads and FV heads in 12 language models. Through detailed ablations, we discover that few-shot ICL performance depends primarily on FV heads, especially in larger models. In addition, we uncover that FV and induction heads are connected: many FV heads start as induction heads during training before transitioning to the FV mechanism. This leads us to speculate that induction facilitates learning the more complex FV mechanism that ultimately drives ICL.
1. Introduction
The paper compares induction and function vector (FV) heads as proposed mechanisms for in-context learning (ICL), finding that FV heads primarily drive few-shot ICL while induction heads may precede them during training. It also explains prior disagreements through metric choice, head correlations, and model scale.
- Motivation and approach: The study compares induction circuits, which identify and copy tokens, with FV heads, which encode task information for ICL.The analysis covers 12 decoder-only models ranging from 70M to 7B parameters and 45 natural-language ICL tasks.
- Distinct mechanisms: Induction and FV heads are distinct: they have low or zero overlap, with induction heads generally appearing earlier in layers and training.Despite this separation, each type behaves more similarly to the other than to random heads in the same network.
- Ablation findings: FV heads are the primary drivers of few-shot ICL: ablating them substantially reduces task accuracy, whereas ablating induction heads has limited effect across all 12 models.The difference becomes more pronounced as model size increases, challenging the view that induction heads are the key ICL mechanism.
- Reconciling prior work: Earlier conclusions differed partly because prior studies used an ICL metric that does not strongly track few-shot performance, overlooked FV–induction correlations, or focused on small models.The paper identifies metric choice as the most significant methodological difference.
- Training dynamics: Many induction heads evolve into FV heads during training, while the reverse transition is never observed.This developmental pattern motivates the speculation that induction may facilitate learning the more complex FV mechanism.
- Broader implications: The findings caution that correlated mechanisms can produce illusory explanations and that the relative importance of FV and induction heads changes substantially with model scale.These results also challenge strong versions of universality in model interpretability.
2. Background & related work
Research on ICL has focused on induction circuits and function vectors as distinct explanatory mechanisms, while differing metrics and experimental controls can produce divergent conclusions about their roles.
- Induction heads identify repeated tokens and attend to the token following an earlier occurrence, thereby predicting the repeated continuation.
- Earlier induction-head studies measured ICL using token-loss difference and found that induction-head emergence and ablation affected that metric.
- Function vectors are compact task representations extracted from specific attention heads and added to model computations to reproduce ICL behavior without demonstrations.
- This paper defines ICL as few-shot learning from prompt demonstrations without parameter updates, distinguishing it from token-loss difference.
- FV heads primarily contribute to few-shot ICL accuracy, whereas induction heads show the opposite pattern on token-loss difference.
- Controlling for correlations between head types and accounting for model scale reconcile the paper’s findings with earlier induction-focused results.
3. Induction heads and function vector heads are distinct but correlated
Induction and FV heads are mostly distinct by top-ranked membership and layer location, yet their scores are strongly correlated across models.
- The analysis compares induction and FV heads before measuring their relative contributions to ICL performance.
- The study reports model parameters, layers, attention heads, and per-head dimensions for the models analyzed.
- Induction heads generally occur in early-middle layers, while FV heads appear in slightly deeper layers.The deeper FV locations may indicate more abstract computations, but the paper labels that interpretation speculative.
- Seven of twelve models show zero overlap between the top 2% induction and FV heads, while the remaining models show only 5-15% overlap.
- Despite minimal set overlap, top FV heads usually fall near the 90-95th percentile of induction scores, and vice versa.
4. Function vector heads drive in-context learning
Systematic ablations show that FV heads are the primary drivers of few-shot ICL accuracy, while induction-head effects largely reflect overlap with FV heads. This conclusion is strongest in larger models and depends on measuring few-shot accuracy rather than token-loss difference.
- Ablation results: Exclusion ablations preserve the top 2% FV heads while removing induction heads, producing minimal ICL impact in models exceeding 1B parameters.The effect is comparable to random ablations in those models.
- Ablation results: Exclusion ablations that remove FV heads while preserving induction heads continue to significantly impair few-shot ICL performance.The FV–induction ablation gap widens with model scale, suggesting that earlier induction effects were driven by heads with both properties.
- Metric comparison: Few-shot ICL accuracy and token-loss difference capture different phenomena: FV heads strongly affect the former, whereas induction heads affect the latter more strongly in larger models.This divergence helps account for contrasting conclusions about which heads matter for ICL.
5. FV heads evolve from induction heads
Training analyses across seven Pythia models show that induction heads emerge early, while FV heads arise later and often develop from initially induction-like heads. During training, these heads lose induction strength as their FV strength increases.
- Analysis setup: The developmental analysis uses eight intermediate checkpoints from seven Pythia models and tracks the mean scores of the top 2% induction and FV heads.Few-shot ICL accuracy is plotted alongside these mechanism scores.
- Training dynamics: Across all Pythia models, induction heads emerge around step 1,000 of 143,000, while FV heads appear around step 16,000.Induction scores rise sharply before plateauing or slightly declining, whereas FV scores increase gradually through training.
- Head transitions: Many heads that ultimately become strong FV heads initially show high induction scores comparable to specialized induction heads.These proto-FV heads emerge around the same time as dedicated induction heads.
- Head transitions: As training progresses, proto-FV heads’ induction scores decline while their FV scores increase, with no observed reverse transition from induction heads to FV heads.The analysis tracks individual heads across training checkpoints using induction and FV scores.
6. Interpretation and discussion
The evidence distinguishes induction and FV mechanisms while showing a developmental connection: many FV heads emerge from earlier induction behavior, and FV heads have the larger effect on few-shot ICL.
- Ablation evidence: FV heads are more important for few-shot ICL than induction heads, especially as model size increases.Ablating FV heads substantially harms ICL, whereas ablating induction heads has comparatively limited effects.
- Interpretation: The induction-to-FV conjecture treats induction heads as an early stepping stone toward the more sophisticated FV mechanism.The conjecture remains incomplete because some FV heads maintain low induction scores throughout training.
- Training dynamics: Many FV heads initially show strong induction behavior before transitioning to FV functionality, while the reverse transition is not observed.The decline in induction scores after FV functionality develops supports a unidirectional developmental relationship.
- Training dynamics: Induction scores rise earlier and more sharply, whereas FV scores emerge later and increase gradually alongside continuing ICL improvement.This timing is consistent with induction contributing to an initial ICL rise and FV mechanisms supporting later improvement.
- Alternative interpretation: An alternative polysemantic-head account is challenged because ablating monosemantic FV heads without significant induction scores still substantially harms ICL.This result is difficult to reconcile with the prediction that combined induction-FV functionality is essential.
7. Conclusion
The paper argues that FV heads, rather than induction heads, play the more crucial causal role in few-shot ICL, while also identifying methodological and scale-dependent qualifications. It further reports that many FV heads develop from earlier induction behavior and leaves several questions open for future work.
- Conclusion: FV heads play a more crucial causal role in few-shot ICL than induction heads.The conclusion attributes earlier disagreement partly to conflating few-shot accuracy with token-loss metrics and overlooking head overlap.
- Conclusion: Many FV heads are former induction heads whose earlier induction behavior may provide an initial foundation for ICL.The paper presents this as support for a conjecture about induction heads preceding FV heads.
- Conclusion: The findings caution that capability metrics and overlapping mechanisms can produce misleading interpretability conclusions when studied in isolation.Few-shot accuracy and token-loss difference can yield different conclusions, and shared head behavior can confound ablations.
- Conclusion: Mechanism importance diverges with scale: FV heads become increasingly crucial, while induction-head impact approaches random-ablation levels.The authors note that both mechanisms contribute meaningfully in smaller models.
- Open questions: Open questions concern why induction may precede FV heads, the role of remaining induction heads, and whether additional mechanisms explain ICL.These questions define boundaries for the paper’s current account.
A.1. Induction scores vs. FV scores
Figure 7 plots the induction score and FV score for each attention head.
- A.1. Induction scores vs. FV scores: The figure compares induction and FV scores across individual attention heads.Use it to inspect whether heads with high scores for one mechanism also score highly for the other.
A.2. Ablations
Figures 8 and 9 track model behavior as increasing numbers of heads from each head type are ablated.
- A.2. Ablations: Figure 8 plots ICL-task accuracy across different quantities of ablated heads, while Figure 9 plots token-loss difference under the same type of comparison.Compare how the two metrics change as ablation quantity increases.
A.3. Random and zero ablations
Ablation and scaling analyses compare induction and FV heads as contributors to ICL, while additional experiments test FV-based task execution and control for head overlap. Across models, FV effects are generally stronger for ICL, though score trends with model scale are not consistent.
- Random and zero ablations measure ICL accuracy after replacing selected head outputs with outputs from other heads or zero vectors.Random ablations also target specific layers under a defined sampling procedure.
- FV heads are increasingly important relative to induction heads for ICL performance as model scale increases.
- Induction scores remain relatively similar across model sizes, while FV scores show no clear scale trend and are especially high in Pythia 1B and 1.4B.Pythia 70M and Llama 2 have relatively low FV scores, and induction-score changes with scale are small or mixed.
- Adding task-specific FV representations to shuffled prompts recovers model performance on uncorrupted prompts in most models, except Pythia 2.8B.The evaluation averages accuracy over 40 ICL tasks using previously unseen examples.
- Ablation comparisons require exclusion controls because the overlap between ablated induction and FV heads increases as more heads are removed.Task-level ablations compare accuracy against task-specific random baselines.
- Induction and FV scores are tracked across training alongside few-shot accuracy to examine how the two mechanisms develop over time.