Activation Patching › Causal Tracing

The broadest use of activation patching is causal tracing, which systematically sweeps all positions and layers to produce a complete map of where information flows through the network. Rather than patching one site at a time and moving on, causal tracing treats the patching experiment as a measurement tool and applies it exhaustively across the model's spatial and temporal dimensions.

In causal tracing, you run the patching experiment for every combination of layer $l$ and token position $p$. The result is a two-dimensional matrix of patch effects, with token positions on one axis and layers on the other. Cells with high effect values indicate that the activation at that (layer, position) coordinate lies on the causal path for the behavior.

This matrix is typically visualized as a heatmap, and the resulting picture is often remarkably structured. When Meng et al. ran causal tracing on GPT-style models for factual associations, they found a striking pattern: the strongest effects were concentrated in the early-to-middle MLP layers at the token position corresponding to the last token of the subject entity. When the model is completing "The Eiffel Tower is located in", the critical information appears to flow through MLP layers processing the "Tower" token position.

This finding suggested that factual associations are stored in MLP weights as key-value pairs, with the subject token's accumulated representation acting as the lookup key. The MLP receives the subject's representation, matches it against patterns learned during pretraining, and retrieves the associated fact. This interpretation was further supported by the ROME paper's direct manipulation of MLP weights to update factual associations. We will revisit it in the Mechanistic Interpretability chapter when we examine circuits in detail.

What makes causal tracing so informative is that it identifies which layers matter and reveals the structure of information flow. The heatmap shows not one bright spot but typically a characteristic shape: an early region where the fact is first "found" in MLP layers at the subject position, and a late region where information arrives at the last token position just before the output is generated. Understanding this shape is understanding the factual retrieval circuit at a high level.

A common choice for creating the corrupted input in causal tracing is adding Gaussian noise to the embeddings of the subject tokens, rather than substituting a different subject entity entirely. This approach has a practical advantage: it corrupts the model's ability to retrieve the fact without changing the prompt's syntactic structure. If you use a different subject (France vs. Germany), the model might answer differently for reasons unrelated to the factual lookup. Different subjects have different tokenization lengths, different frequencies in training data, and different syntactic contexts. A substantially different subject might create confounds that make the causal trace harder to interpret.

Adding noise directly to token embeddings produces a controlled corruption. The subject tokens are still present in the input as tokens, but their embedding representations are scrambled, preventing the model from recognizing the entity. This isolates the effect of "knowing" the subject from any structural differences between two different prompts. The rest of the prompt remains identical, so the syntactic context and token positions are preserved.

Let $e_i$ be the embedding of subject token $i$ before adding positional information. The corrupted embedding is:

$$\hat{e}_i = e_i + \epsilon_i, \quad \epsilon_i \sim \mathcal{N}(0, \sigma^2 I)$$

where:

• $e_i$: the original embedding vector for token $i$
• $\epsilon_i$: Gaussian noise sampled independently for each subject token
• $\sigma$: noise magnitude, typically set to 3 times the empirical standard deviation of the embedding vectors

The noise magnitude $\sigma$ requires calibration. Too small, and the model can still recognize the entity, giving a weak baseline disruption. Too large, and the corrupted input becomes so far out of distribution that the model's behavior is unpredictable and the causal trace becomes noisy. The choice $\sigma = 3\sigma_e$, where $\sigma_e$ is the empirical standard deviation of the embedding distribution across a sample of real tokens, provides a reliable disruption while remaining in a regime where the model continues to process the input coherently.

An important subtlety: when you patch the clean activation at a particular layer and position, you are restoring it into a context where the surrounding activations came from a noise-corrupted input. The attention mechanism at later layers will attend across positions, mixing the patched (clean) activation at one position with corrupted activations elsewhere. This is not a flaw but a feature: the fact that a single clean activation can still recover the behavior despite the mismatched context argues strongly for that activation's causal role.

A key insight from causal tracing is that information doesn't just sit statically in one place. Once a critical activation is established at some layer and position, it propagates forward through later layers. Downstream attention heads read from earlier residual streams and copy information to later positions. This means that in the causal trace matrix, you typically see both "early sites" where information first appears and "late sites" where it arrives at the final prediction position.

The ROME paper introduced explicit terminology for this distinction. Early sites are the first layers where the patch has a large effect, typically middle MLP layers at the subject position. These are the locations where the model first looks up or constructs the factual information. Late sites are the final layers before the prediction, typically at the last token position where the output is generated. Information flows from early sites through attention layers to late sites, where it is used to generate the next token.

Understanding this propagation path is key to understanding the full circuit. The early MLP sites are the "storage" locations: they encode factual associations in their weights and retrieve them given the right activations. The late attention sites are the "routing" mechanism: they read from earlier positions and bring the relevant information to the final prediction position. A complete mechanistic account of factual recall needs both halves of this story.

This propagation structure also explains why residual stream patching at a late layer and position often shows high effects even if the computation happened much earlier. The information has been propagated forward by attention. If you patch the residual stream at layer 11, token position 7, you are providing the network with everything it had during the clean run right before the output, so of course the output recovers. This is a high effect but an uninformative one: it is too close to the output to tell you where the computation happened. This is why causal tracing typically focuses on the earliest high-effect site as the true computational source, rather than the highest-effect site overall.