Activation Patching and Causal Tracing: Mathematical Foundations, Counterfactual Interventions, and Circuit Discovery in Large Language Models

Mechanistic interpretability aims to reverse-engineer the computational graphs implemented by trained neural networks into human-understandable algorithms. In autoregressive Transformer models, standard observational tools such as attention map visualization and linear probing often fail to distinguish between active computational paths and passive correlations. Attention weights indicate where a head routes information, but they do not prove that altering those weights changes model behavior. Similarly, linear probes demonstrate that a representation linearly encodes a feature, but they cannot establish whether downstream layers actually read or depend on that feature.

Activation patching, also formalized as causal mediation analysis, interchange intervention, causal tracing, and representation denoising, resolves this limitation by substituting observational correlation with counterfactual causal intervention. By running paired forward passes on carefully constructed prompt pairs and transplanting specific intermediate activation tensors between runs, activation patching measures the precise causal effect of individual layers, attention heads, MLP blocks, and token positions on the model output.

1. Mathematical Framework of Counterfactual Neural Interventions

Activation patching adapts Judea Pearl's Structural Causal Models (SCM) to the deterministic computation graph of a Transformer.

Let $M$ denote an autoregressive language model parameterized by weights $\theta$. The model processes a sequence of input tokens $x = (x_1, x_2, \dots, x_T)$ across $L$ layers. The residual stream activation at layer $l \in \{0, \dots, L\}$ and token position $p \in \{1, \dots, T\}$ is denoted as $x_l(p) \in \mathbb{R}^d$. The forward computation proceeds through multi-head attention (MHA) and multi-layer perceptron (MLP) blocks:

$$x_l(p) = x_{l-1}(p) + \sum_{h=1}^H z_{l,h}(p) W_{l,h}^O + m_l(p)$$

where $z_{l,h}(p)$ is the output vector of attention head $h$ in layer $l$, $W_{l,h}^O$ is the corresponding output projection matrix, and $m_l(p)$ is the output of the MLP block at layer $l$.

Clean Input:     x_clean   ──► [ Layer 1 ] ──► a_clean ──► [ Layer 2 ] ──► Logit(clean)
                                                   │ (transplant)
                                                   ▼
Corrupted Input: x_corrupt ──► [ Layer 1 ] ──► a_patch ──► [ Layer 2 ] ──► Logit(patched)

To isolate the causal role of an intermediate activation $a \in \mathcal{A}$ (such as a residual stream state $x_l(p)$, an attention head output $z_{l,h}(p)$, or an MLP output $m_l(p)$), we define two inputs:

1. Clean Input ($x_{\text{clean}}$): A prompt eliciting target behavior $y_{\text{clean}}$ (e.g., "When Mary and John went to the store, John gave a drink to Mary" $\rightarrow$ "Mary").
2. Corrupted Input ($x_{\text{corrupt}}$): A minimally altered counterfactual prompt where key information is modified or obscured (e.g., "When Mary and John went to the store, Mary gave a drink to John" $\rightarrow$ "John").

Let $a^{\text{clean}} = a(x_{\text{clean}})$ and $a^{\text{corrupt}} = a(x_{\text{corrupt}})$ represent the activations at node $a$ during the clean and corrupted forward passes, respectively.

An interchange intervention $do(a \leftarrow a^{\text{clean}})$ overrides the activation value at node $a$ during the forward pass on $x_{\text{corrupt}}$ with the pre-computed value $a^{\text{clean}}$, while computing all upstream nodes from $x_{\text{corrupt}}$ and all downstream nodes through standard forward propagation.

2. Denoising vs. Corrupting: Sufficiency vs. Necessity

Activation patching operates in two complementary directions, each testing a distinct causal property.

Clean-into-Corrupted Patching (Denoising / Sufficiency)

In denoising, the model runs primarily on the corrupted input $x_{\text{corrupt}}$, and clean activations are patched into candidate nodes:

$$M(x_{\text{corrupt}} \mid do(a \leftarrow a^{\text{clean}}))$$

This intervention tests whether the activation $a^{\text{clean}}$ is sufficient to restore the correct output logit difference in the presence of corrupted context.

The standard performance metric is the clean-to-corrupted logit difference:

$$\Delta \mathcal{L}(x) = \text{logit}(y_{\text{clean}}) - \text{logit}(y_{\text{corrupt}})$$

The normalized indirect effect (denoising metric) is defined as:

$$\text{Effect}_{\text{denoise}}(a) = \frac{\Delta \mathcal{L}(M(x_{\text{corrupt}} \mid do(a \leftarrow a^{\text{clean}}))) - \Delta \mathcal{L}(M(x_{\text{corrupt}}))}{\Delta \mathcal{L}(M(x_{\text{clean}})) - \Delta \mathcal{L}(M(x_{\text{corrupt}}))}$$

A score of $1.0$ indicates that patching node $a$ alone completely recovers the clean model decision.

Corrupted-into-Clean Patching (Noising / Necessity)

In noising, the base run is the clean input $x_{\text{clean}}$, and a corrupted activation (or Gaussian noise baseline) is patched into node $a$:

$$M(x_{\text{clean}} \mid do(a \leftarrow a^{\text{corrupt}}))$$

This intervention tests whether node $a$ is necessary for the clean behavior. The normalized corruption effect measures how much the clean logit difference degrades:

$$\text{Effect}_{\text{noise}}(a) = \frac{\Delta \mathcal{L}(M(x_{\text{clean}})) - \Delta \mathcal{L}(M(x_{\text{clean}} \mid do(a \leftarrow a^{\text{corrupt}})))}{\Delta \mathcal{L}(M(x_{\text{clean}})) - \Delta \mathcal{L}(M(x_{\text{corrupt}}))}$$

Causal Tracing via Gaussian Noise Corruption

Meng et al. (2022) formalized causal tracing for factual recall (e.g., "The Eiffel Tower is located in [Paris]") by adding independent Gaussian noise $\epsilon \sim \mathcal{N}(0, \sigma^2 I)$ directly to the input token embeddings of the subject entity ("The Eiffel Tower"), completely corrupting the model's factual association. They then sequentially restored individual layer-position states $x_l(p)$ from the uncorrupted run:

$$\text{Average Indirect Effect (AIE)} = \mathbb{E}_{x} \left[ \mathbb{P}\left(y_{\text{clean}} \mid x_{\text{noisy}}, do(x_l(p) \leftarrow x_l^{\text{clean}}(p))\right) - \mathbb{P}\left(y_{\text{clean}} \mid x_{\text{noisy}}\right) \right]$$

This formulation demonstrated that factual retrieval in autoregressive models concentrates in two localized regions: early-to-mid layer MLP modules at the final subject token position (fact retrieval), and late-layer attention heads at the final prompt token position (routing to the output vocabulary).

3. Granularity Levels: From Residual Streams to Circuit Graphs

Activation patching is applied hierarchically to isolate computational subgraphs, moving from coarse localization to fine-grained path interaction.

Hierarchical Localization Workflow:
1. Residual Stream Patching  ──► Pinpoint critical (layer, token position) coordinates
2. Component-Level Patching  ──► Disentangle Attention Heads vs. MLP Blocks
3. Path Patching             ──► Map direct and indirect interaction edges
4. Circuit Verification      ──► Validate minimal sufficient computational subgraph

Level 1: Residual Stream Patching

The residual stream acts as the shared communication bus of the Transformer. Patching entire residual stream vectors $x_l(p)$ across every layer $l \in [0, L]$ and token index $p \in [1, T]$ produces a 2D heatmap showing when and where information is written to the stream.

Level 2: Component-Level Patching

Once critical layers and token positions are localized, individual components are patched in isolation:

• Attention Head Output Patching: $do(z_{l,h}(p) \leftarrow z_{l,h}^{\text{clean}}(p))$ isolates specific attention heads.
• Attention Pattern Patching: $do(A_{l,h} \leftarrow A_{l,h}^{\text{clean}})$ freezes attention weights to test routing versus value vector transformation.
• MLP Activation Patching: $do(m_l(p) \leftarrow m_l^{\text{clean}}(p))$ isolates feed-forward memory lookup modules.

Level 3: Path Patching (Edge Interventions)

Component patching reveals nodes that hold critical information, but does not reveal which downstream components consume that information. Path patching freezes all paths from component $C_1$ to component $C_2$ except for a designated direct edge:

$$x_{l_2}(p) \leftarrow \text{Evaluate}(C_2 \mid \text{input from } C_1 = C_1^{\text{clean}}, \text{all other inputs} = C_{\text{corrupt}})$$

Path patching allows researchers to construct rigorous edge-level dependency graphs, proving that component $C_2$ reads its input directly from $C_1$ rather than through an indirect residual bypass.

4. Case Study: Circuit Discovery in Indirect Object Identification

The most celebrated application of activation patching is the reverse-engineering of the Indirect Object Identification (IOI) circuit in GPT-2 Small by Wang et al. (2022).

Given the template "When [A] and [B] went to the store, [B] gave a bag to [A]", GPT-2 Small reliably outputs [A]. Through systematic activation and path patching across all 144 attention heads, researchers isolated an exact 26-head computational circuit performing structured sub-tasks:

• Duplicate Token Heads (Layers 7-8): Detect repeated name tokens ([B]) and write positional signals into the residual stream.
• S-Inhibition Heads (Layers 7-9): Read duplicate token signals and suppress the attention paid by Name Mover Heads to the duplicated name [B].
• Name Mover Heads and Backup Heads (Layers 9-11): Attend back to the non-duplicated name [A] and directly write positive logit contributions for [A] into the unembedding projection.
• Negative Name Mover Heads (Layers 10-11): Partially dampen overconfident name-mover signals to calibrate output probabilities.

5. Scalability Bottlenecks and Attribution Patching (AtP)

Exhaustive activation patching requires evaluating every candidate node across all layers, token positions, and heads. For a model with $L$ layers, $P$ token positions, and $H$ attention heads, testing each component independently requires:

$$N_{\text{passes}} = \mathcal{O}(L \cdot P \cdot H)$$

forward passes per input example. For a 70-billion parameter model with 80 layers and 64 heads, evaluating a single 1,000-token prompt requires more than 5,000,000 forward passes, rendering exhaustive exact patching computationally intractable.

First-Order Taylor Approximation (Attribution Patching)

Attribution Patching (Nanda, 2023; Kramár et al., 2024) computes a first-order linear approximation of the activation patching effect using exactly two forward passes (clean and corrupted) and one backward pass on the corrupted run.

Let $\mathcal{M}(a)$ denote the metric function evaluated with intermediate activation $a$. Applying a first-order Taylor expansion around the corrupted activation $a^{\text{corrupt}}$:

$$\mathcal{M}(a^{\text{clean}}) - \mathcal{M}(a^{\text{corrupt}}) \approx \left( a^{\text{clean}} - a^{\text{corrupt}} \right)^\top \nabla_{a} \mathcal{M}(a^{\text{corrupt}})$$

where $\nabla_{a} \mathcal{M}(a^{\text{corrupt}})$ is the exact gradient of the metric with respect to activation $a$, obtained via standard backpropagation on the corrupted input.

By computing:

$$\text{AtP}(a) = \sum_{i} \left( a_i^{\text{clean}} - a_i^{\text{corrupt}} \right) \cdot \frac{\partial \mathcal{M}}{\partial a_i}$$

researchers evaluate the approximate causal effect of every node and edge in the network simultaneously. Empirical benchmarks show that Attribution Patching achieves over 90% rank correlation with exact activation patching while reducing compute requirements by three to four orders of magnitude.

6. Methodological Pitfalls and Pathologies

While activation patching is a cornerstone of mechanistic interpretability, improper experimental design introduces systematic artifacts:

Out-of-Distribution (OOD) and Off-Manifold Activations

Combining activations from two distinct forward passes can place downstream layers in intermediate states that never occur during standard inference. Non-linear operations (LayerNorm and GeLU) can exhibit pathological saturated behavior when receiving unnatural linear combinations of clean and corrupted residual components.

Distributed Redundancy and Backup Heads

When a primary functional head (such as a Name Mover Head) is ablated or corrupted, downstream "backup heads" often step in to restore the output signal. In such cases, corrupted-into-clean patching (necessity) may report a near-zero effect because the backup mechanism compensates, masking the component's critical role in unperturbed execution. Denoising patching (sufficiency) helps mitigate this blind spot by verifying whether the component alone can drive the behavior.

Baseline Selection Artifacts

The choice of corrupted baseline fundamentally alters the causal attribution:

• Zero / Mean Ablation: Replaces activations with zeros or dataset averages, frequently introducing severe OOD distribution shifts.
• Gaussian Noise Corruption: Completely destroys structure but can trigger abnormal out-of-domain attention distributions.
• Counterfactual Resample Ablation: Replaces specific entity tokens with semantically valid alternatives (e.g., swapping names), keeping the intermediate representations on-manifold and ensuring high-fidelity causal attribution.

Sources

• Vig et al. (2020) - Investigating Gender Bias in Language Models Using Causal Mediation Analysis
• Geiger et al. (2021) - Inducing Causal Structure for Interpretable Neural Networks
• Meng et al. (2022) - Locating and Editing Factual Associations in GPT
• Wang et al. (2022) - Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small
• Nanda (2023) - Attribution Patching: Activation Patching At Industrial Scale
• Conmy et al. (2023) - Towards Automated Circuit Discovery for Mechanistic Interpretability
• Heimersheim & Nanda (2024) - How to use and interpret activation patching
• Kramár et al. (2024) - AtP: An Efficient and Scalable Method for Localizing LLM Behaviour to Components