Which Attention Heads Matter for In-Context Learning?

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Yin, Kayo, Steinhardt, Jacob
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866918006650568704
author Yin, Kayo
Steinhardt, Jacob
author_facet Yin, Kayo
Steinhardt, Jacob
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_14010
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Which Attention Heads Matter for In-Context Learning?
Yin, Kayo
Steinhardt, Jacob
Machine Learning
Artificial Intelligence
Computation and Language
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.
title Which Attention Heads Matter for In-Context Learning?
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2502.14010