IOI Circuit Analysis with TransformerLens

A set of hands-on notebooks that reverse-engineer how GPT-2 small solves the Indirect Object Identification (IOI) task, using three complementary mechanistic-interpretability techniques: direct logit attribution, activation patching, and path patching.

The IOI task itself is just the vehicle. For a prompt like

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

the model should predict the indirect object (Mary) over the repeated subject (John). Every notebook scores a run by the logit difference logit(IO) − logit(S) at the answer position, and the three techniques are different ways of attributing that logit difference to specific model components.

The focus of this repo is the methods — how each technique is set up, what signal it produces, and how the results chain together to localize a circuit.

---

The three techniques

1. Direct logit attribution (direct_logit_attribution.ipynb)

Question: which components write directly into the final logit difference?

DLA decomposes the residual stream at the answer position into per-component contributions and projects each one onto the answer direction. It is the cheapest technique — a single forward pass, no interventions — and it measures direct effects only (what a head writes straight to the logits, ignoring downstream routing).

The recipe, per attention head:

1. Run the clean prompt with run_with_cache and read each head's output blocks.{L}.attn.hook_z at the answer position.
2. Map it back into the residual stream with that head's output matrix: z @ W_O[L, h].
3. Apply the final LayerNorm scaling with cache.apply_ln_to_stack(...) so the contributions live in the same space as the unembed.
4. Project onto the IO − S unembed direction W_U[:, io] − W_U[:, s].

The per-head projections sum to the total logit difference, so you get a clean ranking of who is directly responsible.

Result: the name-mover head L9H9 dominates the direct contribution to the IO−S logit difference.

The notebook does this first on a single hand-written prompt, then wraps the same computation in logit_diff_for_sample / dla_for_sample and averages it over a real IOIDataset (100 templated prompts from transformer_lens.evals), reporting mean logit diff, accuracy, and the top-10 DLA heads.

2. Activation patching (activation_patching.ipynb)

Question: which components are causally necessary, including indirect effects?

Activation patching (a.k.a. causal tracing) runs the corrupted prompt but splices a clean activation into one location at a time, then measures how much of the clean logit difference is recovered. Where DLA only sees direct writes, patching catches components whose effect is routed through later layers.

Corruption here is the minimal S↔IO swap at the S2 position, which flips the intended answer. Each patch is scored as normalized recovery:

(patched_diff − corrupted_base) / (clean_base − corrupted_base)

0.0 = no effect, 1.0 = fully restores the clean behavior.

Two sweeps, both driven by TransformerLens run_with_hooks + a replace hook that copies from the clean cache:

• Residual-stream sweep over (layer, position) on blocks.{L}.hook_resid_pre — shows where and when the answer-relevant information lives, and how it moves from the S2 position into the END position as depth increases.
• Head sweep over (layer, head) on blocks.{L}.attn.hook_z at the END position — localizes the individual heads that carry the effect.

Results identified: S-inhibition heads at L8H6 / L8H10, name-mover heads around L9H9, and negative name-mover heads at L10H7 / L11H10.

The notebook builds from a single prompt up to a multi-prompt harness. Because the dataset generator emits length-matched prompts (same 14-token template, only single-token names/places/objects swapped), all prompts stack into one aligned batch and share a single clean cache. The harness (patch_resid_pre, patch_heads) mini-batches the sweep, caches only the hook it patches from to bound memory, and averages normalized recovery across the whole dataset.

3. Path patching (path_patching.ipynb)

Question: does component A affect the output specifically through component B?

Activation patching tells you a head matters; path patching tells you which downstream head it talks to. Here it tests the hypothesis surfaced by the previous two notebooks:

S-inhibition heads (L8) feed the name-mover head (L9H9).

The technique isolates a single sender → receiver edge by controlling every other pathway:

1. Do a clean run and a corrupted run, caching both.
2. Re-run the model while freezing all attention heads and MLPs to their clean values — except the sender head (L8H6), which is set to its corrupted value. This lets the sender's effect propagate, but only along real paths, with everything else held fixed.
3. Grab the resulting input to the receiver (the query input of L9's name-mover) and stash it.
4. In a final run, patch just that stashed input into blocks.9.attn.hook_q, leaving the rest of the model clean.
5. Score with the same normalized logit-difference metric.

Result: patching the full block-9 input recovers ~0.24 of the effect; restricting the path to L9H9 alone recovers ~0.10 — direct evidence that the S-inhibition → name-mover edge carries a real fraction of the circuit's behavior.

---

Supporting notebooks

• head_vis.ipynb — attention-pattern visualizations (circuitsvis) for the heads flagged above, to sanity-check what each head is attending to (e.g. name-movers attending from END to the IO token).
• transformer_lens_scratch.ipynb — scratch/reference notebook covering TransformerLens fundamentals: hook points, caching, and QK/OV circuit decomposition (including the induction-head copying + matching example).

How the techniques chain together

The three methods are deliberately sequenced — each one narrows what the next investigates:

DLA                 → finds the heads that directly move the logits (name-movers, L9H9)
Activation patching → finds all causally necessary heads, incl. indirect ones
                      (S-inhibition L8, negative name-movers L10/L11)
Path patching       → confirms the wiring between them (L8 S-inhibition → L9H9)

DLA is fast but blind to routing; activation patching adds causality but only localizes nodes; path patching resolves the edges. Together they reconstruct the IOI circuit rather than just correlating with it.

---

Setup

Requires Python 3.13. A .venv/ is already present; to recreate it:

python3.13 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

requirements.txt pins TransformerLens to a specific commit and installs it from GitHub, plus torch, transformers, circuitsvis, matplotlib, and Jupyter.

Then launch the notebooks:

source .venv/bin/activate
jupyter lab      # or: jupyter notebook

All notebooks boot GPT-2 small on CPU via TransformerBridge.boot_transformers("gpt2", device="cpu") followed by model.enable_compatibility_mode(), so no GPU is required.

Notes

• Metric. Everything is scored with the logit difference logit(IO) − logit(S) at the answer position; patching sweeps report normalized recovery against the clean/corrupted baselines.
• Length-matched prompts. The dataset generator only swaps single-token names/places/objects into one fixed template. Keeping every prompt the same length is what makes position-aligned batching (and averaged patching grids) possible.
• Memory. Full activation caches are large; the harnesses only cache the hook they read from and mini-batch the sweep. Lower batch_size if the kernel dies.