Activation Patching and Causal Tracing

Activation patching (also called causal tracing) is an intervention-based interpretability method that surgically replaces the internal activations of a neural network at a specific layer and token position with activations from a counterfactual forward pass, measuring the causal effect of that localized representation on the model's output distribution. It matters because it transforms the study of neural networks from correlational observation to controlled experimentation, enabling researchers to attribute model behaviors to specific computational sites with mechanistic precision. Remarkably, causal tracing of GPT-2 XL and GPT-J revealed that factual associations such as "The Eiffel Tower is in [Paris]" are stored in a sparse set of mid-layer MLP modules at the subject token, not distributed uniformly across the network as previously assumed.

The core insight

The central insight is that understanding what a network computes is insufficient — you must understand where and when it computes it by breaking the causal graph. Activation patching treats every (layer, token-position) pair as a node in a structural causal model, and the patch operation is an atomic do-intervention: do(h_l^t = h_l^{t,*}), isolating that node's downstream contribution independent of its upstream causes. The "aha moment" is recognizing that a high indirect effect score means that specific computational site is not merely correlated with the output — it is a necessary mediator, a bottleneck in the causal graph through which the relevant information must pass. This reframes mechanistic interpretability from pattern-matching over weights to causal surgery over forward passes.

Common trapCommon Misconception

The most common mistake is conflating high activation magnitude with causal importance — practitioners often inspect which neurons or attention heads have the largest activations and conclude those are the "responsible" components. This is wrong because a unit can be highly active while being entirely redundant (its information is duplicated elsewhere in the residual stream), and a unit with modest activation can be a decisive bottleneck whose removal collapses model performance. The correct mental model is that causal patching measures the information-theoretic bottleneck structure of the computation graph: a site is important if and only if restoring its corrupted value to its clean counterpart recovers output probability, regardless of the absolute magnitude of that value.

How It Works

The Structural Causal Model Formulation

Let $\mathcal{M}$ be a transformer with $L$ layers processing a sequence of $T$ tokens. For each layer $l \in \{0,\ldots,L-1\}$ and token position $t \in \{0,\ldots,T-1\}$ , define $h_l^t \in \mathbb{R}^d$ as the residual stream state. A clean run on prompt $x$ produces activations $\{h_l^t\}$ ; a corrupted run on prompt $x^*$ (typically $x$ with the subject entity noised or replaced) produces $\{h_l^{t,*}\}$ .

\quad \mathrm{IE}(l,t) = \mathbb{P}_{\mathrm{patch}}[o \mid \mathrm{do}(h_l^t \leftarrow h_l^{t,\mathrm{clean}})] - \mathbb{P}_{\mathrm{corrupt}}[o]

Explanation: The indirect effect at site $(l,t)$ is the gain in output probability for the target token $o$ obtained by restoring exactly that one activation to its clean value, while the rest of the network continues from the corrupted run. Here $\mathbb{P}_{\mathrm{patch}}[o]$ is computed by running the corrupted forward pass, intercepting the residual stream at $(l,t)$ , substituting the clean value $h_l^{t,\mathrm{clean}}$ , and propagating forward from that point. Every symbol: $l$ is layer index, $t$ is token position, $o$ is the target output token, $\mathrm{do}(\cdot)$ denotes Pearl's intervention operator.

Total, Direct, and Indirect Effects

Using the Potential Outcomes Framework and Structural Causal Models, we can decompose the total effect $\mathrm{TE}$ into a direct effect $\mathrm{DE}$ (how the corrupted input changes the output while holding mediators fixed) and an indirect effect $\mathrm{IE}$ (how it changes the output through the mediators):

\quad \mathrm{TE}(x, x^*) = \mathbb{P}_{\mathrm{clean}}[o \mid x] - \mathbb{P}_{\mathrm{corrupt}}[o \mid x^*]

Explanation: The total effect is simply the difference in probability of predicting the correct answer between the clean and fully corrupted runs, serving as the normalization baseline for patching experiments.

\quad \mathrm{NIE}(l,t) = \frac{\mathrm{IE}(l,t)}{\mathrm{TE}(x,x^*)}

Explanation: The normalized indirect effect expresses what fraction of the total information loss is recovered by patching site $(l,t)$ , enabling fair comparison across prompts with different baseline probabilities.

MLP vs. Attention Decomposition

Within a single transformer layer, the residual stream update splits into an attention contribution and an MLP contribution:

\quad h_{l+1}^t = h_l^t + \underbrace{a_l^t}_{\text{attention out}} + \underbrace{m_l^t}_{\text{MLP out}}

Explanation: Because the residual stream is additive, each component's causal contribution can be isolated independently by patching $a_l^t$ or $m_l^t$ alone, decomposing NIE into attention-NIE and MLP-NIE without approximation. Patching can further be localized to individual attention heads $h = 1,\ldots,H$ by replacing only $a_l^{t,h}$ :

\quad \mathrm{NIE}_{\mathrm{head}}(l,h,t) = \frac{\mathbb{P}[o \mid \mathrm{do}(a_l^{t,h} \leftarrow a_l^{t,h,\mathrm{clean}})] - \mathbb{P}_{\mathrm{corrupt}}[o]}{\mathrm{TE}(x,x^*)}

Explanation: Head-level NIE identifies specific attention heads that are causally responsible for routing particular information to the final prediction, forming the basis for circuit-level analysis.

researcher · Kevin Meng | b. 1993 | Causal Tracing and ROME

Kevin Meng completed his PhD at MIT working with David Bau and is the primary author of "Locating and Editing Factual Associations in GPT" (Meng et al., NeurIPS 2022), which introduced the causal tracing methodology at scale and demonstrated that factual memories in GPT models are causally localized to mid-layer MLPs. His work directly established the empirical foundation for model editing methods such as ROME and MEMIT. Surprisingly, the causal tracing experiments were initially designed as a sanity check for the editing method, and the clean localization result was an unexpected finding that reshaped the interpretability research agenda.

researcher · David Bau | b. 1974 | Network Dissection and Causal Mediation in Neural Nets

David Bau is a faculty member at UMass Amherst whose earlier work on Network Dissection (Bau et al., CVPR 2017) established unit-level semantic interpretability in vision models; his subsequent collaboration on causal mediation analysis in language models provided the theoretical backbone for activation patching as a rigorous intervention methodology. He brought Pearl's causal hierarchy explicitly into the neural network interpretability toolkit. Bau spent years as a software engineer at Microsoft before returning to academia, which gives his work an unusually practical, tooling-oriented character.

Figure

Causal Tracing Forward Pass Diagram

Figure

NIE Heatmap Schematic