Causal Tracing

Causal tracing relies on three distinct forward passes to isolate factual recall:

• Clean Run: The model processes the original prompt, caching all hidden states.
• Corrupted Run: The subject entity is obfuscated (e.g., by adding noise to its embeddings), corrupting the model's ability to recall the fact.
• Corrupted-with-Restoration Run: The corrupted run is replayed, but a small window of clean hidden states from a specific layer and token position is patched back in. If the model's prediction recovers, the restored states are deemed causally important for the factual association.

The method pinpoints where knowledge is stored by systematically sweeping the restoration window across all layers and token positions. Key findings from the original ROME paper include:

• Early to Mid-Layer MLPs: Factual recall is consistently localized in the feed-forward networks of early to middle layers, not in late layers.
• Last Subject Token: The hidden states at the final token of the subject entity (e.g., 'Eiffel Tower') are the most causally relevant, acting as a key for retrieval.
• Attention is Secondary: Attention layers at the subject token primarily gather context, while the MLP layer performs the actual associative memory lookup.

The core metric for causal importance is the Average Indirect Effect. It quantifies how much a specific hidden state contributes to the correct answer:

• Calculation: AIE = P(correct answer | corrupted-with-restoration) - P(correct answer | corrupted run).
• Interpretation: A high AIE for a specific layer and token position indicates that the state is a necessary intermediate variable in the causal chain linking the subject to the object.
• Statistical Rigor: The effect is averaged over a large corpus of factual statements to ensure the localized site is a general mechanism for factual recall, not an artifact of a single prompt.

Causal tracing is a direct application of causal mediation analysis to neural networks. It treats the model's computation graph as a structural causal model:

• Mediator Variable: A specific hidden state acts as a mediator between the input (subject) and the output (object).
• Interchange Intervention: The restoration step is a formal interchange intervention, replacing the mediator's value under the corrupted treatment with its value under the clean treatment.
• Path-Specific Effect: By isolating a single state, the technique measures the causal effect along a specific computational path, distinguishing it from mere correlation measured by probing classifiers.

The localization provided by causal tracing is the foundational step for Rank-One Model Editing (ROME). Once a specific MLP layer is identified as the causal site for a fact, ROME treats that layer as a linear associative memory:

• Key-Value Store: The MLP's weight matrix is viewed as storing keys (subject encodings) and values (object properties).
• Rank-One Update: A new fact is inserted by applying a targeted rank-one update to the weight matrix, modifying the specific key-value pair without affecting other stored knowledge.
• Specificity: This demonstrates that causal tracing is not just diagnostic but prescriptive, enabling precise surgical modifications to a model's knowledge base.

Causal tracing provides a fundamentally different signal than linear probing or representation analysis:

• Correlation vs. Causation: A probe might show that factual information is present in a layer's representation, but causal tracing shows whether that information is actually used by the model to produce the output.
• Functional Necessity: A state with high probing accuracy but zero indirect effect contains decodable information that the model ignores. Causal tracing filters out this epiphenomenal knowledge.
• Computational Graph: It respects the model's actual forward-pass dependencies, providing a ground-truth map of the functional data flow rather than a static snapshot of representational content.

A direct causal intervention method that replaces a model's internal activation at a specific layer and token position with a cached activation from a different input. By measuring the resulting change in output, researchers can localize where a specific computation occurs. This is the foundational mechanism underlying causal tracing experiments, where clean activations are patched into a corrupted forward pass to restore a fact.

A statistical framework for quantifying how much a model's output depends on a specific intermediate representation. It works by measuring the average treatment effect of intervening on a mediator variable. In the context of causal tracing, this framework formalizes the process of identifying hidden states as mediators that carry the causal effect from an input prompt to a factual answer.

The core data pathway in a transformer architecture where each layer reads from and writes additive updates to a running hidden state. Causal tracing often targets specific positions within this stream because it serves as the primary communication channel. The residual stream accumulates information from attention heads and MLP layers, making it the central object of study for understanding how facts are composed.

A specific neuron in a feed-forward network whose activation is strongly correlated with the expression of a particular factual association. While causal tracing identifies the hidden states responsible for recall, knowledge neuron analysis pinpoints the individual units within those states. These neurons are often identified through activation analysis and can be manipulated to edit stored facts.

A decoding technique that applies the unembedding matrix directly to intermediate residual stream states. This allows researchers to inspect a model's iterative token predictions before the final layer. When combined with causal tracing, the logit lens can reveal how a restored hidden state incrementally shifts the output probability distribution toward the correct factual answer.

A causal technique that removes or zeroes out a model component, such as a neuron or attention head, to measure the resulting drop in performance. In contrast to causal tracing's restorative approach, ablation is a subtractive method. Comparing the results of both techniques provides a more complete picture of a component's functional importance and whether it is strictly necessary for a computation.