Causal Attribution & Patching (Part 1): Internal Mechanisms of Language Models

What They Mean by Causal Attribution (vs. Correlational Probing)

For years, a dominant paradigm in interpretability has been correlational probing. A typical study might train a simple linear classifier (a "probe") on the internal activations of a large language model (LLM) to predict a certain property, like the part-of-speech of a token or the sentiment of a sentence. If the probe achieves high accuracy, the conclusion is that the property is "encoded" in the model's representations. (Probing)

This is an essential first step, but it leaves a critical question unanswered: just because information is present, does that mean the model actually uses it to make its next prediction?

Imagine a car's dashboard. A probe could easily "decode" the car's speed by looking at the radio's volume knob-perhaps the driver turns up the music on the highway. There is a correlation. But the volume knob does not cause the car to accelerate. The speedometer and the engine do. Causal attribution is the effort to find the speedometer, not just another correlated signal.

This is the distinction between information being merely encoded versus being causally implicated in the model's behavior. An internal state can contain linearly decodable information about a concept, but that information might be an artifact of the training data or a precursor to a different computation, rather than a direct input to the final decision.

As Vig et al. (2020) note in their work on causal mediation, many interpretability methods "show only that information is present in the model’s representations, not that it is used." Causal methods, by contrast, rely on interventions. The core idea is simple: if you think component $M$ is causally responsible for behavior $Y$, then changing or disabling $M$ should change $Y$. If it doesn't, your hypothesis is likely wrong, no matter how much information a probe can extract from $M$.

Background & Formalization

The Intervention Model: Corrupt vs. Clean Runs

The foundational tool for causal analysis in LLMs is the patching experiment. This involves running the model on at least two different inputs and swapping internal states between them.

The experiment is to run the model on the corrupted input but, at a specific point in the forward pass, intervene and "patch" in a clean state from the original run. For instance, we might replace the output of the MLP module in layer $l=15$ from the corrupted run with the MLP output from the clean run, and then let the computation proceed.

We measure the effect of this intervention on the model's output logits. Specifically, we look at the change in the logit of the target token (e.g., "Paris"). We can formalize this as a restoration score:

$\Delta \text{logit}(t^\star) = \text{logit}{\text{restore}(M)}(t^\star) - \text{logit}{\text{corrupt}}(t^\star)$

Here, $t^\star$ is our target token ("Paris"). $\text{logit}{\text{corrupt}}(t^\star)$ is the (presumably low) logit for "Paris" on the corrupted run. $\text{logit}{\text{restore}(M)}(t^\star)$ is the logit for "Paris" when we patch in the clean state of a specific component, or mediator, $M$.

A large, positive $\Delta \text{logit}$ is strong evidence that the component $M$ is causally implicated in producing the token "Paris" for the clean prompt. It means that restoring just this one piece of the computation was sufficient to recover much of the correct answer.

The Idea of Causal Mediation

Causal mediation analysis provides a formal language for these interventions. It originates from fields like epidemiology and social sciences, where researchers want to understand how a treatment causes an outcome. For example, how does an educational program (treatment $X$) lead to higher salaries (outcome $Y$)? Is it by increasing a student's confidence, or by teaching them a specific skill (mediator $M$)?

The framework decomposes the Total Effect (TE) of $X$ on $Y$ into two pathways:

$TE = NDE + NIE$

In LLMs, the "treatment" $X$ is the input prompt, and the "outcome" $Y$ is the output distribution. The "mediators" $M$ are the internal components we hypothesize are on the causal path: individual neurons, attention heads, or entire MLP modules. By estimating the indirect effect (NIE) for each component, we can quantify how much of the model's final behavior is attributable to that specific part of the network.

Paper 1 - Causal Mediation in Language Models (NeurIPS 2020)

Vig et al. (2020) were among the first to formally apply causal mediation analysis to LLMs to locate the sources of a specific behavior: gender bias.

• Goal: The researchers wanted to move beyond simply showing that gender is encoded in representations and instead identify which specific model components (attention heads and neuron activations) are causally responsible for producing biased predictions. For example, in a prompt like "The doctor said that," which components cause the model to assign a higher probability to "he" than "she"?
• Method: They framed this as a mediation problem. The "treatment" was the gender of a subject in a sentence. The "outcome" was the probability distribution over subsequent gendered pronouns. They then systematically estimated the Natural Indirect Effect (NIE) for each attention head and neuron in the model, effectively measuring how much of the gender-bias effect was flowing through that component.
• Result: The analysis revealed that gender bias was not a diffuse, entangled property of the entire network. Instead, the causal effects were highly concentrated in a small number of specific attention heads and neurons. For instance, in GPT-2, a handful of heads were consistently identified as "gender-mediating heads." This was a powerful demonstration that complex behaviors can be localized to specific, modular parts of the model.
• Limits & Assumptions: Causal mediation relies on strong statistical assumptions, chiefly the assumption of no unobserved confounding between the mediator and the outcome. In a densely connected neural network, this assumption is difficult to formally satisfy. Furthermore, the identified components are specific to the model, its training data, and the particular task being analyzed. The results may not generalize perfectly to different prompts or model families.

Paper 2 - Linear Mechanisms in Practice (NAACL 2024)

While mediation analysis tells us where an effect is happening, it doesn't always tell us how. The work by Todd et al. (2024) on vector-arithmetic mechanisms provides a compelling hypothesis for the "how," suggesting that the underlying computation can be surprisingly simple and linear.

• Claim: The paper demonstrates that for several common relational tasks, LLMs of various sizes (from 124M to 176B parameters) often implement the solution via a simple, additive update in their Feed-Forward Network (FFN), also known as MLP, layers. This is particularly true when retrieving knowledge stored during pre-training, as opposed to information present in the immediate context.

  For example, to solve the task "Capital of France?", the model might process the representation for "France", $r(\text{France})$, and an MLP layer will compute something akin to:
  $MLP^{(l)}(r(\text{France})) \approx r(\text{France}) + \vec{d}_{\text{capital}}$ is a consistent direction in activation space that transforms a country's representation into its capital's. They find similar additive mechanisms for tasks like uppercasing words and changing verb tense.

• Connection to Patching: This finding provides a powerful explanatory framework for why patching interventions work so well. If a task-relevant computation is just the addition of a specific vector $\vec{d}$ by an MLP module, then a causal patching experiment that restores the output of that MLP should have an enormous effect. It's not just restoring some abstract "information"; it's restoring the result of a precise, near-linear arithmetic operation that is critical to solving the task.

Intuition

Think of the MLP layers as a massive key-value store. The attention layers route the correct query (e.g., the representation of the last token of "Eiffel Tower") to the right MLP layer. That MLP layer then retrieves the associated "value" vector (e.g., a vector that points towards "Paris") and adds it to the residual stream. This additive update pushes the model's overall representation towards the correct answer.

• Caveat: The authors are careful to note that this is not a universal mechanism. LMs are complex, non-linear systems. However, for a significant class of knowledge-retrieval and relational tasks, this simple linear update model appears to be a common and effective strategy implemented by the network.

Paper 3 - Causal Tracing & ROME (NeurIPS 2022 project)

The ROME project (Meng et al., 2022) provides a practical and powerful synthesis of these ideas, using causal interventions not just to locate information but to surgically edit it.

Causal Tracing: Finding the Path

To edit a fact, you first need to know where it lives. ROME introduces a method called causal tracing to do exactly this. It is a direct application of the corrupt-and-restore patching protocol described earlier, executed at scale.

• Method: The model is run on a clean prompt (e.g., "The Eiffel Tower is in") and a corrupted prompt with random token embeddings. Then, the algorithm iterates through every single state in the clean forward pass (every layer, every token position, for both attention and MLP outputs). At each step, it patches that single clean state into the corrupted run and measures the restoration of the probability of the correct answer ("Paris").
• Localization Results: By aggregating the results, a clear picture emerges. The causal effect for retrieving this fact is not distributed evenly. It is highly localized at:
  1. MLP Modules: The effect is overwhelmingly concentrated in the MLP layers, not the attention layers.
  2. Mid-Layers: The critical MLP modules are typically found in the middle layers of the network (e.g., layers 15-20 in a model like GPT-J).
  3. Last Subject Token: The effect is localized to the computation happening at the position of the last token of the subject (the token "Tower" in our example).

This process creates a causal map of factual recall, pointing to a very specific set of modules responsible for storing and retrieving the fact.

ROME: Rewriting Facts with Rank-One Edits

Once a fact is localized to a specific MLP module, ROME (Rank-One Model Editing) performs a direct edit.

• Method: An MLP layer computes $y = W_{out} \sigma(W_{in} x)$. The ROME method treats this layer as a linear key-value memory. It solves for a minimal change to one of the weight matrices (e.g., $W_{in}$) that will map the subject's representation (the "key") to a new target representation (the "value"). This change can be formulated as a rank-one update, meaning the modification matrix $\Delta W$ has a rank of one. This is a very low-complexity and targeted change.
• Evaluation: A successful edit isn't just about changing one fact. The ROME project introduces a crucial evaluation framework based on three criteria, tested on the CounterFact dataset:
  1. Efficacy: Does the model now output the new, edited fact?
  2. Specificity: Does the model still correctly answer questions about other, unrelated subjects? (i.e., did the edit cause collateral damage?).
  3. Generalization: Does the edit generalize to paraphrases of the original prompt? (e.g., "Where is the Eiffel Tower located?").

The results show that these rank-one edits are remarkably effective, achieving high success across all three criteria.

Synthesis - A Unifying View of Patching

These three anchors are not disparate findings; they are three parts of the same story, painting a progressively clearer picture of how models implement knowledge.

Together, they support a "linear pathway + targeted intervention" picture. The model uses attention to route information to the right place, and a specific MLP module at a specific token position performs a vector-arithmetic operation to retrieve a fact stored in its weights.

This picture breaks down when tasks require more complex, non-linear reasoning, multi-step compositional logic, or when the primary mechanism is not a memory lookup but a dynamic computation heavily reliant on attention patterns over the local context.

Caveat

Be meticulous in your reporting. Document the exact prompts, model versions, layers, and token positions tested. Causal effects can be highly specific, and what holds for one model or prompt template may not hold for another. Beware of over-claiming based on a small set of examples.

My View

For practitioners, the most effective workflow seems to be hierarchical. Start with a broad, mediation-style analysis or causal trace to map the computational path and identify the critical few components. Then, use the linear/additive mechanism hypothesis from Todd et al. (2024) to design simple, directional patching experiments to quickly test hypotheses about what the component is doing. Reserve surgical, rank-one edits like ROME for cases where you need to make a durable change to the model's behavior and are prepared to rigorously evaluate the collateral damage.

Intuition:
A key insight unifying these papers is that many "memory" tasks in LLMs seem to be implemented locally and linearly. The FFNs, despite their non-linear activation function, can be driven into a regime where they behave like simple additive key-value stores. Causal tracing and ROME succeed because they correctly identify this implementation strategy and intervene at precisely the right layer and token position where this "memory lookup" occurs.

Caveat:
The success of these methods on factual recall should not be misinterpreted as evidence that the entire LLM is a simple, linear system. These techniques excel on atomized, single-hop tasks. They are less informative about multi-step reasoning, creative generation, or behaviors that emerge from complex interactions between many components. Furthermore, the causal pathways in a base model may be significantly altered by instruction tuning and RLHF, which layer on new, complex behaviors that may not be as easily localized.

Limits & Open Questions

This line of research, while incredibly promising, is still in its early stages. Key open questions remain:

• Portability: How well do localized mechanisms transfer across model sizes and families? Is the fact "The Eiffel Tower is in Paris" stored in a similar location and way in GPT-J, Llama-3, and Claude 3?
• Effect of Finetuning: How does instruction tuning or RLHF reshape these causal pathways? Does it create new, dedicated circuits for following instructions that override or modulate the base model's factual recall mechanisms? Are these new circuits as editable?
• Compositionality: How can we move from editing single, atomic facts to editing compositional knowledge or complex rules? Editing "All birds can fly" is much harder than editing "The ostrich is a bird," as it requires a change that propagates through the model's conceptual space.
• Beyond Factual Recall: What are the causal mechanisms for more abstract tasks, like identifying logical fallacies, summarizing a long document, or adopting a specific persona? Are they as localizable and editable?

Answering these questions will be critical as we move from understanding isolated mechanisms to building a more holistic, causal theory of how these powerful and complex models work.

References

Meng, K., Bau, D., Andonian, A., & Belinkov, Y. (2022). Locating and Editing Factual Associations in GPT. Project page based on work presented at NeurIPS 2022.

Todd, M. E., Mamou, J., & Lieber, O. (2024). Language Models Implement Simple Word2Vec-style Vector Arithmetic. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL).

Vig, J., Geiger, A., & Belinkov, Y. (2020). Investigating Gender Bias in Language Models Using Causal Mediation Analysis. In Advances in Neural Information Processing Systems (NeurIPS).