Ablation: deleting a component, and why the delete matters

Note

Generated from notebooks/applications/ablation.ipynb, which you can run yourself. The outputs below are the ones it produced.

To test whether a component matters, remove it and see what breaks. The catch is that “remove” is not one operation. You can zero the component, replace it with its average, or resample it from a different prompt, and these can disagree completely about whether the same head is important.

This notebook runs all three on the same head and shows them disagreeing.

Key questions

• What does it mean to delete a component from a network that never sees zeros?
• Do zero, mean and resample ablation agree about which head matters?
• How does ablation differ from activation patching?

Structure

1. Set up
2. Establish the baseline
3. Compare zero, mean and resample
4. Ablate whole components, not just heads
5. Ablation versus patching
6. 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. No extras: pip install murano-interp

1. Set up

import torch

from murano import MuranoModel, Pipeline, keys
from murano.nodes import Node, NodeSet
from murano.steps import (
    Ablate,
    KLDivergenceStep,
    LoadPrompts,
    LogitDiffStep,
    Logits,
    Save,
    Sweep,
)

OUTPUT_DIR = "murano_outputs/ablation"

# 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 baseline

Every number below is read against the model’s undamaged logit difference.

HEAD = Node(9, "self_attn", head=9)

baseline = Pipeline(
    [
        LoadPrompts(task.clean),
        Logits(model),
        LogitDiffStep(task.correct, task.incorrect, model=model),
    ]
).run()
clean_score = baseline[keys.LOGIT_DIFF].value
print(f"clean logit difference: {clean_score:+.4f}")
print(f"target component      : {HEAD}")

clean logit difference: +2.3971
target component      : L9.self_attn.h9

3. Compare zero, mean and resample

Ablate supports three replacements, and the choice is a modelling decision:

• zero writes zeros. Simple, and badly off-distribution: no real forward pass ever produces a zero head output, so the model is being asked about an input it has never seen.
• mean writes the component’s average over the batch. On-distribution, but it erases what varies across prompts, which is often exactly the signal.
• resample writes the component’s value from a different prompt. Here we resample from the corrupt prompts, which ask the model to name the other person.

We score each two ways: the logit difference (did the behavior break?) and the KL divergence from the clean output distribution (how much did anything change?).

methods = {
    "zero": {"method": "zero"},
    "mean (over the batch)": {"method": "mean"},
    "resample (from corrupt)": {"method": "resample", "source": task.corrupt},
    "resample (shuffle batch)": {"method": "resample", "seed": 0},
}

# One Sweep, two harvested columns: the swept chain runs once per method and each
# ablated forward pass is scored both ways.
damage = Pipeline(
    [
        Sweep(
            over=list(methods),
            steps=lambda name: [
                Ablate(model, targets=[HEAD], **methods[name]),
                LogitDiffStep(
                    task.correct,
                    task.incorrect,
                    logits_key=keys.ABLATED_LOGITS,
                    model=model,
                ),
                KLDivergenceStep(p_key=keys.FINAL_LOGITS, q_key=keys.ABLATED_LOGITS),
            ],
            read=[keys.LOGIT_DIFF, keys.KL_DIV],
        )
    ]
).run(baseline.copy())

swept = damage[keys.SWEEP]
divergence = swept.column(keys.KL_DIV)

print(f"{'method':<26} {'logit diff':>11} {'change':>9} {'KL':>8}")
print("-" * 56)
for name, score in swept.scores.items():
    print(
        f"{name:<26} {score:>+11.4f} {score - clean_score:>+9.4f} "
        f"{divergence[name]:>8.4f}"
    )

method                      logit diff    change       KL
--------------------------------------------------------
zero                           +2.3897   -0.0073   0.0284
mean (over the batch)          +2.3582   -0.0389   0.0341
resample (from corrupt)        +1.6128   -0.7843   0.0313
resample (shuffle batch)       +2.2947   -0.1024   0.0854

The three methods do not agree, and the disagreement is the point.

Zero ablation says L9H9 does nothing (the logit difference barely moves), even though attention.ipynb showed its attention pattern is a textbook name mover and its OV circuit copies names 8 times out of 8.

Resampling from the corrupt prompts says L9H9 matters a great deal. Replacing its output with what it computes on a prompt whose answer is the other name costs most of the behavior.

Why the gap? Zeroing takes the head off the data manifold, and the rest of the network partly compensates for a signal that is merely missing. Resampling keeps the head’s output plausible while making it carry the wrong information, which is the counterfactual the task is actually about.

The lesson generalizes: an ablation encodes a hypothesis about what the component is for. Zero asks “what if this head were absent?” Resample asks “what if this head believed a different prompt?” Only the second is a question about the circuit.

Note also that the KL column stays small in every row. The output distribution barely moves in aggregate, while the answer flips. A distribution-level metric can miss a task-level failure entirely.

4. Ablate whole components, not just heads

targets accepts any component address, or a NodeSet of several. The same step deletes a single head, an entire attention block, or an MLP.

targets = {
    "one head": [HEAD],
    "two heads (NodeSet)": NodeSet(
        [Node(9, "self_attn", head=9), Node(10, "self_attn", head=0)]
    ),
    "the whole L9 attention": [Node(9, "self_attn")],
    "the L8 MLP": [Node(8, "mlp")],
}

scope = Pipeline(
    [
        Sweep(
            over=list(targets),
            steps=lambda name: [
                Ablate(
                    model,
                    targets=targets[name],
                    method="resample",
                    source=task.corrupt,
                ),
                LogitDiffStep(
                    task.correct,
                    task.incorrect,
                    logits_key=keys.ABLATED_LOGITS,
                    model=model,
                ),
            ],
            read=keys.LOGIT_DIFF,
        )
    ]
).run(baseline.copy())

print(f"{'target':<24} {'logit diff':>11} {'change':>9}")
print("-" * 46)
for name, score in scope[keys.SWEEP].scores.items():
    print(f"{name:<24} {score:>+11.4f} {score - clean_score:>+9.4f}")

target                    logit diff    change
----------------------------------------------
one head                     +1.6128   -0.7843
two heads (NodeSet)          +1.1563   -1.2408
the whole L9 attention       +1.2007   -1.1963
the L8 MLP                   +2.7033   +0.3062

5. Ablation versus patching

The two interventions are mirror images, and they answer different questions.

	base run	what is written in	question answered
Ablate	clean	a replacement (zero, mean, or another prompt)	is this component necessary?
Patch	corrupt	the clean activation	is this component sufficient?

Necessity and sufficiency are not the same property, and a redundant circuit can be sufficient without any single part being necessary. That is exactly why zeroing L9H9 looked harmless: other heads cover for it.

activation_patching.ipynb runs the other direction.

6. Save

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

print("saved to:", run_dir)

saved to: murano_outputs/ablation

What next

• activation_patching.ipynb runs the reverse intervention: restore a component instead of deleting it.
• metrics.ipynb covers KL divergence and the other ways to score the damage.