Which Attention Heads Matter for In-Context Learning?
Keywords: interpretability, in-context learning, large language models, mechanistic interpretability, induction heads
TL;DR: We compare the role of induction heads and function vector (FV) heads, we find that FV heads drive most of in-context learning (ICL), challenging the popular belief that induction implements ICL, FV heads might evolve from induction heads.
Abstract: Large language models (LLMs) exhibit impressive in-context learning (ICL) capability, enabling them to generate relevant responses from a handful of task demonstrations in the prompt.
Prior studies have suggested two different explanations for the mechanisms behind 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 induction heads and FV heads in 12 language models.
Our study reveals that in all 12 models, few-shot ICL is driven primarily by FV heads: ablating FV heads decreases few-shot ICL accuracy significantly more than ablating induction heads, especially in larger models. We also find 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 heads facilitate the learning of the more complex FV mechanism for ICL.
Finally, the prevalence of FV and induction heads varies with architecture, which questions strong versions of the
"universality" hypothesis: findings from interpretability research are not always generalizable across models.
Primary Area: interpretability and explainable AI
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 12492
Loading