Activation patching, in code
By the end: You can run an activation patching experiment, choose a sensible metric, and read the result as a map of where information flows.
- Clean and corrupted inputs
- A metric: logit difference
- Patching the residual stream by layer and position
- Patching individual heads
- What patching does and does not show
New to interpretability? A trained network is a very large pile of numbers that produces good answers while nobody can say why. Interpretability is the laboratory work of finding out: feeding in controlled inputs and watching which parts respond, testing whether a specific fact can be read back out of an internal representation, switching a component off to see which behaviour disappears with it. It is closer to studying an unfamiliar organism than to reading source code. Why it matters
Reading activations tells you what is present in a model. Patching tells you what the output actually depends on, and it is the technique most interpretability claims rest on. This section runs one experiment end to end: two prompts differing by a single token, a metric that measures the decision you care about, and activations copied from one run into the other until you can see which locations carry the difference. It closes on what the result does and does not license you to say, which is where most of the mistakes are made.
The experiment#
Pick a behaviour. A well-studied one: in 'When Mary and John went to the store, John gave a drink to', GPT-2 predicts 'Mary', the name that was not repeated. Make a corrupted version that flips the answer by changing one token. Then copy activations from the clean run into the corrupted run, one location at a time, and see which ones restore the clean answer.
from functools import partial
import torch
clean = model.to_tokens("When Mary and John went to the store, John gave a drink to")
corrupt = model.to_tokens("When Mary and John went to the store, Mary gave a drink to")
mary, john = model.to_single_token(" Mary"), model.to_single_token(" John")
def logit_diff(logits):
return (logits[0, -1, mary] - logits[0, -1, john]).item()
_, clean_cache = model.run_with_cache(clean)
print(logit_diff(model(clean)), logit_diff(model(corrupt))) # positive, then negativeWhy logit difference
It measures exactly the decision you care about, Mary against John, and is roughly linear in the model's internals. Accuracy is too coarse and probability saturates. The two prompts must have the same number of tokens so positions line up.
Choosing the corruption is the decision that shapes everything afterwards, and it is made before any code runs. Swapping one name asks which components carry that name. Replacing the name with a random token asks a different question, and will light up components that handle unusual tokens rather than components that handle names. Replacing it with noise asks a third. None is wrong; they are three different experiments, and a result only means anything alongside a statement of which corruption produced it. The rule of thumb is that the corruption should change the answer and nothing else you can name.
Patch every layer and position#
def patch(resid, hook, pos):
resid[:, pos, :] = clean_cache[hook.name][:, pos, :]
return resid
n_pos = clean.shape[1]
result = torch.zeros(model.cfg.n_layers, n_pos)
for layer in range(model.cfg.n_layers):
for pos in range(n_pos):
logits = model.run_with_hooks(
corrupt,
fwd_hooks=[(f"blocks.{layer}.hook_resid_pre", partial(patch, pos=pos))],
)
result[layer, pos] = logit_diff(logits)
# plot result as a heatmap: layers down, positions acrossTypically the information sits at the position of the changed name in early layers and moves to the final position in later layers. The heatmap shows when and where it travels. Patching the output of each attention head at the final position then narrows it down to the heads that move it.
The same experiment, small enough to run#
The three-token task below has a known answer, which is the point: you can check what patching says against what the task provably requires. Each input is an item, a second item, and a cue naming which of the two to answer about. The answer is the cued item's colour. The forward pass takes an optional patch argument that overwrites one position at one named location with a value copied from elsewhere, which is the whole of activation patching in four lines.
import torch, torch.nn as nn
torch.manual_seed(0)
N_ITEMS, D, N_HEADS = 48, 32, 4
CUE_FIRST, CUE_SECOND = N_ITEMS, N_ITEMS + 1
def colour(item): return item % 4 # the attribute the task uses
def shape(item): return (item // 4) % 3 # in the token, never used by the task
class Block(nn.Module):
def __init__(self):
super().__init__()
self.attn = nn.MultiheadAttention(D, N_HEADS, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(D, 4 * D), nn.ReLU(), nn.Linear(4 * D, D))
def forward(self, x):
x = x + self.attn(x, x, x, need_weights=False)[0]
return x + self.mlp(x)
class Toy(nn.Module):
def __init__(self):
super().__init__()
self.embed = nn.Embedding(N_ITEMS + 2, D)
self.pos = nn.Parameter(torch.randn(3, D) * 0.02)
self.blocks = nn.ModuleList([Block(), Block()])
self.out = nn.Linear(D, 4)
def forward(self, tokens, cache=None, patch=None):
def at(name, x):
if cache is not None:
cache[name] = x.detach().clone()
if patch is not None and patch[0] == name:
x = x.clone(); x[:, patch[1]] = patch[2] # copy one position in
return x
x = at("embeddings", self.embed(tokens) + self.pos)
for i, blk in enumerate(self.blocks):
x = at(f"after_block_{i}", blk(x))
return self.out(x[:, -1])
def make_batch(n):
a = torch.randint(0, N_ITEMS, (n,)); b = torch.randint(0, N_ITEMS, (n,))
cue = torch.randint(0, 2, (n,))
toks = torch.stack([a, b, torch.where(cue == 0, CUE_FIRST, CUE_SECOND)], 1)
return toks, torch.where(cue == 0, colour(a), colour(b))
model = Toy()
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
for step in range(800):
toks, labels = make_batch(256)
loss = nn.functional.cross_entropy(model(toks), labels)
opt.zero_grad(); loss.backward(); opt.step()
with torch.no_grad():
t, l = make_batch(4000)
print(f"held-out accuracy: {(model(t).argmax(-1) == l).float().mean():.3f}")
# held-out accuracy: 1.000Now the clean and corrupted pair. They differ in exactly one token, and that token flips the answer, which is what makes the difference between the two runs attributable to one thing. The metric is the difference between the two candidate logits.
# Item 4 is colour 0, item 5 is colour 1. One token differs between the two runs.
clean = torch.tensor([[4, 9, CUE_FIRST]]) # the answer should be colour 0
corrupt = torch.tensor([[5, 9, CUE_FIRST]]) # the answer should be colour 1
ANSWER, ALTERNATIVE = colour(4), colour(5)
def logit_diff(logits):
"""Positive means the clean answer is winning. Exactly the decision we care about."""
return (logits[0, ANSWER] - logits[0, ALTERNATIVE]).item()
clean_cache = {}
with torch.no_grad():
hi = logit_diff(model(clean, cache=clean_cache))
lo = logit_diff(model(corrupt))
print(f"clean run {hi:+.2f}")
print(f"corrupted run {lo:+.2f}")
# clean run +15.53
# corrupted run -28.12Those two numbers are the ends of the scale. Patching a location restores some fraction of the distance from the corrupted value back to the clean one, and reporting that fraction rather than the raw logit makes locations comparable across prompt pairs with different absolute margins.
print("\nrecovered share of the clean logit difference, one location at a time")
print(" pos 0 (item 1) pos 1 (item 2) pos 2 (cue)")
for name in ("embeddings", "after_block_0", "after_block_1"):
row = []
for pos in range(3):
with torch.no_grad():
out = model(corrupt, patch=(name, pos, clean_cache[name][:, pos]))
row.append(100 * (logit_diff(out) - lo) / (hi - lo))
print(f"{name:16s}" + "".join(f"{v:13.1f}%" for v in row))recovered share of the clean logit difference, one location at a time
pos 0 (item 1) pos 1 (item 2) pos 2 (cue)
embeddings 100.0% 0.0% 0.0%
after_block_0 81.0% -0.2% 21.4%
after_block_1 0.0% 0.0% 100.0%Read the first column downwards. At the embeddings, patching position 0 restores everything, because at that depth position 0 is the only place the difference exists. After one block it restores 81% and the final position has picked up 21.4%: the copy is under way. After two blocks position 0 restores nothing, because by then the answer has been moved and position 0 is no longer read. The middle column is zero everywhere, correctly, since item 2 is not involved in this pair at all. That is what a clean patching result looks like, and having seen one makes a messy one on a real model much easier to interpret.
The control that makes it a result#
Patching found where the answer travels. It has not yet shown that the model is using the colour rather than the token. Each item also has a shape, which the task never asks about, and items 0, 4 and 8 all share colour 0 while differing in shape. If the model were keying on token identity rather than colour, swapping among those three would move the answer.
# The control: change the shape of item 1 and hold its colour fixed.
print("\nitems 0, 4 and 8 are all colour 0 and differ in shape")
for tok in (4, 0, 8, 5):
with torch.no_grad():
d = logit_diff(model(torch.tensor([[tok, 9, CUE_FIRST]])))
print(f" item {tok}: colour {colour(tok)}, shape {shape(tok)} -> logit diff {d:+7.2f}")
# items 0, 4 and 8 are all colour 0 and differ in shape
# item 4: colour 0, shape 1 -> logit diff +15.53
# item 0: colour 0, shape 0 -> logit diff +23.57
# item 8: colour 0, shape 2 -> logit diff +13.13
# item 5: colour 1, shape 1 -> logit diff -28.12Changing the shape moves the margin around, from 13.13 to 23.57, and never changes the decision. Changing the colour takes it to -28.12 and flips it. So shape is present in the activations, since it visibly perturbs them, and the output does not depend on it. That pair of facts is the difference between probing and patching stated in four numbers, and it is worth quoting in an interview because it is concrete.
One pair is an anecdote#
Everything above came from a single prompt pair, which is exactly the mistake this field makes most often. The fix costs almost nothing on a model this size: run the same procedure over two hundred pairs with different items and report the spread rather than one number.
import statistics
torch.manual_seed(1)
rows = {name: [[] for _ in range(3)] for name in ("embeddings", "after_block_0", "after_block_1")}
for _ in range(200):
a, b = torch.randint(0, N_ITEMS, (2,)).tolist()
other = torch.randint(0, N_ITEMS, (1,)).item()
if colour(a) == colour(other):
continue # the pair must flip the answer
cl = torch.tensor([[a, b, CUE_FIRST]]); co = torch.tensor([[other, b, CUE_FIRST]])
ans, alt = colour(a), colour(other)
diff = lambda lg: (lg[0, ans] - lg[0, alt]).item()
cc = {}
with torch.no_grad():
h = diff(model(cl, cache=cc)); l = diff(model(co))
for name in rows:
for pos in range(3):
with torch.no_grad():
out = model(co, patch=(name, pos, cc[name][:, pos]))
rows[name][pos].append(100 * (diff(out) - l) / (h - l))
print("\nover 200 prompt pairs: median recovered share, and the 10th-90th percentile")
for name in rows:
cells = []
for pos in range(3):
v = sorted(rows[name][pos]); n = len(v)
cells.append(f"{statistics.median(v):5.0f}% [{v[n//10]:4.0f},{v[9*n//10]:4.0f}]")
print(f"{name:16s}" + " ".join(cells))
# over 200 prompt pairs: median recovered share, and the 10th-90th percentile
# embeddings 100% [ 100, 100] 0% [ 0, 0] 0% [ 0, 0]
# after_block_0 85% [ 75, 95] 0% [ -0, 1] 14% [ 5, 21]
# after_block_1 0% [ 0, 0] 0% [ 0, 0] 100% [ 100, 100]The cost of that replication is worth noticing, because it is what decides whether you can afford it on a real model. Patching every location for one prompt pair is one forward pass per location: three depths by three positions is nine passes here, and twelve layers by fifteen positions is 180 on GPT-2 small. Two hundred pairs multiplies that by two hundred. On a toy model the whole thing finishes in five seconds; on a real model the same experiment is where an afternoon goes, which is the honest reason people report one pair. Batch the pairs rather than looping them and most of that cost disappears.
Reading it honestly#
- Patching shows a location is sufficient to carry the difference between these two prompts. Other prompts may use other routes.
- The choice of corruption decides what you find. A different corruption asks a different question.
- Components can compensate for each other, so knocking one out may show less effect than it really has.
- Restoring the clean answer is not the same as destroying it. Patching clean into corrupted measures sufficiency; patching corrupted into clean measures necessity. Say which you ran.
- A recovered share above 100% or below 0% is normal and means the patch overshot or moved things the wrong way. It is information, not a bug.
What to carry, and what to skip#
- Carry: clean prompt, minimally corrupted prompt, a logit difference, and a normalised recovered share. Those four decisions are the whole method and an interviewer will ask for them in that order.
- Carry: the shape control above. Being able to say in one sentence how you would show a representation is not used separates a strong answer from a recited one.
- Skip: patching every location before you have a hypothesis. Patch the residual stream by position first to find the route, then narrow to components along it.
- Skip: building a plotting pipeline. A printed table is enough to read a result and much faster to trust.
Where people go wrong
Running patching on one prompt pair and describing 'the circuit' for the behaviour. It is the circuit for that pair until shown otherwise.
If this comes up in an interview
Walk me through an activation patching experiment.
Choose a behaviour and a clean prompt, make a corrupted prompt that changes the answer by a minimal edit, pick a metric such as the logit difference between the two answers, then copy activations from the clean run into the corrupted run one location at a time and measure how much of the clean behaviour each restores.
Then they push further
Why is this stronger evidence than probing?
A probe shows information can be decoded from a location. Patching intervenes: if replacing that activation changes the output, the model's computation actually depends on it there.
What could make a patching result misleading?
An unrepresentative corruption, a metric that saturates, one prompt standing in for a whole behaviour, and backup components that take over when one is changed. Replicate across many prompt pairs and corruptions.
What gets marked up. Strong answers explain clean, corrupted and metric, stress the causal nature, and volunteer the limitations.
Build this: Reproduce the residual stream heatmap above on twenty prompt pairs with different names, average them, then patch individual head outputs at the final position and list the five heads with the largest effect.
This is one section of Interpretability, 26 hours across 6 sections. Read it in full to get the whole topic in order.