In-context Learning and Induction Heads

TL;DR

This paper hypothesizes that induction heads are the primary mechanism behind in-context learning in large transformer models, supported by evidence from models of varying sizes and development stages.

Contribution

It provides the first comprehensive analysis linking induction heads to in-context learning, with causal evidence for small models and correlational evidence for larger models.

Findings

Induction heads develop concurrently with in-context learning ability.

Induction heads are likely the mechanistic source of in-context learning.

Evidence spans small to large transformer models.

Abstract

"Induction heads" are attention heads that implement a simple algorithm to complete token sequences like [A][B] ... [A] -> [B]. In this work, we present preliminary and indirect evidence for a hypothesis that induction heads might constitute the mechanism for the majority of all "in-context learning" in large transformer models (i.e. decreasing loss at increasing token indices). We find that induction heads develop at precisely the same point as a sudden sharp increase in in-context learning ability, visible as a bump in the training loss. We present six complementary lines of evidence, arguing that induction heads may be the mechanistic source of general in-context learning in transformer models of any size. For small attention-only models, we present strong, causal evidence; for larger models with MLPs, we present correlational evidence.