What do induction heads do?

What is in-context learning -> the ability of transformers to learn from a few examples provided in the input context, without updating the model parameters. This phenomenon allows models to adapt to new tasks and data distributions on-the-fly, leveraging the information presented in the prompt. Give examples.

What is induction heads -> attention heads that are capable of recognizing and continuing patterns in sequences, particularly in the context of in-context learning. They can identify repeated structures or n-grams in the input data and use this information to make predictions about subsequent elements in the sequence.

A lot of previous research focused on understanding attention heads in isolation in in-context tasks. Theoreticians like to understand things crisply so we focus on attention-only networks where things are a bit simpler. A leading theory was that they simply count occurrences of n-grams in the context: math with fancy numerator etc. and denominator maybe? visual example where things are counted and previous occurrences are highlighted. This is a very simple mechanism that can explain some of the behavior of induction heads.

In particular, cite all the works that think about induction heads as counting n-grams.

However, as we shall see this is not the whole story. First, we will have a couple of observations from classical distribution learning scenarios:

1. counting is the MLE solution – so an extremal predictor
2. the optimal predictor is different – explain Bayes smoothing under a prior distribution
3. highlight that counting can be obtained as limits of some parameters but it’s actually a very brittle predictor

Now, we start exposing what we know about induction heads. Instead of thinking about the mechnaism directly, let’s look at the parameterization of a self-attention. This will be revealing in understanding what is actually implemented.

The first observation is the label-symmetry. There is no reason to overfit to a semantic content. Therefore we expect the network to simply have identity parameters comparing different positions and nothing more. The reason why we expect networks to copy previous content need some justification though. Then we explain why partial matches on first index instead of second index does not change anything so one should expect weights to be identities multiplied (i.e., same scalar for every position) with the same constant.

Then, we show that this is indeed the case in experiments via plots.

Then, we massage the object ans show how induction heads can interpolate between n-grams, allowing them to generalize beyond the specific patterns seen in the context. This interpolation capability enables the model to make predictions about sequences that are not explicitly present in the input, enhancing its overall performance in in-context learning tasks.

We have an interactive figure based on the visualizations from the paper here