Activation Patching and Causal Tracing

#ai #llm #machinelearning #programming

Activation patching answers a causal question that no amount of staring at activations can: is this component necessary for this behaviour? You run the model twice on two carefully chosen prompts, copy one internal tensor from the second run into the first, and measure how far the output moves.

The experiment in one paragraph

Take two prompts that differ in exactly one respect that changes the right answer. Run the model on both, caching every internal activation. Now run it a third time on the first prompt, but at one chosen site — layer L, token position p, and possibly a single head — overwrite the activation with the one cached from the second prompt. If the output flips towards the second prompt’s answer, that site carries information the model uses to decide. If nothing moves, it does not.

The technique is causal mediation analysis borrowed from statistics; Vig and colleagues brought it into NLP in 2020, and Meng, Bau and colleagues made it widely known in 2022 as causal tracing in the ROME work on locating factual associations in GPT. Sweeping every (layer, position) pair produces a heatmap, and that heatmap is the image most people have seen without knowing what generated it.

Building the clean and corrupted pair

Everything rests on the pair. A sloppy pair produces a beautiful heatmap that means nothing.

Same token length, ideally the same tokens except one. If the two prompts tokenise to different lengths, position p is not the same thing in both runs and the patch is incoherent. Check the token ids, not the strings — a leading space changes tokenisation.
The difference must change the answer. “The Eiffel Tower is in” against “The Colosseum is in” differs in the subject and in the correct completion. That is the useful kind of pair.
Decide what “corrupted” means. Three common choices: substitute a different real subject, replace the subject embedding with Gaussian noise, or replace it with the embedding of a randomly chosen token. Noise is the original causal tracing recipe; a real alternative subject is generally cleaner, because noise pushes activations off the distribution the model was trained on and off-distribution behaviour is not the behaviour you are trying to explain.

Which direction to patch, and why it matters

There are two experiments here and they answer different questions. People routinely report one and describe the other.

Direction	Description
denoising	Run on the corrupted prompt, patch in the clean activation. Asks: is this site sufficient to restore the behaviour? This is the classic causal-tracing direction and it tends to light up broadly, because many sites can restore an answer.
noising	Run on the clean prompt, patch in the corrupted activation. Asks: is this site necessary? Much stricter, and much sparser results, because the model often has redundant paths and breaking one changes little.
zero ablation	Replace the activation with zeros rather than with another run's value. Cheap, but zero is not a neutral value in a residual stream — it is a specific off-distribution point — so a large effect may just mean you broke the model.
mean ablation	Replace with the mean activation over a distribution of prompts. Better behaved than zero ablation because it stays closer to the manifold, and it is the sensible default when you have no natural corrupted counterpart.

Report which one you ran. A site that is sufficient but not necessary is a real and interesting finding, and it is not the same finding as “this is where the fact is stored”.

Choosing the metric

The output moved — by how much? Three metrics, in increasing order of how much they tell you:

Probability of the correct token. Intuitive but saturating: once the probability is near one, a large logit change registers as nothing.
Logit difference between the clean answer and the corrupted answer. Linear in the quantity the model actually computes, insensitive to the softmax squashing everything at the top, and the standard choice in circuit work for exactly those reasons.
Normalised recovery. Rescale the logit difference so that the unpatched corrupted run is 0 and the clean run is 1. Now a patch that recovers 0.6 is comparable across prompts, which is what you need to average over a dataset instead of showing one cherry-picked example.

The code

Forward hooks only, so this works on any HuggingFace causal model. The one detail that catches everyone: a decoder layer’s forward returns a tuple, so a hook that returns a bare tensor silently breaks the next layer.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "gpt2"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name).eval()
layers = model.transformer.h          # GPT-2; Llama-style: model.model.layers

clean     = "The Eiffel Tower is located in the city of"
corrupted = "The Colosseum is located in the city of"

ids_clean = tok(clean, return_tensors="pt")["input_ids"]
ids_corr  = tok(corrupted, return_tensors="pt")["input_ids"]
assert ids_clean.shape == ids_corr.shape, "pair must tokenise to equal length"

answer_clean = tok(" Paris")["input_ids"][0]
answer_corr  = tok(" Rome")["input_ids"][0]

def logit_diff(logits):
    """Positive means the model prefers the clean answer."""
    last = logits[0, -1]
    return (last[answer_clean] - last[answer_corr]).item()

# --- 1. cache every layer's residual-stream output on the CORRUPTED prompt ---
cache = {}

def make_cacher(i):
    def hook(module, args, output):
        cache[i] = output[0].detach().clone()   # output[0] = hidden states
    return hook

handles = [l.register_forward_hook(make_cacher(i)) for i, l in enumerate(layers)]
with torch.no_grad():
    base_corr = logit_diff(model(ids_corr).logits)
for h in handles:
    h.remove()

with torch.no_grad():
    base_clean = logit_diff(model(ids_clean).logits)

# --- 2. patch one (layer, position) at a time into the CLEAN run -------------
def make_patcher(i, pos):
    def hook(module, args, output):
        hidden = output[0].clone()
        hidden[0, pos, :] = cache[i][0, pos, :]
        return (hidden,) + tuple(output[1:])     # MUST stay a tuple
    return hook

seq_len = ids_clean.shape[1]
results = torch.zeros(len(layers), seq_len)

for i in range(len(layers)):
    for pos in range(seq_len):
        h = layers[i].register_forward_hook(make_patcher(i, pos))
        with torch.no_grad():
            ld = logit_diff(model(ids_clean).logits)
        h.remove()
        # 1.0 = unaffected, 0.0 = fully flipped to the corrupted answer
        results[i, pos] = (ld - base_corr) / (base_clean - base_corr)

print(f"clean {base_clean:+.2f}   corrupted {base_corr:+.2f}")
for i in range(len(layers)):
    row = "  ".join(f"{v:.2f}" for v in results[i])
    print(f"layer {i:>2}  {row}")

Read the output as a grid. Columns are token positions, rows are layers, and a cell near 0 means patching there destroyed the clean answer — that site was carrying the information. On a factual-recall prompt you should see the effect concentrated at the subject tokens in early-to-middle layers and again at the final position in later layers, which is the pattern the causal-tracing literature describes.

This loop is n_layers × seq_len forward passes. On GPT-2 small with a ten-token prompt that is 120 passes and takes seconds on a CPU. On a 32-layer model with a 200-token prompt it is 6,400 passes, which is why the approximations below exist.

Path patching and the cheaper approximations

Patching a residual stream site tells you the information is there at that point. It does not tell you which route carried it, because the residual stream is a sum of everything written into it so far.

Path patching is the refinement: instead of overwriting the residual stream, you patch the input to one specific downstream component while leaving every other consumer of that site reading the original value. That isolates the edge rather than the node, and it is what turns a heatmap into a wiring diagram. It is how the indirect-object identification circuit was resolved into named heads with named roles rather than a collection of important layers.

Attribution patching is the cheap approximation: use a first-order gradient estimate of what the patch would have done, which needs two forward passes and one backward pass for the entire sweep rather than one pass per site. It is a linear approximation of a non-linear system, so it is unreliable where the effect is large. The standard use is as a filter — screen everything with attribution patching, then verify the top candidates with real patches.

What a patching result does and does not license

A patching sweep gives you a genuinely causal claim about a specific prompt distribution. It supports: on prompts of this shape, the model’s output depends on the activation at these sites.

It does not support several things it is regularly used to support. It does not show that a fact is stored at the site — only that the computation routes through it, which is also true of a site that merely relays. Hase and colleagues made this concrete in 2023 by showing that where causal tracing localises a fact and where an edit to change that fact succeeds are not the same place, which breaks the intuitive chain from tracing to editing. It does not generalise beyond the prompt distribution you built, and templated pairs are a narrow distribution. And a null result is weak evidence: redundancy means a component can be genuinely involved and still removable without effect, because another path picks up the work.

The falsifying test for a patching-derived story is the one worth building into the harness from the start. If your account says heads A and B jointly implement the behaviour, then patching both simultaneously should have an effect that your account predicts — and patching a matched set of control heads with similar activation statistics should not. Run that comparison, or the story is just the heatmap restated in words.