The IOI circuit & activation patching

The IOI task

IOI stands for Indirect Object Identification. In "When John and Mary went to the store, John gave a drink to ___", the correct completion is Mary. The paired corrupted prompt swaps names so the correct answer flips. That pairing makes the task experimentally clean.

The metric is Logit difference : correct-name logit minus incorrect-name logit. A higher value means the model prefers the right answer more strongly.

import torch

logits = torch.tensor([3.2, 1.7])  # Mary, John
logit_diff = logits[0] - logits[1]
print(logit_diff.item())
# 1.5

From attribution to causality

Direct logit attribution asks which heads write in directions that help the answer logit. That is useful but correlational. A component can point in the right direction without being necessary.

Activation patching swaps an activation from the clean run into the corrupted run. Cache activations on a clean prompt, then run the corrupted prompt with one activation replaced. If the model's answer recovers, that activation carried causal information for the task.

def patch_head(activation, hook, head, clean_cache):
    activation[:, :, head] = clean_cache[hook.name][:, :, head]
    return activation

Noising and denoising ask opposite questions: does corrupting a clean run break behavior, or does restoring a corrupted run recover behavior? They license different conclusions. If restoring one activation recovers the behavior (denoising, what this notebook mostly does), that activation is sufficient to carry the task signal past the corruption. If corrupting one activation breaks the behavior (noising), that activation is necessary. Keep track of which experiment you ran before writing the conclusion.

The IOI circuit as a story

The discovered GPT-2-small circuit is roughly:

1. Duplicate-token heads (helped by induction heads) notice, at the second occurrence of the subject name, that this name is repeated.
2. S-inhibition heads (a cluster in layers 7-8) move that "this name is duplicated" signal to the final position, where it makes later heads pay less attention to the repeated name.
3. Name-mover heads (9.9, 9.6, and 10.0) attend to the remaining name and copy it into the output logits.
4. The weird-biology twist: negative name-mover heads (10.7 and 11.10) actively push against the correct answer, and backup name-mover heads partially take over when you ablate the main name movers. The model self-repairs under intervention.

This is the week where "a model has an algorithm" becomes concrete.

Pair-session guide

Core work is ARENA 1.4.1 sections 1 through 3: Model & Task Setup, Logit Attribution, and Activation Patching. Stretch work is section 4 ( path patching : patching a specific sender-to-receiver path instead of a whole activation) and section 5 (Full Replication: Minimal Circuits), plus a bonus section on anomalies like the negative and backup heads. Use a strict naming convention: clean_tokens, corrupted_tokens, clean_cache, corrupted_cache, patched_logits. Most bugs are swapped variables.

What you should see

You should see the residual-stream patching heatmap over layer and position light up in two localized spots: at the S2 position (the second occurrence of the subject name) in early layers, then jumping to the END position (the final token, where the prediction happens) around layers 7-8. Per-head patching should pick out the name-mover heads 9.9, 9.6, and 10.0 as strongly positive, with 10.7 and 11.10 strongly negative. The goal is explaining what intervention made the metric move.

Where to go next

Week 10 turns from proving circuits to reading model state with probes. The ARENA notebook's own recommended material for this week:

• Interpretability in the Wild by Wang et al. is the IOI paper itself; after the session you can read it as a report on experiments you have run.
• The authors' video walkthrough of the paper (an interview with Neel Nanda) explains why the task was chosen and how the circuit was found; ARENA encourages watching it.
• Some Lessons Learned from Studying Indirect Object Identification in GPT-2 small is the honest postmortem, including what the paper's methodology misses.