Info
Last Execution: 2026-08-19
| Package | Version |
|---|---|
| nnsight | 0.8 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0 |
Induction Heads and the In-Context Learning Phase Change¶
Introduction¶
🔌 Somewhere very early in training, a language model learns to do something it could not do at all a moment before: continue a pattern it has never seen, purely from the pattern's earlier occurrence in the same context. Olsson et al. report that this ability arrives abruptly, that it coincides with a visible bump in the training loss, and that what forms in that window is a specific mechanism — the induction head.
An induction head implements the rule [A][B] ... [A] → [B]. At a token that has appeared
before, it looks back to the earlier occurrence, attends to whatever came next that time, and
copies it to the output. That is a two-head circuit: a previous-token head in an early layer
writes "the token before me was X" into each position, and the induction head reads that as its
key. Elhage et al. lay out the algebra; Olsson et al. make the empirical claim that induction
heads are where in-context learning comes from.
The claim has two halves, and they are measured very differently:
- Mechanistic. A per-head induction score — on a sequence of random tokens repeated twice, how much attention probability sits on "the token that followed this one last time".
- Behavioural. An in-context learning score — how much better the model predicts token 500 of a document than token 50. A model that uses context well gets much better as context accumulates; a model that only knows unigram statistics does not.
The claim is that both change in the same window of training.
We test it on the Pythia suite, which publishes 154 intermediate checkpoints per model on the
Hub. In nnsight a checkpoint is a constructor argument — revision="step1000" — so a training
sweep is the same code run against a list of strings. We use pythia-160m-deduped (12 layers ×
12 heads) throughout.
📗 Primary result from In-context Learning and Induction Heads (Olsson et al., Anthropic 2022). Mechanism from A Mathematical Framework for Transformer Circuits (Elhage et al., 2021). Checkpoints from Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling (Biderman et al., ICML 2023).
Setup¶
If using Colab, install NNsight and datasets:
!pip install -U nnsight datasets
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight datasets
from IPython.display import clear_output
import torch
import torch.nn.functional as F
import nnsight
from nnsight import TransformersModel
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"
MODEL = "EleutherAI/pythia-160m-deduped"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = TransformersModel(
MODEL,
task="text-generation",
attn_implementation="eager", # SDPA silently returns attn_weights=None
dtype=torch.float32,
device_map=DEVICE,
dispatch=True,
)
clear_output()
n_layers = model.config.num_hidden_layers
n_heads = model.config.num_attention_heads
print(f"{MODEL}: {n_layers} layers x {n_heads} heads")
EleutherAI/pythia-160m-deduped: 12 layers x 12 heads
⚠️ attn_implementation="eager" is not optional here. Under the default SDPA kernel the
attention probabilities are never materialised and the tensor we are about to read comes back as
None. Everything in this notebook is an attention probability, so we ask for the eager path.
We load in float32. Pythia's checkpoints are published in fp32, and at 160M parameters the
whole sweep peaks around 3 GB of GPU memory.
1. What an induction head is¶
The attention probabilities are computed inside GPTNeoXAttention.forward; they are not the
module's output. nnsight exposes intermediate values inside a forward through .source,
which prints the real annotated source of the method with a hookable op name beside each line.
It works outside a trace and needs no forward pass, so it is the fastest way to find out what a
module you have never met computes.
print(model.gpt_neox.layers[0].attention.source)
* def forward(
0 self,
1 hidden_states: torch.FloatTensor,
2 attention_mask: torch.FloatTensor,
3 layer_past: Cache | None = None,
4 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
5 **kwargs: Unpack[FlashAttentionKwargs],
6 ):
7 input_shape = hidden_states.shape[:-1]
8 hidden_shape = (*input_shape, -1, 3 * self.head_size)
9
self_query_key_value_0 -> 10 qkv = self.query_key_value(hidden_states).view(hidden_shape).transpose(1, 2)
view_0 -> + ...
transpose_0 -> + ...
qkv_chunk_0 -> 11 query_states, key_states, value_states = qkv.chunk(3, dim=-1)
12
13 cos, sin = position_embeddings
apply_rotary_pos_emb_0 -> 14 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
15
16 # Cache QKV values
17 if layer_past is not None:
layer_past_update_0 -> 18 key_states, value_states = layer_past.update(key_states, value_states, self.layer_idx)
19
ALL_ATTENTION_FUNCTIONS_get_interface_0 -> 20 attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
21 self.config._attn_implementation, eager_attention_forward
22 )
23
24 # Compute attention
attention_interface_0 -> 25 attn_output, attn_weights = attention_interface(
26 self,
27 query_states,
28 key_states,
29 value_states,
30 attention_mask,
31 scaling=self.scaling,
32 dropout=0.0 if not self.training else self.attention_dropout,
33 **kwargs,
34 )
35
36 # Reshape outputs and final projection
attn_output_reshape_0 -> 37 attn_output = attn_output.reshape(*input_shape, -1).contiguous()
contiguous_0 -> + ...
self_dense_0 -> 38 attn_output = self.dense(attn_output)
39
40 return attn_output, attn_weights
41
The line we want is attention_interface_0, which returns (attn_output, attn_weights). So
layers[L].attention.source.attention_interface_0.output[1] is the [batch, head, dest, src]
tensor of post-softmax attention probabilities.
The stimulus. To see the mechanism we build a list of 16 unrelated words and repeat it,
prefixed with <|endoftext|>. The second copy is perfectly predictable from the context and
nothing else — no amount of English knowledge tells you that "guitar" follows "monkey", but the
first copy does.
⚠️ Pythia's tokenizer does not add a BOS token, and the first position of a sequence tends to
absorb leftover attention mass. We prepend <|endoftext|> (id 0) explicitly so that position
exists.
WORDS = ["river", "table", "monkey", "guitar", "copper", "velvet", "planet", "ladder",
"ocean", "pepper", "window", "tiger", "candle", "forest", "silver", "rocket"]
core = torch.tensor([model.tokenizer.encode(" " + w)[0] for w in WORDS])
demo_ids = torch.cat([torch.tensor([0]), core, core])[None] # [1, 2*16 + 1]
demo_labels = [f"{i} {model.tokenizer.decode([t])}" for i, t in enumerate(demo_ids[0])]
with torch.no_grad(), model.trace(input_ids=demo_ids.to(DEVICE)):
_, layer3 = model.gpt_neox.layers[3].attention.source.attention_interface_0.output
_, layer4 = model.gpt_neox.layers[4].attention.source.attention_interface_0.output
demo_attn = torch.stack([layer3[0, 0], layer4[0, 10]]).save() # L3H0 and L4H10
print("attention probabilities:", tuple(demo_attn.shape))
print("rows sum to:", demo_attn.sum(-1).min().item(), "-", demo_attn.sum(-1).max().item())
attention probabilities: (2, 33, 33) rows sum to: 0.9999998807907104 - 1.000000238418579
The row sums are 1.0, which confirms we are reading post-softmax probabilities and not scores. That is the one-line wiring check worth doing before any interpretation.
Now the patterns themselves, for two heads: L3H0, and L4H10. They are the two halves of the circuit — we picked them because the rest of §1 and §4 will show they are the strongest of their kind in this model, but look at the pictures before the labels.
patterns = demo_attn.float().cpu().numpy()
fig = px.imshow(
patterns,
facet_col=0,
x=demo_labels, y=demo_labels,
color_continuous_scale="Blues", zmin=0, zmax=1,
labels=dict(x="source (attended to)", y="destination (attending from)", color="attn prob"),
title="Two heads on a twice-repeated word list",
height=520,
)
fig.layout.annotations[0].text = "L3H0"
fig.layout.annotations[1].text = "L4H10"
fig.show()
T = len(core)
induction_head = patterns[1]
hits = 0
print("L4H10, destinations inside the second copy:")
for d in range(T + 1, 2 * T + 1):
target = d - (T - 1)
top = int(induction_head[d].argmax())
hits += top == target
if d < T + 6:
print(f" at {demo_labels[d]:>11} largest attention on {demo_labels[top]:>15}"
f" (induction target {demo_labels[target]:>11}, p={induction_head[d, target]:.3f})")
print(f" ...")
print(f" the induction target is the argmax for {hits} of {T} destinations")