Replicating ROME's Causal Tracing

Recovering the causal-tracing sites from Meng et al. (2022) with a native activation-restoration primitive and no original-paper code, then testing whether the same structure appears in an 8B model

Topics	Language models · Circuits & causal tracing
Models	GPT-2 XL (48L) · Qwen3-8B-Base (36L, GQA)
References	Meng et al. 2022, "Locating and Editing Factual Associations in GPT" (ROME) · causalab causal_sufficiency

Try it yourself. Open the finished run in your own workspace and pick up where it leaves off.

The task and the corruption

Causal tracing asks where in a transformer a single fact is stored. It starts from a prompt whose subject the model completes correctly:

Base prompt. "The Space Needle is in downtown ___" → "Seattle"

The method then corrupts the fact and watches the model recover it. Gaussian noise (3σ of the embedding distribution) is added to the subject tokens' input embeddings, which collapses the prediction: P(Seattle) falls from near one to almost zero. With the run corrupted, the analysis restores one clean hidden state at a time — taking the value a single state held in the uncorrupted run and writing it back into the corrupted run, leaving everything else noised — and records how much of the lost probability returns. That recovered probability is the average indirect effect (AIE) of that state.

Sweeping the restore over every token position and every layer gives a position-by-layer map of where the fact lives. The trace is run three ways, restoring a different component each time:

Residual stream — a single hidden state per (position, layer).
MLP — a 10-layer window of MLP outputs, centered at each layer.
Attention — a 10-layer window of attention-module outputs.

ROME reports two sites. An early site at the last subject token in middle layers, which the paper attributes to mid-layer MLPs writing the fact, and a late site at the last token in upper layers, where attention moves the fact to the position that makes the prediction. This recipe reproduces that trace using only causalab's native causal_sufficiency primitive — 3σ subject-embedding noise, single-pass clean-state restoration, scored by P(first object token), no code from the ROME repository. In the heatmaps below, ember cells recover more of the corrupted prediction, cream cells recover almost none, and the ringed cell marks the site named in the text.

GPT-2 XL: the replication

Run against GPT-2 XL on 1,000 facts from CounterFact that the model knows (greedy generation names the object first), averaged over 3 noise seeds. Mean clean P(object) is 26.0% against the paper's 27.0%, so the model and the known set line up before any tracing begins. Two of the paper's three claims reproduce; the sharpest one does not.

Residual two-site structure — reproduces. The early site appears at the last subject token, peak AIE 6.9% at layer 13 (paper 8.7% at layer 15), alongside the late readout site at the last token.
Late attention site — reproduces. Restoring attention windows recovers the prediction most at the last token, peaking at layer 31 (paper centers layer 32).
Mid-layer MLP site — does not reproduce. The MLP-only trace at the last subject token decays monotonically from the embedding rather than bumping at the paper's mid layers, so this run finds no clean evidence for ROME's strongest mechanistic claim.

GPT-2 XL residual-stream causal trace, 6 token positions by 48 layers — Residual stream (paper Figure 2a). Each cell is the average recovery (AIE) from restoring one clean residual state, by ROME relative token position (rows) and layer (columns), over 1,000 facts × 3 seeds. The global maximum is the trivial readout band at the last token in the final layers; the ringed cell marks the non-trivial discovery, the early site at the last subject token (peak 6.9% at layer 13, paper 8.7% at layer 15).

GPT-2 XL MLP-window causal trace heatmap — MLP windows (paper Figure 2b). This is the one place the replication does not hold. The MLP-only trace at the last subject token peaks at layer 4 (ringed) and decays monotonically from the embedding, with no localized mid-layer bump like the paper's ~layer 17. A 10-layer MLP window centered low overlaps the corrupted embedding, so restoring it returns the corrupted information directly. The monotonic decay is the signature of that restore-near-corruption artifact, not of mid-layer MLP recall. Note the scale: peak recovery here is 3.4%, far below the residual and attention maps.

GPT-2 XL attention-window causal trace heatmap — Attention windows (paper Figure 2c). Recovery concentrates at the last token in mid-to-late layers, peaking at layer 31 (ringed; paper centers layer 32), consistent with attention copying the subject's information to the final position for prediction. Effect inside the subject span is comparatively weak.

One number sits outside tolerance and is worth stating plainly: the corrupted floor is 1.1% against the paper's 8.5%, roughly 8× too low, and the average total effect is 24.9% against 18.6%. The 3σ noise over-corrupts relative to the original protocol, which both inflates the total effect and worsens the restore-near-corruption artifact in the MLP map. The clean prediction still matches the paper, so the localization comparison stands; the magnitude comparison does not.

Qwen3-8B-Base: the same trace on a modern model

ROME reports numbers only for GPT-2 XL. The identical task, corruption, and three traces were run on Qwen3-8B-Base (36 layers, grouped-query attention) to test whether the structure transfers. The corruption is decisive here — 3σ noise drops P(object) from 0.422 to 0.015, an average total effect of 0.406, well above the threshold needed for the indirect effects to be interpretable. The two-site structure appears, with the same honest caveat on the MLP site.

Early site — generalizes. Residual restoration peaks at the last subject token at layer 8 (AIE 0.231), and MLP restoration peaks there at layer 5 (AIE 0.233): the same component-and-position structure as GPT-2 XL, shifted shallower in relative depth.
Late attention site — generalizes. Attention restoration peaks at the last token at layer 26 (AIE 0.220), the same late-model position GPT-2 XL shows.
The same confound applies. Qwen3's MLP peak sits at the shallow layers adjacent to the corruption, so it carries the same restore-near-corruption ambiguity as GPT-2's. The structure transfers; the mid-layer MLP-recall claim is no better resolved here than in the original.

Qwen3-8B-Base residual-stream causal trace, 6 positions by 36 layers — Residual stream. Read the same way as the GPT-2 maps. The ringed cell is the early site at the last subject token (layer 8, AIE 0.231); the strongest band remains the last-token readout in the final layers (global peak 0.406 at layer 35). The early-subject site is the non-trivial match to ROME's Figure 2a.

Qwen3-8B-Base MLP-window causal trace heatmap — MLP windows. Recovery concentrates on the last subject token at shallow-to-mid layers, peaking at layer 5 (ringed, AIE 0.233), the ROME Figure-2b position. As in GPT-2 XL, the peak sits near the corruption, so it carries the same restore-near-corruption ambiguity rather than serving as clean evidence for mid-layer MLP recall.

Qwen3-8B-Base attention-window causal trace heatmap — Attention windows. The dominant site is late: recovery concentrates on the last token at layers ~20–30, peaking at layer 26 (ringed, AIE 0.220). A weaker effect at the subject is visible but far smaller. Read with the MLP map, attention moves the recalled fact to the last position while the subject-side components supply it.

Try causal tracing on a task other than factual recall

Causal tracing is not specific to facts. The recipe corrupts the tokens that carry the answer, restores one clean state at a time, and reads off where the answer is reconstructed. Any task with a clean prompt, a localizable input span to corrupt, and a single-token target can be traced the same way. The site map that comes back is a hypothesis about where that behavior lives, and a different task should localize differently. A few candidates the same pipeline can take with only a new task package:

Subject–verb agreement. Trace prompts like "The keys on the table ___" where the verb must agree with the head noun across an intervening distractor, corrupting the subject noun and scoring are over is. The interesting question is whether the controlling state sits at the head noun or the distractor, and how late agreement is resolved.
In-context retrieval. Trace a small key–value prompt such as "Anna lives in Paris. Bob lives in Rome. Bob lives in ___", corrupting the queried name and scoring the stored value. This should localize to attention moving the bound value rather than to MLP recall, a direct contrast to the factual-recall map.
Two-step arithmetic. Trace "5 + 3 = 8, 8 + 2 = ___" by corrupting the operands and scoring the result token, to see whether the intermediate sum is reconstructed at a single position and depth or smeared across layers.

Replicating ROME's Causal Tracing

The task and the corruption

GPT-2 XL: the replication

Qwen3-8B-Base: the same trace on a modern model

Try causal tracing on a task other than factual recall

Related recipes