Activation patching: which components carry the signal?

Take a prompt the model gets right (clean) and one it gets wrong (corrupt). Run the corrupt prompt, but splice in a single component’s activation from the clean run. If the model’s answer snaps back toward correct, that component was carrying the information the corrupt prompt destroyed.

Repeat for every component and you have a causal map, with no assumption that the component writes to the logits directly.

Key questions

Which heads restore the behavior when patched from the clean run?
Do they match the ones direct logit attribution found, or are they different?

Structure

Set up
Establish the two baselines
Patch a single head
Sweep every head
Save

Model and data. GPT-2 small in float32. It is small enough that bfloat16 rounding visibly degrades its predictions, so the notebooks load it explicitly in float32. Eight indirect-object prompt pairs from murano.tasks.ioi.

Requirements. pip install murano-interp[plot]

1. Set up

import torch

from murano import MuranoModel, NodeSet, Pipeline, keys
from murano.nodes import Node
from murano.plotting import plot_sweep, save_figure
from murano.steps import (
    LoadPaired,
    LogitDiffStep,
    Logits,
    Patch,
    RecoveredMetricStep,
    Save,
    Sweep,
)

OUTPUT_DIR = "murano_outputs/activation_patching"

# GPT-2 is small enough that bfloat16 rounding degrades its predictions.
model = MuranoModel("gpt2", dtype=torch.float32)
print(f"{model.model_id}: {model.n_layers} layers, {model.n_heads} heads")

gpt2: 12 layers, 12 heads

This notebook uses the indirect-object-identification task from murano.tasks. Each clean prompt names two people and then has the subject give a drink, so the next token should be the other name, the indirect object. The corrupt prompt swaps who gives, which flips the answer.

clean  : When Mary and John went to the store, John gave a drink to ___  ->  " Mary"
corrupt: When Mary and John went to the store, Mary gave a drink to ___  ->  " John"

The metric throughout is the logit difference: the correct name’s logit minus the incorrect name’s. Positive means the model prefers the right answer.

from murano.tasks import ioi

task = ioi(n=8)
print("clean  :", task.clean[0], "->", repr(task.correct[0]))
print("corrupt:", task.corrupt[0], "->", repr(task.incorrect[0]))
print(f"{len(task.clean)} prompt pairs")

clean  : When Mary and John went to the store, John gave a drink to -> ' Mary'
corrupt: When Mary and John went to the store, Mary gave a drink to -> ' John'
8 prompt pairs

2. Establish the two baselines

LoadPaired writes both prompt sets into Results. We run the model twice, once per set, into different keys, and score each. Everything after this is measured against these two numbers.

baseline = Pipeline(
    [
        LoadPaired(task),
        Logits(model),
        Logits(
            model,
            prompts_key=keys.CORRUPT_PROMPTS,
            logits_key=keys.CORRUPT_LOGITS,
            mask_key=keys.CORRUPT_MASK,
            targets=None,
        ),
        LogitDiffStep(task.correct, task.incorrect, output_key="ld_clean", model=model),
        LogitDiffStep(
            task.correct,
            task.incorrect,
            logits_key=keys.CORRUPT_LOGITS,
            mask_key=keys.CORRUPT_MASK,
            output_key="ld_corrupt",
            model=model,
        ),
    ]
).run()

print(f"clean   logit difference: {baseline['ld_clean'].value:+.4f}  (right name)")
print(f"corrupt logit difference: {baseline['ld_corrupt'].value:+.4f}  (wrong name)")

clean   logit difference: +2.3971  (right name)
corrupt logit difference: -3.1112  (wrong name)

3. Patch a single head

Patch runs the corrupt prompt as its base and splices in activations from the clean run. That is the “denoising” direction: start broken, add back one piece, see how much behavior returns.

RecoveredMetricStep normalizes the result:

recovered = (patched - corrupt) / (clean - corrupt)

so 0 means the patch changed nothing and 1 means it fully restored clean behavior.

Results.copy() forks the baseline, so the two forward passes above are not re-run for every head we try.

patched = Pipeline(
    [
        Patch(model, targets=[Node(8, "self_attn", head=6)]),
        LogitDiffStep(
            task.correct,
            task.incorrect,
            logits_key=keys.PATCHED_LOGITS,
            mask_key=keys.PATCHED_MASK,
            output_key="ld_patched",
            model=model,
        ),
        RecoveredMetricStep("ld_clean", "ld_corrupt", "ld_patched"),
    ]
).run(baseline.copy())

print(f"patching L8H6: recovered {patched[keys.RECOVERED].value:+.3f}")

patching L8H6: recovered +0.240

4. Sweep every head

That was a fork, three steps, and a read. Doing it once per head is what Sweep is: it takes the items to sweep, a callback building the chain for one item, and the key to harvest. It forks the incoming Results per item, so each head is measured against the same two baselines and the intermediate keys never leak back into the pipeline.

144 patched forward passes on GPT-2 small take a couple of seconds, so there is no reason to guess which heads to try.

swept = Pipeline(
    [
        Sweep(
            over=NodeSet.product(
                range(model.n_layers), ["self_attn"], heads=range(model.n_heads)
            ),
            steps=lambda head: [
                Patch(model, targets=[head]),
                LogitDiffStep(
                    task.correct,
                    task.incorrect,
                    logits_key=keys.PATCHED_LOGITS,
                    mask_key=keys.PATCHED_MASK,
                    output_key="ld_patched",
                    model=model,
                ),
                RecoveredMetricStep("ld_clean", "ld_corrupt", "ld_patched"),
            ],
            read=keys.RECOVERED,
        )
    ]
).run(baseline.copy())

recovered = swept[keys.SWEEP]
print(f"{len(recovered.scores)} heads scored, harvested from {recovered.primary!r}")

144 heads scored, harvested from 'recovered'

# The scores straddle zero, so plot_sweep centers a diverging scale on it: a head
# with no effect is drawn in the neutral color rather than shaded as if it had a sign.
fig = plot_sweep(recovered, title="Recovered logit difference by patched head")
save_figure(fig, f"{OUTPUT_DIR}/plots/recovered_heads.png")
fig

Activation patching: which components carry the signal? figure 1

ranked = sorted(recovered.scores.items(), key=lambda item: -abs(item[1]))

print(f"{'head':<10} {'recovered':>12}")
print("-" * 23)
for node, score in ranked[:8]:
    print(f"{f'L{node.layer}H{node.head}':<10} {score:>+12.3f}")

head          recovered
-----------------------
L10H7            -0.460
L5H5             +0.268
L8H10            +0.263
L8H6             +0.240
L7H9             +0.224
L11H10           -0.222
L7H3             +0.155
L3H0             +0.129

Reading the map

The heads that light up here are not the same as the ones direct logit attribution ranked highest, and that difference is the point.

Heads in layers 7 and 8 (such as L8H6, L8H10, L7H9) recover a lot of behavior even though they contribute little directly to the logits. They matter because of what they feed downstream.
L10H7 recovers a strongly negative fraction: patching it in from the clean run makes the corrupt run worse, consistent with the negative name mover found by attribution and ablation.
Layer 5 contains a duplicate-token head that notices the repeated name.

So attribution answers “who wrote the answer” and patching answers “who carried the information”. circuit_discovery.ipynb connects the two.

5. Save

run_dir = Save(
    output_dir="murano_outputs",
    run_name="activation_patching",
    model_id=model.model_id,
)(swept)[keys.OUTPUT_DIR]

print("saved to:", run_dir)
print("the sweep itself:", sorted(p.name for p in (run_dir / "sweep").iterdir()))

saved to: murano_outputs/activation_patching
the sweep itself: ['sweep.json']

What next

ablation.ipynb runs the reverse intervention and asks about necessity rather than sufficiency.
circuit_discovery.ipynb connects the heads found here to the ones attribution found.