Kiru Lab  /  Mechanistic Interpretability  /  Causal Methods

Ablation and Lesion Studies

Remove a component and measure the damage. The oldest method in neuroscience, applied to weights.

Hands-On Lab  ·  about 40 minutes

Ablation asks the necessity question: if this component is removed, does the behavior fail? The technique is borrowed directly from lesion studies in neuroscience, and it inherits both their power and their interpretive traps.

The baseline you choose is a hypothesis

  • Zero ablation sets the output to zero. Simple, but zero is not a neutral value — it is far outside the normal activation distribution, so you may be measuring distribution shock rather than lost function.
  • Mean ablation substitutes the average activation over a dataset. Stays in distribution, removes the input-specific signal, and is usually the better default.
  • Resample ablation substitutes the activation from a different input. Preserves the marginal distribution and is the most conservative option.
  • Path patching ablates a specific connection between two components rather than a component entirely, which is what you want when the question is about composition.
Mean ablation of a single head
def ablate_head(model, layer, head, mean_activation):
    block = model.blocks[layer]
    d_head = model.config.d_head
    lo, hi = head * d_head, (head + 1) * d_head

    def hook(module, inputs, output):
        output[:, :, lo:hi] = mean_activation[lo:hi]
        return output

    return block.attn.register_forward_hook(hook)   # caller removes the handle
The lesion fallacy, imported wholesale

Neuroscience learned this the hard way: damage to a region degrading a function does not mean the region implements that function. It might be relaying, gating, supplying energy, or simply adjacent to what actually matters. The same caution applies exactly to ablating a head. "Necessary for" is a much weaker claim than "implements", and conflating them is the most common error in interpretability writeups.

Report ablation results as effect sizes across a distribution of inputs and seeds, not as single dramatic examples. A head whose ablation costs 0.3 nats on one prompt and 0.0 on forty others has told you about that prompt.

Hold on to

  • Mean or resample ablation avoids the artifacts of zero ablation
  • "Necessary for" does not mean "implements"
  • Report distributions of effect, not single striking examples

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Ablate every attention head one at a time and rank them by effect on a task. Compare zero versus mean ablation rankings.
    Hint

    Compute the mean activation over a dataset first, then substitute it head by head.

    Solution

    The two rankings usually agree at the very top and diverge in the middle. Zero ablation inflates the apparent importance of heads whose typical output is far from zero, because you are measuring distribution shock as well as lost function — the model has been pushed somewhere it never goes. Mean ablation keeps the activation in-distribution and isolates the input-specific contribution, which is the thing you actually wanted to measure.

    import torch

scores = {}
for layer in range(model.config.n_layers):
    for head in range(model.config.n_heads):
        for mode in ("zero", "mean"):
            handle = ablate_head(model, layer, head,
                                 torch.zeros_like(mean_act) if mode == "zero" else mean_act)
            with torch.no_grad():
                scores[(layer, head, mode)] = metric(model(ids))
            handle.remove()

for mode in ("zero", "mean"):
    ranked = sorted(((v, k) for k, v in scores.items() if k[2] == mode))
    print(mode, [k[:2] for _, k in ranked[:5]])
  2. Take your top head, ablate it, and check whether another head's behavior changes to compensate.
    Hint

    Record every other head's output with the target head intact, then again with it ablated.

    Solution

    On some tasks another head measurably increases its contribution when the primary is removed — backup behavior. Where it happens, the single-ablation effect understates the primary head's role, because the network repaired itself within one forward pass. Report both numbers: effect of ablating alone, and effect of ablating together with the backup. Reporting only the first is how a real circuit gets dismissed as unimportant.

    import torch

baseline = head_outputs(model, ids)                  # (layers, heads, ...)
handle = ablate_head(model, PRIMARY_LAYER, PRIMARY_HEAD, mean_act)
with torch.no_grad():
    after = head_outputs(model, ids)
handle.remove()

delta = (after - baseline).norm(dim=-1)
top = delta.flatten().topk(5)
print("heads that changed most:", top.indices.tolist(), top.values.tolist())

In the Bio Mirror

Lesions, Compensation, and the Limits of Ablation

Sign in to track your progress through the lab.

← Activation Patching Probing: What a Classifier on Activations Can and Cannot Tell You →