Causal tracing of MLP outputs (GPT-2 small)¶

Meng et al. (2022) corrupt a fact's subject tokens with noise, then restore one hidden state at a time from the clean run and measure how much of the correct answer's probability comes back. This entry measures the MLP-output grid on GPT-2 small with TorchLens and again with plain PyTorch hooks on the same facts and noise draws. The entry passes: the two implementations agree exactly (max abs difference 0 over all 72 cells), both wrong-site controls recover almost nothing, and the largest effect sits at the last subject token as the paper reports, but at layer 0 on GPT-2 small rather than in the paper's middle-layer band on GPT-2 XL.

The figure is drawn from the saved causal_tracing.json by the shared gallery renderer; the script's --figure-only flag redraws it without rerunning the model.

Oracles¶

The average total effect of the corruption (clean minus corrupted answer probability) is 0.132. The single-layer AIE grid, rounded to four places (rows are token categories, columns layers 0 to 11):

Single-layer restoration on GPT-2 small recovers far less than the paper's ten-layer windows on GPT-2 XL, and the subject-token effect concentrates in layer 0, the MLP that reads the corrupted embedding directly. The paper's middle-layer band does not appear at this size and window; --window 10 reproduces the paper's window rule for comparison.

Controls¶

Artifact and replay¶

The clean trace of fact 0 is saved as a portable causal_tracing_clean_fact0.tlspec (24.3 MB, sha256 08ce6135552cbb8e...). The script reloads it with tl.load, takes the layer-0 MLP value at the last subject token from the reloaded trace, reruns the restoration on a fresh corrupted trace, and gets the in-memory effect for that cell exactly (0.0013961, difference 0.0).

Basis-grid row¶

The counts are the load-bearing lines (no blanks, comments or docstrings) of the script's torchlens region (capture the clean and corrupted runs, read the MLP values, fork, attach the restoring edits, rerun) and its plain region (the same with forward hooks). Shared setup (tokenising, corrupting, the restoration writer, scoring) is the same for both. The TorchLens side is longer here: it keeps the traces for the artifact and the controls, which the plain-hooks side does not need. No count against the original reference repository was measured.

The stamp¶

{"commit": "e9122c7bc3767c6ee529f7578ca880bc8b483b36", "date_utc": "2026-11-03T20:15:32Z",
 "device": "cpu", "libraries": {"numpy": "2.2.6", "torchvision": "0.28.0+cpu",
 "transformers": "5.18.0"}, "machine": "x86_64",
 "platform": "Linux-5.15.0-139-generic-x86_64-with-glibc2.31", "python": "3.10.19",
 "torch": "2.13.0+cpu", "torch_threads": 2, "torchlens": "2.36.3", "wall_seconds": 871.2}

Vary it¶

The script takes --facts N (measure more of the 1,000 facts; the first 10 still decide the verdict) and --window W (restore W consecutive MLP layers at once, the paper's window of 10 on GPT-2 XL; the verdict always uses one layer). Both add a reported grid beside the verdict.

The examples below run the same restoration on a two-layer stand-in GPT-2 with random weights, so they finish in seconds: TorchLens forks the corrupted trace, attaches an edit to one MLP's down_out facet that writes the clean value at the last subject token, and reruns; a plain forward hook does the same and must give the same answer probability.

# recipe-test: tool=transformers expected=True
import torch
from torch import nn
from transformers import GPT2Config, GPT2LMHeadModel

import torchlens as tl

config = GPT2Config(n_layer=2, n_head=2, n_embd=16, vocab_size=32, n_positions=16, use_cache=False)
config._attn_implementation = "eager"
torch.manual_seed(0)
lm = GPT2LMHeadModel(config).eval()


class EmbedsGPT2(nn.Module):  # the corrupted run is an ordinary input: token embeddings
    def __init__(self, lm):
        super().__init__()
        self.lm = lm

    def forward(self, embeds):
        return self.lm(inputs_embeds=embeds).logits[:, -1]


wrapped = EmbedsGPT2(lm).eval()
ids = torch.tensor([[3, 7, 1, 9, 4]])
subject, target = [1, 2], 5  # the subject's positions and the answer token
with torch.no_grad():
    clean_embeds = lm.transformer.wte(ids)
noise = torch.randn(4, len(subject), 16, generator=torch.Generator().manual_seed(0))
corrupted = clean_embeds.repeat(4, 1, 1)
corrupted[:, subject] += 3.0 * noise

clean = tl.trace(wrapped, clean_embeds, save="all")
corrupt = tl.trace(wrapped, corrupted, save="all")
values = {
    layer: clean.modules[f"lm.transformer.h.{layer}.mlp"].facets["down_out"].value.detach()
    for layer in range(2)
}


def restorer(layer):
    def restore(out, *, ctx=None):
        patched = out.clone()
        patched[:, subject[-1]] = values[layer][0, subject[-1]]
        return patched

    return restore


def p_answer_torchlens(layers):
    restored = corrupt.fork("restore")
    hooks = [
        (tl.facet("down_out").in_module(f"lm.transformer.h.{layer}.mlp"), restorer(layer))
        for layer in layers
    ]
    restored.do(hooks, model=wrapped, inputs=corrupted)
    return float(restored.output_ops[-1].out.softmax(-1)[:, target].mean())


def p_answer_plain(layers):
    handles = [
        lm.transformer.h[layer].mlp.register_forward_hook(
            lambda module, args, out, layer=layer: restorer(layer)(out)
        )
        for layer in layers
    ]
    with torch.no_grad():
        p = float(wrapped(corrupted).softmax(-1)[:, target].mean())
    for handle in handles:
        handle.remove()
    return p


one_layer = {layer: (p_answer_torchlens([layer]), p_answer_plain([layer])) for layer in range(2)}
RESULT = all(abs(ours - plain) <= 1e-6 for ours, plain in one_layer.values())

A window restores several layers in one rerun; widening it to both layers of the stand-in is one more edit on the same fork, and the plain hooks agree again:

# recipe-test: tool=transformers expected=True
window = [0, 1]
RESULT = abs(p_answer_torchlens(window) - p_answer_plain(window)) <= 1e-6