Which Attention Heads Matter for In-Context Learning?

Kayo Yin, Jacob Steinhardt

TL;DR: We compare the role of induction heads and function vector (FV) heads in in-context learning (ICL), we find that FV heads drive most of few-shot ICL which challenges the popular belief that induction implements ICL.

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 and compare induction heads and FV heads in 12 language models. Through detailed ablations, we find that few-shot ICL is driven primarily by FV 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 facilitates learning the more complex FV mechanism for ICL.

Lay Summary:

Large language models are surprisingly good at picking up new tasks just by seeing a few examples in their input. This skill is called in-context learning. But what allows these models to do this so well?

We looked inside these models to find out which internal components are most important for this learning ability. We focused on two types of “attention heads”—tiny components that help the model decide what information in the input to focus on. One type, called "induction heads", helps the model find and copy patterns from the examples it’s shown. The other type, called "function vector (FV) heads", seems to help the model understand the overall task it’s being asked to do, given examples.

By turning off each type of attention head in several models and seeing how that affects their performance, we discovered that FV heads are especially important for in-context learning, especially in larger models. We also found that many FV heads actually start out as induction heads when the model is first training, which suggests that the simpler induction heads may be a stepping stone for models to acquire FV heads.

By figuring out the mechanisms in AI models that are responsible for their learning abilities, we can design smarter, safer, and more understandable AI systems in the future.

Link To Code: github.com/kayoyin/icl-heads

Primary Area: Social Aspects->Accountability, Transparency, and Interpretability

Keywords: interpretability, in-context learning, large language models, mechanistic interpretability, induction heads

Submission Number: 13318