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")
L4H10, destinations inside the second copy: at 17 river largest attention on 17 river (induction target 2 table, p=0.004) at 18 table largest attention on 0 <|endoftext|> (induction target 3 monkey, p=0.325) at 19 monkey largest attention on 4 guitar (induction target 4 guitar, p=0.870) at 20 guitar largest attention on 0 <|endoftext|> (induction target 5 copper, p=0.311) at 21 copper largest attention on 6 velvet (induction target 6 velvet, p=0.868) ... the induction target is the argmax for 13 of 16 destinations
Both panels are stripes; the offsets are the whole story.
L3H0 puts ~0.96 of its attention one position back, at every position — the stripe sits immediately below the diagonal. It is a previous-token head, and it takes no notice of the repetition: the second copy looks like the first.
L4H10 spends the whole first copy on <|endoftext|> (0.98–1.00 of its mass — the sink is where
attention goes when a head has nothing to do), and then in the second copy lights up along a stripe
displaced 15 positions from the diagonal: destination d attends to source d - 15, which for
a 16-token list repeated twice is exactly "the token that followed this one last time". The
induction target is the argmax for 13 of the 16 destinations; at the second copper the largest
attention is on velvet, p = 0.87. The three misses are all at the start of the second copy: the
first token has no repeated bigram to match on yet, and at two more the match is weak enough that
the sink wins outright — though even there the head puts p ≈ 0.31 on the right token. An induction
head is a soft match, not a lookup table.
Note what L4H10 is not doing: it does not attend to the earlier copy of the same token, which would be a stripe at offset 16 (heads that do that exist too, and are called duplicate-token heads). It attends one position later. That off-by-one is the mechanism, and it is also why L3H0 is in the picture — the key L4H10 matches its query against is what the previous-token head wrote. §4 returns to this.
The induction score. Words are legible but they are also words — a head could get some of this right from English. The standard stimulus is a sequence of uniformly random token ids, repeated twice, where nothing but the repetition predicts anything.
The score for one head is the mean attention probability on the induction offset. With a
sequence [BOS] + x[0:T] + x[0:T], a destination at position T+1+k (the k-th token of the
second copy) should attend to position k+2 (the token after the first copy of that token),
and (T+1+k) - (k+2) = T-1. So the whole measurement is one diagonal of the attention matrix,
at offset -(T-1), averaged over destinations that lie inside the second copy.
Taking the diagonal inside the trace matters. Written the obvious way, the measurement is one forward pass per layer and the whole attention tensor comes back to the host each time:
The slower, clearer version:
scores = torch.zeros(n_layers, n_heads)
for layer in range(n_layers):
with torch.no_grad(), model.trace(input_ids=rep_ids.to(DEVICE)):
_, attn = model.gpt_neox.layers[layer].attention.source.attention_interface_0.output
full = attn.save() # [8, 12, 257, 257] -- 25 MB, per layer
diagonal = torch.diagonal(full, offset=-(T_REP - 1), dim1=-2, dim2=-1)
scores[layer] = diagonal[..., 2:T_REP + 2].mean(dim=(0, -1)).cpu()
The version below reads all twelve layers in one pass and reduces inside the forward, so the
300 MB of attention probabilities is never materialised as a whole and only a [12, 8, 12, 130]
tensor — half a megabyte — leaves the model.
T_REP, B_REP = 128, 8
generator = torch.Generator().manual_seed(0)
rand_core = torch.randint(1000, 50254, (B_REP, T_REP), generator=generator)
rep_ids = torch.cat([torch.zeros(B_REP, 1, dtype=torch.long), rand_core, rand_core], dim=1)
def head_diagonals(model, ids, offset):
"""One diagonal of every head's attention matrix -> [layer, batch, head, position]."""
n_layers = model.config.num_hidden_layers
with torch.no_grad(), model.trace(input_ids=ids.to(DEVICE)):
per_layer = []
for layer in range(n_layers):
_, attn = model.gpt_neox.layers[layer].attention.source.attention_interface_0.output
per_layer.append(torch.diagonal(attn, offset=offset, dim1=-2, dim2=-1))
stacked = torch.stack(per_layer).save()
return stacked
def induction_score(model):
diag = head_diagonals(model, rep_ids, -(T_REP - 1))
return diag[..., 2:T_REP + 2].mean(dim=(1, -1)).float().cpu() # [layer, head]
scores = induction_score(model)
fig = px.imshow(
scores,
color_continuous_scale="Blues", zmin=0, zmax=1,
labels=dict(x="head", y="layer", color="induction score"),
title="Induction score, pythia-160m-deduped (final checkpoint)",
height=480,
)
fig.show()
print(f"chance level (uniform over {rep_ids.shape[1]} positions): {1 / rep_ids.shape[1]:.4f}")
print(f"median head: {scores.median():.4f}")
top = scores.flatten().topk(6)
for value, index in zip(top.values.tolist(), top.indices.tolist()):
print(f" L{index // n_heads}H{index % n_heads:<3} {value:.3f}")
chance level (uniform over 257 positions): 0.0039 median head: 0.0031 L4H10 0.914 L4H6 0.886 L5H6 0.864 L4H11 0.843 L4H8 0.823 L6H6 0.706
The picture is bimodal, not graded. A handful of heads in layers 4–6 are above 0.7; the median head is at 0.003, which is the chance level for a 257-position sequence. Nothing sits in between in any interesting quantity.
That matters for what follows. Because the population splits so cleanly, "does this checkpoint
have induction heads?" is a question with a yes/no answer rather than a threshold argument, and
max over heads is a fair one-number summary of a checkpoint.
2. The same measurement, across training¶
Pythia publishes 154 checkpoints per model as git branches on the Hub: step0, step1,
step2, step4, … step512, then step1000, step2000, … step143000. Pythia trains with a
batch of 1024 sequences × 2048 tokens, so one step is 2.097M tokens and step143000 is 300B.
TransformersModel takes revision= and forwards it to from_pretrained, so a checkpoint is a
string:
TransformersModel(MODEL, revision="step1000", ...)
Before plotting anything across 18 checkpoints, check that this argument does what it says. The
last checkpoint should be the released model, so revision="step143000" and the default
main should produce identical logits on the same input.
def final_checkpoint_logits(ids):
checkpoint = TransformersModel(MODEL, revision="step143000", task="text-generation",
attn_implementation="eager", dtype=torch.float32,
device_map=DEVICE, dispatch=True)
with torch.no_grad(), checkpoint.trace(input_ids=ids):
logits = checkpoint.output.logits.save()
return checkpoint.revision, logits
revision, checkpoint_logits = final_checkpoint_logits(demo_ids.to(DEVICE))
with torch.no_grad(), model.trace(input_ids=demo_ids.to(DEVICE)):
released_logits = model.output.logits.save()
clear_output()
print(f"model.revision : {revision}")
print(f"logits torch.equal(main) : {torch.equal(checkpoint_logits, released_logits)}")
model.revision : step143000 logits torch.equal(main) : True
Bit-identical. The revision plumbing is real, and every number below is a plot of actual weights.
One trap, and it is expensive. model.trace() snapshots the entire calling frame's locals
in order to build the traced block's scope, and that snapshot outlives the with block — it dies
with the frame. So del checkpoint inside the function that traced it will not free the weights,
and a checkpoint loop written at top level holds every model it has loaded. Writing the
per-checkpoint work as a function, as above and below, is what keeps memory flat: the snapshot
goes away when the function returns. At 160M parameters this is 0.65 GB and easy to miss; at 70B
it is the difference between running and not.
The sweep below does all four measurements in a single pass per checkpoint, because the expensive part is the load:
- induction score — the repeated-random-token diagonal, as above;
- previous-token score — the same code with
offset=-1on natural text (§4); - in-context learning score — per-document, on 512 Pile documents (§5);
- mean loss — on the same documents, as the control (§3, §6).
⚠️ This downloads one fp32 checkpoint per revision — 649 MB each, about 12 GB for the 18 below —
and takes roughly 10 minutes on one GPU. Shrink STEPS if you want the shape of the result
faster; the interval that matters is (512 → 1000].
from datasets import load_dataset
CTX, N_DOCS, BATCH = 512, 512, 4
EARLY, LATE = slice(44, 54), slice(499, 511)
pile = load_dataset("NeelNanda/pile-10k", split="train")
documents = []
for record in pile:
ids = model.tokenizer(record["text"])["input_ids"]
if len(ids) >= CTX:
documents.append(ids[:CTX])
if len(documents) == N_DOCS:
break
documents = torch.tensor(documents)
clear_output()
print("documents:", tuple(documents.shape))
documents: (512, 512)
STEPS = [0, 2, 4, 8, 16, 32, 64, 128, 256, 512,
1000, 2000, 4000, 8000, 16000, 32000, 64000, 143000]
TOKENS_PER_STEP = 1024 * 2048
def measure(step):
"""Everything we need from one checkpoint. A function, so the model is freed on return."""
checkpoint = TransformersModel(MODEL, revision=f"step{step}", task="text-generation",
attn_implementation="eager", dtype=torch.float32,
device_map=DEVICE, dispatch=True)
induction = induction_score(checkpoint)
prev_token = head_diagonals(checkpoint, documents[:8], -1)[..., 5:].mean(dim=(1, -1)).float().cpu()
early, late, total = [], [], []
for i in range(0, documents.shape[0], BATCH):
chunk = documents[i:i + BATCH].to(DEVICE)
with torch.no_grad(), checkpoint.trace(input_ids=chunk):
logits = checkpoint.output.logits.save()
nll = F.cross_entropy(logits[:, :-1].flatten(0, 1), chunk[:, 1:].flatten(),
reduction="none").view(chunk.shape[0], -1).cpu()
early.append(nll[:, EARLY].mean(1))
late.append(nll[:, LATE].mean(1))
total.append(nll.mean(1))
del logits, nll
return {"induction": induction, "prev_token": prev_token,
"icl": torch.cat(late) - torch.cat(early), "loss": torch.cat(total)}
results = {}
for step in STEPS:
results[step] = measure(step)
clear_output(wait=True)
print(f"{'step':>8} {'tokens':>9} {'max induction':>14} {'heads>0.5':>10}"
f" {'max prev-token':>15} {'ICL':>8} {'loss':>7}")
for done, record in results.items():
print(f"{done:>8} {done * TOKENS_PER_STEP:>9.3g} {record['induction'].max():>14.4f}"
f" {int((record['induction'] > 0.5).sum()):>10}"
f" {record['prev_token'].max():>15.3f}"
f" {record['icl'].mean():>+8.3f} {record['loss'].mean():>7.3f}")
step tokens max induction heads>0.5 max prev-token ICL loss
0 0 0.0059 0 0.010 -0.018 11.088
2 4.19e+06 0.0059 0 0.010 -0.018 11.087
4 8.39e+06 0.0059 0 0.010 -0.019 11.007
8 1.68e+07 0.0061 0 0.010 -0.016 10.672
16 3.36e+07 0.0065 0 0.011 +0.025 9.977
32 6.71e+07 0.0073 0 0.012 +0.047 9.469
64 1.34e+08 0.0069 0 0.011 -0.028 8.651
128 2.68e+08 0.0067 0 0.019 -0.150 7.255
256 5.37e+08 0.0060 0 0.170 -0.190 6.116
512 1.07e+09 0.0154 0 0.502 -0.104 5.017
1000 2.1e+09 0.9613 4 0.913 -0.376 3.900
2000 4.19e+09 0.9751 4 0.940 -0.320 3.318
4000 8.39e+09 0.9613 5 0.931 -0.272 3.032
8000 1.68e+10 0.9501 5 0.923 -0.253 2.871
16000 3.36e+10 0.9362 6 0.902 -0.251 2.782
32000 6.71e+10 0.9224 6 0.883 -0.248 2.720
64000 1.34e+11 0.9366 6 0.862 -0.225 2.673
143000 3e+11 0.9140 8 0.831 -0.215 2.701
That table is the result, and it is worth reading down the max induction column before looking
at any plot.
Ten checkpoints, covering the first billion tokens of training — step0 (random init) through
step512 — sit between 0.006 and 0.015, at or barely above chance.
The very next published checkpoint, step1000, is at 0.961, with four heads above 0.5. Every
checkpoint after that is between 0.91 and 0.98. There is no checkpoint at which induction is
half-formed.
labels = [f"step{s}" for s in STEPS]
max_induction = torch.tensor([results[s]["induction"].max() for s in STEPS])
n_strong = [int((results[s]["induction"] > 0.5).sum()) for s in STEPS]
fig = go.Figure()
fig.add_bar(x=labels, y=n_strong, name="heads with score > 0.5", yaxis="y2", opacity=0.3)
fig.add_scatter(x=labels, y=max_induction, mode="lines+markers", name="max induction score")
fig.update_layout(title="Induction score across Pythia checkpoints",
xaxis_title="checkpoint (the grid is not linear in training)",
yaxis=dict(title="max induction score", range=[0, 1]),
yaxis2=dict(title="heads > 0.5", overlaying="y", side="right", range=[0, 12]),
height=460)
fig.show()
heatmap = torch.stack([results[s]["induction"].flatten() for s in STEPS])
fig = px.imshow(heatmap.T, x=labels, color_continuous_scale="Blues", zmin=0, zmax=1,
labels=dict(x="checkpoint", y="head (layer * 12 + head)", color="induction score"),
title="Per-head induction score, all 144 heads", height=560, aspect="auto")
fig.show()
The per-head heatmap says the transition is not a gradual recruitment of many heads either. The
same small set — L4H10, L4H11, L5H6, L4H6 — goes from floor to saturated together, and they are
the heads that still carry induction at step143000, 300 billion tokens later. A few more heads
join at low scores over the rest of training, but the circuit that exists at the end is
essentially the circuit that appeared in that one interval.
3. What this checkpoint grid can and cannot resolve¶
A step function drawn on a coarse grid is the oldest trap in training-dynamics work: a smooth curve sampled sparsely looks abrupt too. Two things have to be said before this counts as evidence.
The grid. Pythia's published checkpoints are log-spaced to step512 and then linear at every
1000 steps. There is no checkpoint between step512 and step1000 — not one we skipped, one
that does not exist. The gap is 488 steps, a factor of 1.95 in training tokens.
The control. If the abruptness were an artifact of that gap, then everything would look abrupt across it. So we ask where the largest single-interval drop in overall loss is, on the same checkpoints and the same documents.
intervals = list(zip(STEPS[:-1], STEPS[1:]))
mean_loss = {s: results[s]["loss"].mean().item() for s in STEPS}
peak_induction = {s: results[s]["induction"].max().item() for s in STEPS}
rises = [max(peak_induction[b] - peak_induction[a], 0.0) for a, b in intervals]
biggest_induction = intervals[int(torch.tensor(rises).argmax())]
drops = [mean_loss[a] - mean_loss[b] for a, b in intervals]
biggest_loss = intervals[int(torch.tensor(drops).argmax())]
print(f"{'interval':>18} {'x tokens':>9} {'d(max induction)':>17} {'d(mean loss)':>13}")
for (a, b), rise, drop in zip(intervals, rises, drops):
print(f"{f'({a} -> {b}]':>18} {b / max(a, 1):>8.2f}x {peak_induction[b] - peak_induction[a]:>+17.4f}"
f" {-drop:>+13.4f}")
print(f"\nlargest induction rise : ({biggest_induction[0]} -> {biggest_induction[1]}]"
f" {max(rises):+.4f}, {max(rises) / sum(rises):.0%} of all positive change")
print(f"largest loss drop : ({biggest_loss[0]} -> {biggest_loss[1]}] {-max(drops):+.4f}")
print(f"loss across the induction interval: "
f"{mean_loss[biggest_induction[0]]:.3f} -> {mean_loss[biggest_induction[1]]:.3f}")
interval x tokens d(max induction) d(mean loss)
(0 -> 2] 2.00x +0.0000 -0.0019
(2 -> 4] 2.00x +0.0001 -0.0798
(4 -> 8] 2.00x +0.0002 -0.3347
(8 -> 16] 2.00x +0.0004 -0.6951
(16 -> 32] 2.00x +0.0008 -0.5079
(32 -> 64] 2.00x -0.0005 -0.8181
(64 -> 128] 2.00x -0.0002 -1.3960
(128 -> 256] 2.00x -0.0007 -1.1384
(256 -> 512] 2.00x +0.0094 -1.0997
(512 -> 1000] 1.95x +0.9459 -1.1168
(1000 -> 2000] 2.00x +0.0138 -0.5819
(2000 -> 4000] 2.00x -0.0138 -0.2860
(4000 -> 8000] 2.00x -0.0112 -0.1607
(8000 -> 16000] 2.00x -0.0139 -0.0893
(16000 -> 32000] 2.00x -0.0138 -0.0616
(32000 -> 64000] 2.00x +0.0142 -0.0469
(64000 -> 143000] 2.23x -0.0226 +0.0276
largest induction rise : (512 -> 1000] +0.9459, 96% of all positive change
largest loss drop : (64 -> 128] -1.3960
loss across the induction interval: 5.017 -> 3.900
Two things in that table.
The rise is concentrated. The (512 → 1000] interval is not merely the largest single-interval
increase in induction; it is 96% of all the positive change across the entire trajectory.
Everything before it is noise around chance, and everything after it is drift of ±0.02.
The control passes. If the abruptness came from the coarse grid, the same interval would look
abrupt for every quantity. It does not. The largest single-interval drop in mean loss is at
(64 → 128], −1.40 nats — a different interval entirely. Across the induction interval the loss
falls 5.02 → 3.90, −1.12 nats, an unremarkable step on this curve. The step function is specific
to induction, not to the sampling.
And here is what we cannot say. The transition is bracketed to (512, 1000] — less than one
doubling of training, 1.07e9 → 2.10e9 tokens — and no tighter. Whether it takes 50 steps or 400
inside that window is not a question these checkpoints can answer, because the checkpoints that
would answer it were never published, and no reweighting of this analysis changes that. Finer
resolution means training a model and saving your own.
4. Previous-token heads form first¶
The induction snap itself is trapped in the unpublished gap, but its prerequisite is not.
Elhage et al. describe the induction circuit as two heads composed: a previous-token head writes
"the token before me was X" into position i, and the induction head uses that as the key it
matches its query against. If the story is right, previous-token heads should be in place
before induction heads appear — and that ordering is resolvable on this grid, because the two
curves are separated by more than one checkpoint.
The previous-token score is the same diagonal measurement with offset=-1, on natural text
rather than repeated tokens: how much of a head's attention sits on the immediately preceding
token.
max_prev = torch.tensor([results[s]["prev_token"].max() for s in STEPS])
fig = go.Figure()
fig.add_scatter(x=labels, y=max_prev, mode="lines+markers", name="max previous-token score")
fig.add_scatter(x=labels, y=max_induction, mode="lines+markers", name="max induction score")
fig.update_layout(title="The prerequisite arrives first",
xaxis_title="checkpoint", yaxis_title="score", height=450)
fig.show()
print(f"{'step':>8} {'max prev-token':>15} {'heads>0.3':>10} {'max induction':>14}")
for s in STEPS:
prev = results[s]["prev_token"]
print(f"{s:>8} {prev.max():>15.3f} {int((prev > 0.3).sum()):>10}"
f" {results[s]['induction'].max():>14.3f}")
step max prev-token heads>0.3 max induction
0 0.010 0 0.006
2 0.010 0 0.006
4 0.010 0 0.006
8 0.010 0 0.006
16 0.011 0 0.007
32 0.012 0 0.007
64 0.011 0 0.007
128 0.019 0 0.007
256 0.170 0 0.006
512 0.502 4 0.015
1000 0.913 4 0.961
2000 0.940 5 0.975
4000 0.931 5 0.961
8000 0.923 6 0.950
16000 0.902 6 0.936
32000 0.883 7 0.922
64000 0.862 5 0.937
143000 0.831 8 0.914
The prerequisite arrives first, and this time the grid is fine enough to see it.
The previous-token score is still 0.019 at step128, is 0.170 at step256, and reaches 0.502 at
step512, with four heads above 0.3, while the induction score is still 0.015. The head is L3H0,
and it is the same head at every checkpoint from step512 onward. A full doubling of training
before induction appears, the model already has the component that induction consumes.
This also independently supports "abrupt". Induction does not simply follow its own precursor: the
precursor ramps smoothly across step128 → step512 while induction stays at the floor, and then
induction jumps in a single interval.
5. The behavioural half: in-context learning¶
Olsson et al. define the in-context learning score as the loss on the 500th token of a context minus the loss on the 50th. It is a difference within a single forward pass, so it isolates how much a model gains from having more context, and it is not confounded by the model simply getting better at everything.
We use small windows rather than single positions — tokens 500–511 minus tokens 45–54 — and score each document separately, on 512 Pile documents (Pythia's own training distribution).
⚠️ The between-document variance of this metric is enormous relative to the effect we are looking for. An unpaired estimate on a few dozen documents will put individual checkpoints on the wrong side of zero. The fix is to score the same 512 documents at every checkpoint and test the checkpoint-to-checkpoint difference per document — a paired test, in which the between-document variance cancels.
from scipy.stats import ttest_rel
icl = {s: results[s]["icl"] for s in STEPS}
icl_mean = [icl[s].mean().item() for s in STEPS]
icl_sem = [(icl[s].std() / len(icl[s]) ** 0.5).item() for s in STEPS]
fig = go.Figure()
fig.add_scatter(x=labels, y=icl_mean, mode="lines+markers", name="ICL score",
error_y=dict(type="data", array=icl_sem))
fig.add_hline(y=0, line_dash="dot")
fig.update_layout(title="In-context learning score (loss at token ~500 minus loss at token ~50)",
xaxis_title="checkpoint", yaxis_title="ICL score (lower = more in-context learning)",
height=450)
fig.show()
print(f"{'interval':>18} {'delta':>9} {'sem':>7} {'t':>8} {'p':>10}")
paired = []
for a, b in intervals:
difference = icl[b] - icl[a]
t, p = ttest_rel(icl[b].numpy(), icl[a].numpy())
paired.append((difference.mean().item(), a, b, t, p))
print(f"{f'({a} -> {b}]':>18} {difference.mean():>+9.4f}"
f" {difference.std() / len(difference) ** 0.5:>7.4f} {t:>+8.2f} {p:>10.1e}")
best = min(paired)
print(f"\nlargest ICL improvement: ({best[1]} -> {best[2]}] "
f"delta={best[0]:+.4f} t={best[3]:+.2f} p={best[4]:.1e}")
interval delta sem t p
(0 -> 2] -0.0000 0.0001 -0.26 7.9e-01
(2 -> 4] -0.0008 0.0033 -0.23 8.2e-01
(4 -> 8] +0.0025 0.0140 +0.18 8.6e-01
(8 -> 16] +0.0418 0.0285 +1.47 1.4e-01
(16 -> 32] +0.0218 0.0157 +1.39 1.6e-01
(32 -> 64] -0.0752 0.0225 -3.35 8.7e-04
(64 -> 128] -0.1218 0.0309 -3.94 9.2e-05
(128 -> 256] -0.0399 0.0282 -1.41 1.6e-01
(256 -> 512] +0.0854 0.0327 +2.61 9.3e-03
(512 -> 1000] -0.2716 0.0537 -5.06 6.0e-07
(1000 -> 2000] +0.0556 0.0267 +2.09 3.8e-02
(2000 -> 4000] +0.0486 0.0203 +2.39 1.7e-02
(4000 -> 8000] +0.0185 0.0161 +1.15 2.5e-01
(8000 -> 16000] +0.0027 0.0142 +0.19 8.5e-01
(16000 -> 32000] +0.0027 0.0144 +0.18 8.5e-01
(32000 -> 64000] +0.0232 0.0130 +1.79 7.5e-02
(64000 -> 143000] +0.0093 0.0163 +0.57 5.7e-01
largest ICL improvement: (512 -> 1000] delta=-0.2716 t=-5.06 p=6.0e-07
The largest single-interval improvement in in-context learning, across all seventeen intervals, is
(512 → 1000]: d = −0.272 ± 0.054 per document, t = −5.06, p = 6e−07 — the induction interval,
and more than twice the size of the next largest.
That next largest is (64 → 128], d = −0.122, p = 9e−05. It is real, and it is before any
induction head exists. It also coincides with the biggest drop in overall loss (§3), so it looks
like a model getting generically better rather than a circuit switching on. Part of in-context
learning is not induction.
⚠️ The error bars are the reason for the paired design. The standard error on a single checkpoint's ICL score is about ±0.07 — larger than several of the interval differences we just resolved at p < 0.01. Between-document variance dominates the metric; it cancels only because every checkpoint is scored on the same 512 documents and the test is on the per-document difference.
One more thing the plot shows: the ICL score peaks at step1000 (−0.376) and then gets weaker
for the remaining 300 billion tokens, ending at −0.215. The phase change overshoots.
6. The correlation is not the evidence¶
The tempting next move is to correlate the two trajectories across checkpoints and report the coefficient. It comes out high. It is also not evidence, and the control that shows why costs one line: correlate in-context learning against overall loss instead, which improves over training for every reason there is.
import numpy as np
induction_mass = np.array([results[s]["induction"].sum().item() for s in STEPS])
icl_gain = -np.array(icl_mean) # higher = more in-context learning
loss_gain = -np.array([mean_loss[s] for s in STEPS]) # higher = better model
print(f"r(total induction, ICL gain) = {np.corrcoef(induction_mass, icl_gain)[0, 1]:+.2f} <- the result")
print(f"r(-mean loss, ICL gain) = {np.corrcoef(loss_gain, icl_gain)[0, 1]:+.2f} <- the control")
after = [i for i, s in enumerate(STEPS) if s >= 1000]
print(f"\nafter the transition (step1000 onward, n={len(after)}):")
print(f" total induction mass {induction_mass[after][0]:.2f} -> {induction_mass[after][-1]:.2f}")
print(f" ICL score {icl_mean[after[0]]:+.3f} -> {icl_mean[after[-1]]:+.3f}")
print(f" r(total induction, ICL gain) = {np.corrcoef(induction_mass[after], icl_gain[after])[0, 1]:+.2f}")
r(total induction, ICL gain) = +0.72 <- the result r(-mean loss, ICL gain) = +0.90 <- the control after the transition (step1000 onward, n=8): total induction mass 4.65 -> 11.67 ICL score -0.376 -> -0.215 r(total induction, ICL gain) = -0.92
r = +0.72 looks like a reproduction. It is not one.
The control — in-context learning against overall loss, a quantity with no mechanistic content whatsoever — comes out at +0.90, stronger than the effect. Both trajectories improve over training, as does everything else about the model, so a whole-trajectory correlation across checkpoints cannot distinguish "induction heads produce in-context learning" from "training produces both". Had we reported the +0.72 and stopped, we would have reported a statistic that our own null hypothesis generates more of.
This is the most useful thing in the notebook, so it is worth being blunt: the correlation coefficient is not the evidence. The evidence is the interval coincidence — the single largest improvement in in-context learning falls in the one interval where induction heads appear, at p = 6e−07 — together with the loss control in §3 showing that this interval is unremarkable for everything except induction.
The second printout gives a further reason to distrust the coefficient. After the transition, the
two quantities move in opposite directions: from step1000 to step143000 total induction
mass more than doubles (4.65 → 11.67, heads above 0.2 going 5 → 19) while the ICL score weakens
(−0.376 → −0.215), r = −0.92. This is consistent with Olsson et al., whose claim concerns a
co-occurring phase change and not a monotone relationship — but a reader who sampled only
checkpoints after step1000 would have run the same correlation and reported the opposite sign.
Caveats¶
- We measure prefix matching, not copying. Olsson et al. define an induction head by two properties: it attends back to the token that followed the earlier occurrence (prefix matching) and its output raises the logit of the token it attended to (copying). The induction score here is entirely the first. A head with this attention pattern that writes something else into the residual stream would score 0.96 and not be an induction head; checking the other half means looking at the OV circuit, or ablating.
- Nothing here is causal. Two curves move together in one interval. The strongest evidence in the paper is an intervention — knock the induction heads out at test time and in-context learning collapses — and we have not run it. That is the natural next notebook, not this one.
- One model, one seed.
pythia-160m-deduped, 12 × 12. Olsson et al. study 34 transformers. Re-running all of the above onpythia-70m-dedupedis a one-line change toMODEL, and is the cheapest check that the interval is not a property of this particular training run. - The bracket is one doubling wide, and that is worth repeating because it is the easiest thing to overstate.
- The windows are arbitrary. Olsson et al. say as much about tokens 50 and 500; we average small windows around those positions to reduce variance. We also truncate contexts at 512 tokens, a quarter of Pythia's training context, so "late in context" here is shallower than the model was trained for.
- We do not see the loss bump. The paper reports a visible bump in the training loss during the phase change. Our loss column is held-out loss at eighteen checkpoints, and at this resolution it is smooth through the transition — a limit of the sampling, not a contradiction.
Conclusion¶
🎉 The result reproduces, in the form the paper states it rather than the form that is easiest to
plot. Induction heads in pythia-160m-deduped go from chance to saturated inside a single published
interval; that interval carries 96% of all the change in induction across 300 billion tokens; the
largest improvement in in-context learning anywhere on the trajectory is the same interval at
p = 6e−07; and the previous-token head the circuit depends on is already half formed a doubling
earlier. The naive correlation between the two trajectories is beaten by a control correlation with
overall loss, so it is the coincidence of intervals, not the coefficient, that carries the claim.
Three pieces of nnsight did the work. revision="stepN" turns a checkpoint sweep into a loop over
strings, with no separate checkpoint-loading code path. .source reaches the attention
probabilities, which are an intermediate inside GPTNeoXAttention.forward and not any module's
output — and print(module.source) finds the op name without a forward pass. And reducing inside
the trace keeps 300 MB of attention on the GPU, which is why the sweep runs in about 3 GB.
⚠️ One idiom to carry away: put the per-checkpoint work in a function. A trace snapshots its
calling frame's locals, so a model loaded and traced at top level is not freed by del until that
frame exits. At 160M parameters that is invisible; on a sweep of larger checkpoints it is the
difference between a run and an OOM.
Reclaiming the ~12 GB of checkpoints this notebook downloads:
from huggingface_hub import scan_cache_dir
cache = scan_cache_dir()
repo = next(r for r in cache.repos if r.repo_id == "EleutherAI/pythia-160m-deduped")
strategy = cache.delete_revisions(*[rev.commit_hash for rev in repo.revisions])
print(strategy.expected_freed_size_str)
strategy.execute()
Related: Attention Heads for reading attention
patterns in general, .source for values computed inside a
forward, Dual-Route Induction for what induction heads look
like in a modern model, and Grokking Progress Measures for the
other classic phase-change reproduction.
References¶
- Olsson, Elhage, Nanda, Joseph, DasSarma, Henighan, Mann, Askell, Bai, Chen, Conerly, Drain, Ganguli, Hatfield-Dodds, Hernandez, Johnston, Jones, Kernion, Lovitt, Ndousse, Amodei, Brown, Clark, Kaplan, McCandlish, Olah, In-context Learning and Induction Heads, Transformer Circuits Thread, 2022
- Elhage, Nanda, Olsson, Henighan, Joseph, Mann, Askell, Bai, Chen, Conerly, DasSarma, Drain, Ganguli, Hatfield-Dodds, Hernandez, Jones, Kernion, Lovitt, Ndousse, Amodei, Brown, Clark, Kaplan, McCandlish, Olah, A Mathematical Framework for Transformer Circuits, Transformer Circuits Thread, 2021
- Biderman, Schoelkopf, Anthony, Bradley, O'Brien, Hallahan, Khan, Purohit, Prashanth, Raff, Skowron, Sutawika, van der Wal, Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling, ICML 2023
- Gao, Biderman, Black, Golding, Hoppe, Foster, Phang, He, Thite, Nabeshima, Presser, Leahy,
The Pile: An 800GB Dataset of Diverse Text for Language
Modeling, 2020 — the corpus, via the
NeelNanda/pile-10ksample