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.