Causal tracing is the experiment that put a physical address on a fact. Introduced by Meng and colleagues in the 2022 ROME paper — Locating and Editing Factual Associations in GPT — it asks a sharp question: when a language model completes “The Space Needle is located in the city of” with Seattle, where inside the network does the ‘Space Needle → Seattle’ association actually live? The method answers it by corrupting the subject, watching the prediction collapse, and then restoring one internal hidden state at a time to see which restoration brings the correct answer back. It is a specialized form of activation patching, and its headline finding — that factual recall concentrates in mid-layer MLPs at the last subject token — is what made principled model editing possible. This article builds it from first principles: the corruption, the three runs, the indirect-effect metric, the severed traces, a worked example, the link to editing, and the critiques.
What causal tracing is
Causal tracing is a causal-mediation analysis of a single factual prediction. The network is treated as a causal graph whose nodes are hidden states, and the goal is to find which nodes mediate the flow of information from the subject tokens to the predicted object. Unlike a correlational probe, which only shows a state contains subject information, tracing shows the prediction depends on that state.
It is a close cousin of general activation patching, specialized in three ways. The corruption is not a swapped minimal pair but additive Gaussian noise on the subject’s token embeddings; the metric is the probability of the correct object token, not a logit difference; and the sweep restores clean states into the corrupted run, so it is a denoising experiment throughout. Together these turn a single fact into a heatmap over (layer, token) pinpointing where it is stored.
The corruption: noising the subject
The setup begins with a prompt whose subject tokens are identified — for “The Space Needle is located in…” the subject is The Space Needle. Instead of replacing words, causal tracing adds noise directly to the subject’s embedding vectors at the input:
for each subject-token position i:
e_i' = e_i + ε_i , ε_i ~ N(0, ν² I)
ν = 3 × (std. dev. of embeddings over a token sample)The scale ν is set a few standard deviations wide so the model effectively ‘forgets’ which entity the subject is and the probability of the correct object drops toward chance. Noising the embeddings rather than swapping tokens keeps sequence length and every position fixed, which makes a later clean-state restoration well-defined: every corrupted hidden state has an exact clean counterpart of identical shape at the same coordinates.
Three runs: clean, corrupted, restored
Every trace is built from three forward passes over the same prompt:
1. CLEAN run normally → p_clean = P(o) [high]
2. CORRUPTED noise the subject → p_corr = P(o) [low]
3. RESTORED noise the subject,
THEN copy clean h^(l)_t back → p_rest(l,t) = P(o | do) [?]The first two are baselines: the clean run shows the model knows the fact, the corrupted run shows the noise destroyed it. The third is the intervention. During an otherwise corrupted pass, you overwrite exactly one hidden state h^(l)_t — the residual-stream vector at layer l, token t — with its clean value, and let the rest of the computation proceed on the corrupted activations. If that single clean vector pulls the correct object back to the top, the fact’s recall route passes through (l, t).
The metric: total effect and indirect effect
Two scalars summarize a trace. The total effect is how much the corruption cost overall; the indirect effect of a state is how much of that damage restoring it repairs:
TE = p_clean − p_corr
IE(l, t) = p_rest(l, t) − p_corr
restored fraction = IE(l, t) / TEIE is a probability gap: the extra mass the correct object gains over the corrupted baseline purely because one state was returned to its clean value. A large IE at (l, t) is a direct causal claim — recall of this fact is mediated by that state. Sweeping every layer and token yields the causal trace heatmap, an L × T grid whose bright band shows where the fact lives and, read across layers, how the information marches from the subject tokens toward the final read-out position.
Averaging over many facts
A single prompt is one noisy measurement, and the Gaussian corruption adds literal randomness, so no serious claim rests on one trace. The ROME analysis averages the indirect effect over a large dataset of known facts and over multiple noise samples per fact, giving the average indirect effect:
AIE(l, t*) = mean over facts, noise seeds of IE(l, t)Here t* is a token role — ‘last subject token’, ‘last token’ — not an absolute index, since prompts differ in length. Aggregating turns the per-fact heatmaps into one population-level map of the network’s factual machinery, smoothing out the noise variance so the surviving signal reflects the architecture rather than a lucky seed.
The finding: an early site and a late site
The averaged trace of a decoder-only transformer has a striking two-peak shape. There is a strong early site: a band of high indirect effect in the mid-layer MLPs at the last token of the subject. And there is a weaker late site: high indirect effect in the upper layers at the final token, just before the prediction is read out.
The natural reading is a two-stage mechanism. At the last subject token, mid-layer feed-forward blocks act as an associative key–value store: they recognize the assembled subject and write the associated attributes into the residual stream. Later, attention at the final position retrieves that enriched representation and moves it to where the next-token distribution is computed. The early site is the fact being recalled; the late site is it being used.
Severed traces: is it the MLP or the attention?
Restoring the whole residual stream at (l, t) is coarse: it repairs whatever passed through without saying which submodule did the work. To decompose it, causal tracing runs severed variants that restore only one contribution while forcing the other to stay corrupted:
| Trace variant | What is restored | What it isolates |
|---|---|---|
| Full | whole residual state h^(l)_t | total mediation at this site |
| MLP-severed | state, but MLPs at t recompute on corrupt input | effect not due to those MLPs |
| Attn-severed | state, but attention recomputes on corrupt input | effect not due to attention |
The comparison is diagnostic. At the early site, cutting the MLPs’ contribution collapses the indirect effect — the restoration only helps because the mid-layer MLPs recompute with clean input — while cutting attention barely changes it. That is direct evidence that the MLP modules, not attention, store the factual association at the last subject token.
A worked numeric example
Take the Space Needle prompt on a mid-size model. The clean run assigns p_clean = P(“Seattle”) = 0.91. Noising the subject embeddings collapses it to p_corr = 0.04, so the total effect is TE = 0.87. Now restore single clean states and re-read the probability:
restore h at (layer 6, last-subject-tok): p_rest = 0.62
IE = 0.62 − 0.04 = 0.58 → 0.58/0.87 ≈ 67% restored
restore h at (layer 2, last-subject-tok): p_rest = 0.07
IE = 0.03 → ~3% restored (info not yet assembled)
restore h at (layer 6, first-subject-tok): p_rest = 0.09
IE = 0.05 → ~6% restored (wrong position)One vector at layer 6 on the last subject token recovers two-thirds of the lost probability; the same layer one token earlier does almost nothing, and an early layer nothing at all. That contrast, sharp in both layer and position, is the localization signal, and why the heatmap has a compact hot spot rather than a diffuse glow.
From localization to editing: ROME
Causal tracing was never an end in itself — it was reconnaissance for editing. If a fact is written by the mid-layer MLP at the last subject token, changing it should be possible by a surgical change to that one weight matrix. ROME treats the MLP’s second (down-projection) layer as a linear associative memory — a store of key→value pairs — and applies a closed-form rank-one update that inserts a new key→value binding for the subject while minimally disturbing the thousands of others.
The key is the subject’s representation at that layer; the value is optimized so the model predicts the new object. Because the edit is rank-one and aimed at the traced site, it can make “The Space Needle is in…” output Paris without retraining and without shredding unrelated knowledge. MEMIT extends the idea to thousands of edits across several layers. The heatmap says where; the rank-one update acts exactly there.
The critique: localization is not editability
The tidy story — trace finds the site, edit succeeds there — has been productively challenged. Hase and colleagues showed that where causal tracing localizes a fact and where you can best edit it are not the same place: you can often inject a new fact by updating layers the trace assigns little indirect effect, and the correlation between an fact’s trace peak and its best edit layer is weak.
The lesson is careful, not a dismissal. Causal tracing faithfully answers the question it asks — which states, when restored, recover the prediction under this noise corruption — but that is about information flow during recall, not about which weights most efficiently store the fact for editing. The two diverge because a transformer is redundant and distributed. Read causal tracing as strong evidence about mediation, and treat ‘therefore edit here’ as a hypothesis to test, not a corollary. Two caveats round it out: the map depends on the noise scale ν, and the random corruption demands averaging over many facts and seeds.
Running it on a CPU-class small model
Causal tracing is gradient-free and cheap — a genuinely practical laptop experiment. Each trace is three forward passes plus one per restored (l, t) cell, roughly 2 + L × T passes per fact, no backward pass, memory bounded by one cached clean-activation set. On a sub-1B SLM a 15-token forward pass is tens of milliseconds on CPU, so a single-fact sweep finishes in seconds and an averaged map over a few hundred facts is a minutes-scale job.
The recipe: pick a small model that knows some facts (GPT-2 XL is the ROME reference, but smaller models work for the mechanics); assemble subject–relation–object prompts it completes correctly; identify subject-token spans; add N(0, ν²) noise to their embeddings; and loop a forward hook that restores clean residual states cell by cell, recording P(object). TransformerLens and nnsight expose named activations and hook APIs that make the whole sweep — severed variants included — a short script.
IE(l,t) = p_rest(l,t) − p_corr, normalized by the total effect p_clean − p_corr and averaged over many facts and noise seeds. The averaged map reveals a two-stage mechanism — an early site where mid-layer MLPs at the last subject token recall the fact, and a late site where attention at the final position uses it — and severed traces pin the recall on the MLPs specifically. That localization is what made ROME’s rank-one editing possible. But respect the caveats: the map depends on the noise scale, needs averaging, and — as the localization-versus-editability critique shows — where a fact is mediated is not always where it is best edited. The whole sweep is gradient-free and runs in seconds on a CPU with a small model.