IOI circuit in GPT-2 small

• Pre-trained models. Steps 0-3 trained models from scratch. From now on, we use ones other people trained. The shift in mindset is "I don't get to choose the data, the architecture, or the seed; I just have weights and I have to figure out what they're doing."
• TransformerLens. Neel Nanda's library wraps any HuggingFace transformer and gives you clean access to every internal tensor. One-line model loading, one-line attention pattern access, one-line hooks for patching activations.
• Activation patching. A controlled experiment on the model's internals: swap a specific activation with one from a different prompt and measure whether the prediction changes. This converts circuit-finding from "the pattern looks suggestive" to "without this component the model fails in a predictable way."
• Logit difference. The standard IOI metric: how much does the model prefer the right answer over the wrong one, in raw logits? logit(correct) - logit(distractor). We track this throughout.

Everything else (attention, residual stream, circuits, etc.) you already know from steps 1-3.

---

Take the prompt:

"When John and Mary went to the store, John gave a drink to ___"

A fluent speaker fills in Mary. The reasoning: John is the subject of "gave," he can't be the indirect object too, so the recipient must be the other name introduced earlier - Mary.

GPT-2 small handles this correctly. The interesting question is how. The model has 12 layers and 144 attention heads. Some specific heads, working together, must be implementing "look at the names mentioned earlier, suppress the one that's also the subject of the current verb, predict the other one."

Wang et al. fully reverse-engineered this circuit. The names they gave to the head groups:

1. Duplicate token heads - notice that one name appears twice (John ... John).
2. Previous token heads - communicate "what came before me" along the sequence.
3. Induction heads - yes, the same induction heads from step 3; they do supporting work here.
4. S-inhibition heads - read the duplicate-token signal and write a "don't attend to S" message into the residual stream.
5. Name mover heads - at the final position, read the inhibition message and the list of candidate names, attend to the non-subject name (the IO), and copy it to the output.

This is a 5-component circuit with about 15 specific heads. We won't replicate all of it - only the cleanest, most striking part: finding the name mover heads via activation patching.

---

Suppose you've found some heads whose attention patterns look like they might be relevant. Are they actually doing the work? Visualisation alone can't tell you - a head might look meaningful but be downstream-ablatable with no effect.

The standard test:

1. Clean run: feed the model the original prompt. Cache every activation. Record the logit diff (positive - model gets it right).
2. Corrupted run: feed the model a near-identical prompt where the right answer is now the wrong one. E.g. swap the names: "When Mary and John went to the store, Mary gave a drink to ___" (correct answer is now John, not Mary). Logit diff on the original labels (logit(Mary) − logit(John)) is now strongly negative.
3. Patched run: re-run the corrupted prompt, but at one specific point in the model, replace that activation with its clean-run value. Then continue the forward pass. Measure the new logit diff.

If the patched activation is irrelevant: logit diff stays bad (corrupted-like). The activation didn't matter.

If the patched activation is the thing carrying the IO information: logit diff jumps back toward the clean value. The activation matters causally.

By repeating step 3 for every (layer, head) pair, you get a heatmap of causal importance. The bright cells are your circuit.

This is the tool for converting "this head's pattern looks right" into "this head is causally necessary." Wang et al. introduced it in this form; it's now the bread and butter of mech interp on real models.

---

The model. GPT-2 small (12 layers, 12 heads each, d_model=768). Loaded pre-trained via TransformerLens. We do not touch the weights.

The prompts. We use a small batch of IOI templates with different name pairs (John/Mary, James/Tom, …) for statistical robustness. For each, we construct:

• Clean: "When John and Mary went to the store, John gave a drink to" → expected answer " Mary".
• Corrupted (name-swap): "When Mary and John went to the store, Mary gave a drink to" → expected answer is now " John" (because Mary is now the subject and John is the indirect object).

We use name-swap corruption rather than the ABC variant from the original paper. ABC has a slightly cleaner causal interpretation but requires a third name that tokenises to one GPT-2 token. Name-swap guarantees identical token lengths automatically and gives a larger logit-diff swing (clean ~ +3 vs corrupt ~ −3 instead of vs ~ 0), which makes the patching results more visually clear. The trade-off is that name-swap corrupts multiple signals at once, so the recovery numbers should be interpreted as "this head transmits the IO identity" rather than the more nuanced statements the paper makes.

The metric. logit_diff = logit(IO) - logit(S), averaged over the batch, using the clean labels throughout.

• Clean logit diff: ~3-4 (model strongly prefers IO).
• Corrupted logit diff: ~−3 (model now prefers S, because S is the IO of the swapped sentence).

The patching sweep. For each of GPT-2 small's 144 attention heads (12 layers × 12 heads), we patch the head's output (z) at the final token position from the clean run into the corrupted run. We measure the new logit diff and express it as a fraction of recovery:

recovery = (patched_diff - corrupted_diff) / (clean_diff - corrupted_diff)

0% = patching had no effect. 100% = patching fully restored the clean prediction.

The expected result. A 12×12 heatmap of recovery percentages. Most cells will be near zero. A small number of cells in the middle-to-late layers will light up strongly - these are the name mover heads (commonly L9H9, L9H6, L10H0, L10H10 in GPT-2 small).

---

1. Go to colab.research.google.com.
2. File → Upload notebook → upload ioi_circuit.ipynb.
3. Runtime → Change runtime type → GPU. T4 is fine.
4. Runtime → Run all. First run takes(most is downloading GPT-2 and the patching sweep). Subsequent runs are faster (the model is cached).

If your Colab session keeps dying or out-of-memorying, swap GPT-2 small for a smaller model (e.g. gpt2, attn-only-1l, or pythia-70m) - HookedTransformer.from_pretrained supports many.

---

You now have a real tool - activation patching - that you can apply to any hypothesis about any pre-trained model. From here:

• Step 5 (Sparse autoencoders) will introduce automatic feature finding - instead of manually constructing prompts and patching one head at a time, an SAE recovers an entire dictionary of interpretable features at once. It closes the loop with step 1 (superposition).

Further reading

• Interpretability in the Wild: A Circuit for Indirect Object Identification in GPT-2 small - the original Wang et al. paper. Long but skimmable.
• Causal Scrubbing - Redwood Research's much stricter version of activation patching. The current gold standard for "is this really the circuit?"
• Path Patching - patches paths between specific components, not just whole activations. Lets you distinguish "head A directly affects head B" from "head A indirectly affects head B through everything else."
• ARENA Chapter 1.4 - guided exercises that reproduce more of the IOI circuit step by step.

A note on what mech interp actually looks like at this stage

This step is the most realistic depiction of day-to-day mech interp work. It's a lot of "hypothesise → set up clean/corrupted prompts → run a patching sweep → stare at a heatmap → repeat." Most papers in the field are extensions of this loop. Steps 1-3 were unusually clean experiments; from here on, the noise is real and arguing about evidence is most of the job.