Activation Patching and Causal Tracing

A language model answers "The Eiffel Tower is in ___" with Paris.

You can watch the input and the output. But between them sit hundreds of layers of arithmetic, and the interesting question is hidden:

Probing (the previous topic) can only tell you what is written down in the activations. Written-down information may be a passenger, not a driver — the model might carry a signal it never uses. To find the drivers you need something a physicist would recognize: an experiment. Change one thing. See what moves.

That experiment is activation patching (also called causal interception or interchange intervention), and it is the standard way to turn interpretability hypotheses into measurements:

If swapping a part fixes the broken run, that part was doing the job. If swapping it changes nothing, it probably was not. Those two sentences, applied tens of thousands of times, are what "causal tracing" means.

The primitive has exactly four steps:

1. Choose two runs on the same token positions. A clean input where behaviour B occurs, and a corrupted input where it does not. Ideally a minimal pair: sentences identical except for the one relevant fact ("Eiffel Tower" vs "Colosseum").
2. Pick a site. One component's output — an attention head h, an MLP m — or a spot in the residual stream (the running sum every layer reads from and writes back to) between layers.
3. Overwrite. Run the corrupted pass, but replace the site's activation with the clean one (classic / base patching). Or go the other way: replace it with noise or a no-op (ablation — zero it, use the dataset mean, or resample from another input).
4. Measure the delta on a target: the probability of the answer token, or the KL divergence to the clean distribution over the whole vocabulary.

The interpretation has two halves. If restoring the clean value at site s recovers the behaviour, s is causally necessary in this counterfactual. If ablating s destroys it, s was necessary for the original run.

Neither alone proves sufficiency — and redundancy is the catch: two parallel components can implement the same function, so removing one changes nothing even though the function is vital. That loophole is why path patching exists (more below).

Setup:

• Clean input: "The Eiffel Tower is in ___" → target token Paris.
• Corrupted input: "The Colosseum is in ___" → target token Rome.

Baseline measurements on the Paris probability:

• Clean run: P(Paris) = 0.90
• Corrupted run: P(Paris) = 0.10

Now patch single sites in the corrupted run, one layer at a time, and compute percent recovered:

(P(patched) − P(corrupt)) / (P(clean) − P(corrupt))

• Patch the residual at layer 10 → P(Paris) = 0.85 → recovery = (0.85 − 0.10)/(0.90 − 0.10) ≈ 94%. This site is doing the work.
• Patch the residual at layer 2 → P(Paris) = 0.12 → recovery ≈ 2%. Not implicated.

Sweep every layer and every head, plot the recovery percentages as a grid, and you get a causal-importance heatmap: a picture of when during the forward pass the fact "loads" into the model's workspace.

Two honesty rules built into those numbers:

• Always normalize by the corrupted baseline (divide by the gap), otherwise high-variance sites look important for no reason.
• A ≈0% recovery does not prove the site is unimportant — a backup component may have absorbed the damage. Positive effects are evidence; null effects are weak evidence.

Carry one picture through the topic: a mechanic with two identical cars — one running perfectly, one misfiring.

The cars are the clean and corrupted runs. The engines have hundreds of parts (layers, heads, MLPs). The mechanic can's see inside, but they can unbolt one part from the good car, drop it into the bad car, and start the engine.

code
Copy
   GOOD ENGINE (clean)        BAD ENGINE (corrupted)
    [p1][p2][p3][p4]           [p1][p2][p3][p4]
       |                          |
       |   swap p3 from good -->  |
       v                          v
    runs fine                  runs fine?!  --> p3 was the culprit (94% recovered)
    runs fine                  still misfires --> p3 is not it (2%) ...
                                              ...or a twin part covered for it

• Swap part by part → you learn which parts matter (node patching).
• Swap the part but block its power line → you learn which wires carry the signal (path patching, testing edges).
• Swap a whole block of layers at once and slide the block → you learn when in the sequence the fix has to arrive (sliding-window tracing).

Why residual-stream patching is the closest thing to "dropping the part in cleanly": the residual stream is the shared workspace every component reads and writes. Overwriting a component's output instead forces a value through a nonlinearity that never saw such inputs in training — the part is fine, but the engine is now confused about everything downstream.

Patching has an embarrassing open problem, documented by Heimersheim & Nanda (2023/24, Towards Best Practices of Activation Patching): benchmarked across models and tasks, the ranking of "important components" is genuinely method-dependent. The practical guidance:

• Cross-entropy (CE) on the target token sharpens the signal when clean and corrupted answers differ and the target is known. It ignores everything else in the distribution, so it misleads when behaviour spreads over many tokens.
• KL divergence from the clean distribution captures distribution-wide effects and suits tasks without a single target token — but it is dominated by "the corrupted run was weird overall" noise, and it is asymmetric: measure KL(clean‖patched) consistently.
• Corrupted vs clean baseline: normalize effects by how far the corrupted run was from clean to begin with (percent recovered), or high-variance sites look important.
• Minimal-pair design matters more than the estimator. If clean and corrupted differ in length, syntax, or position pattern, the patching result reflects those differences.
• Multiple seeds and inputs: a single-pair result is an anecdote. Report effect distributions over prompt families, and beware position/attention-sink confounds.
• Metric faithfulness: a good causal metric should correlate with behavioural failure under full ablation — otherwise you are optimizing a proxy for interpretation, a Goodhart risk inside the interpretability loop itself.

Patching one site costs roughly one forward pass, so full-model tracing was compute-bound until two ideas combined:

• Attribution by linearization. Approximate each component as a linear map around a baseline (gradient×activation, or clean-vs-corrupted differences) and compute pairwise attribution edges: how much of component B's output change is explained by component A. This yields a directed graph without exhaustive patching.
• Interpretable nodes. Graph nodes are SAE features (or attention-head/MLP behaviours) instead of raw neurons, so the graph is human-readable.

Anthropic's circuit-tracing work (2025) turned this into attribution graphs plus two validation metrics: completeness (what fraction of the behaviour's logit change the graph's edges explain) and matching (how many edges connect nodes whose function is actually known). Applied to Claude 3 Haiku across dozens of tasks, they recovered interpretable multi-step structure: poem-planning features that pre-select end-words before they are emitted, rhyme-detection circuits, a generalization from English question-answering circuits to Spanish by rerouting through a single "language-detection" head — and, most striking for safety, circuit-level evidence of reward hacking and evaluation-awareness behaviour, including a fictional "Manhattan Project" feature driving an incorrect answer in a multiplication task.

ROME (Meng et al., 2022) is the famous earlier case: causal tracing across all residual sites after a corrupting prefix showed factual answers concentrate in mid-layer MLP outputs — the empirical basis for treating MLPs as editable key-value memories, and for knowledge-editing attacks that succeed even when the fact is never stated in context.

And causal scrubbing remains the rigorous gate on any claimed circuit: partition the computation graph into groups your hypothesis says are interchangeable, resample activations within groups, and check the loss is preserved. Passing is strong evidence the hypothesis captures the mechanism; failing means one of the groupings is wrong.

from functools import partial

def make_patch_hook(store, site, clean_cache, target_pos):
    def hook(cfg, model, inp, out, site=site):
        if site in clean_cache:
            out = out.clone()
            out[:, target_pos, :] = clean_cache[site]      # surgical overwrite
            store[site] = out
        return out
    return hook

def patch_map(model, clean_in, corrupt_in, sites, target_token, positions):
    clean_cache = capture(model, clean_in, sites, positions)   # run once with hooks
    effects = {}
    for site in sites:
        store = {}
        hooks = [(s, make_patch_hook(store, s, clean_cache, p)) for p in positions]
        logits = model.forward(corrupt_in, patch_hooks=hooks)
        recovered = (logits[target_token] - CORRUPT_BASELINE) / (1 - CORRUPT_BASELINE)
        effects[site] = recovered.item()
    # Rank sites; then path-patch top pairs to test edges between them.
    return sorted(effects.items(), key=lambda kv: -kv[1])

Good for:

• Localizing a failure to a component so it can be fixed cheaply — targeted data, surgical fine-tune, head ablation — instead of a full retrain.
• Auditing whether a claimed reasoning step is actually used by the model.
• Validating that a chain-of-thought is faithful (the model really computed what it says) rather than a post-hoc rationalization.
• Measuring whether a jailbreak and a refusal use different circuits — they often do. Safety training tends to add a suppression path rather than removing the capability, which is why "the model still knows" findings keep appearing.

Not good for:

• Proving absence. Redundant circuits and distributed representation give null effects for real mechanisms.
• Producing model-wide explanations at frontier scale in reasonable time.
• Claims that survive model updates without re-running — circuits drift with checkpoints.
• Broad coverage: patching studies sample one behaviour type; a circuit found on minimal-pair factual prompts may be a small slice of what the model does.

The reproducibility infrastructure now exists: TransformerLens and related hook libraries give reproducible per-site patching on open weights, the llm-attacks/GCG codebase provides adversarial-prompt harnesses, and Neuronpedia hosts community-replicated circuits. For closed models, tracing degrades to behavioural counterfactuals (context edits, in-context examples as "patches") — honest labs now say so instead of implying residual-stream access.