Localizing Factual Recall in Transformer MLP Layers

Motivation

Large language models can recall thousands of factual associations, from capital cities to chemical formulas, yet we still lack a mechanistic account of where and how these facts are stored.¹ Understanding the internal structure of factual recall is important for AI safety: if we can localize knowledge to specific components, we gain tools for auditing model beliefs, detecting hallucinations, and performing targeted edits without retraining.

This post walks through causal tracing, a technique introduced by Meng et al. (2022), and its follow-up method Rank-One Model Editing (ROME). We will cover the core ideas, the math behind the intervention, and a minimal implementation.

---

Background: Transformer MLP Layers as Key-Value Memories

The standard transformer block

Recall that each transformer layer $\ell$ applies two sub-layers to a residual stream. Given an input vector $\mathbf{h}^{(\ell)}$, the MLP sub-layer computes:

$$\mathbf{m}^{(\ell)} = W_{\text{proj}}^{(\ell)} \, \sigma\!\left(W_{\text{fc}}^{(\ell)} \, \mathbf{h}^{(\ell)} + \mathbf{b}_{\text{fc}}^{(\ell)}\right)$$

where $\sigma$ is a nonlinearity (typically GELU) and $W_{\text{fc}}, W_{\text{proj}}$ are the two weight matrices.²

The key-value interpretation

Geva et al. (2021) observed that MLP layers behave like key-value memories:

• The rows of $W_{\text{fc}}$ act as keys that match input patterns.
• The columns of $W_{\text{proj}}$ act as values that are retrieved when a key fires.
• The nonlinearity acts as a soft gating mechanism, selecting which memories to read.

This means a single MLP layer can be thought of as storing a sparse lookup table over learned patterns. A factual association like “The Eiffel Tower is located in Paris” may correspond to a small number of key-value pairs across a handful of layers.

“The MLP layers of a transformer are not merely nonlinear feature extractors. They function as distributed associative memories, with each neuron contributing a fragment of the retrieval.” — Geva et al., 2021

Causal Tracing

Setup

The core question is: which hidden states are causally responsible for a model’s factual prediction? Causal tracing answers this by running three forward passes:³

1. Clean run — pass the prompt through normally and record all hidden states $\mathbf{h}_i^{(\ell)}$ at every token position $i$ and layer $\ell$.
2. Corrupted run — replace the subject tokens with noise (e.g., Gaussian perturbation of their embeddings) and run the model again. The correct answer probability typically drops.
3. Corrupted-with-restoration — start from the corrupted run but patch in the clean hidden state at a single position $(i, \ell)$. Measure how much the correct answer probability recovers.

The indirect effect

We quantify the average indirect effect (AIE) of restoring a particular state:

$$\text{AIE}(i, \ell) = P\!\left[\,y^* \mid \text{do}\!\left(\mathbf{h}_i^{(\ell)} \leftarrow \mathbf{h}_{i,\text{clean}}^{(\ell)}\right)\right] - P\!\left[\,y^* \mid \text{corrupted}\right]$$

where $y^*$ is the correct next token. A high AIE at position $(i, \ell)$ means that state is a critical mediator for the factual prediction.

What the results show

The causal tracing map reveals a consistent pattern across many GPT-style models:

Component	Token position	Layers	Effect
MLP output	Last subject token	Mid layers (15—25 in GPT-J)	Strong decisive restoration
Attention output	Last subject token	Early layers	Moderate partial effect
MLP output	Non-subject tokens	All layers	Negligible
Full residual	First token	Late layers	Weak indirect effect

The dominant finding is that MLP outputs at the last subject token in middle layers carry the bulk of the causal effect.⁴ This is the site where factual recall is concentrated.

Rank-One Model Editing

Problem formulation

Given a factual association we want to change (e.g., replacing “Paris” with “Rome” for the Eiffel Tower query), ROME directly modifies the MLP weight matrix at the critical layer $\ell^*$ identified by causal tracing.

We model the MLP as a linear associative memory:

$$W_{\text{proj}} \, \mathbf{k} = \mathbf{v}$$

where $\mathbf{k}$ is the key vector (the hidden state when processing the subject) and $\mathbf{v}$ is the value vector (encoding the fact to be recalled). To insert a new association $(\mathbf{k}_1, \mathbf{v}_1)$, we compute a rank-one update:

$$\hat{W}_{\text{proj}} = W_{\text{proj}} + \frac{(\mathbf{v}_1 - W_{\text{proj}} \, \mathbf{k}_1) \, \mathbf{k}_1^\top}{\mathbf{k}_1^\top \mathbf{k}_1}$$

This is essentially the Sherman-Morrison formula applied to insert a single key-value pair while preserving existing associations as much as possible.⁵

Computing the key vector

The key vector $\mathbf{k}_1$ is not simply the hidden state from one prompt. Instead, ROME averages over multiple contexts to extract a context-invariant subject representation:

import torch
from transformer_lens import HookedTransformer

def compute_subject_key(
    model: HookedTransformer,
    subject: str,
    prompts: list[str],
    layer: int,
) -> torch.Tensor:
    """Average the MLP input at the last subject token across prompts."""
    keys = []
    for prompt in prompts:
        tokens = model.to_tokens(prompt)
        _, cache = model.run_with_cache(tokens)

        # Find the last token position of the subject
        subject_tokens = model.to_tokens(subject, prepend_bos=False)
        subject_len = subject_tokens.shape[1]
        last_subject_pos = find_subject_end(tokens, subject_tokens)

        h = cache[f"blocks.{layer}.mlp.hook_pre"][0, last_subject_pos]
        keys.append(h)

    return torch.stack(keys).mean(dim=0)

Computing the value vector

The target value vector $\mathbf{v}_1$ is optimized so that, when inserted, the model produces the desired new fact. This is done by gradient descent on the value vector directly:

$$\mathbf{v}_1 = \arg\min_{\mathbf{v}} \; \mathbb{E}_{p \,\sim\, \mathcal{P}} \left[-\log P\!\left(o_{\text{new}} \mid p;\, \hat{W}_{\text{proj}}(\mathbf{v})\right)\right] + \lambda \|\mathbf{v}\|^2$$

where $\mathcal{P}$ is a set of paraphrase prompts and $o_{\text{new}}$ is the desired output.⁶

Pseudocode for the full edit

function ROME_edit(model, subject, layer, new_fact):
    # Step 1: Compute context-invariant key
    prompts <- generate_diverse_prompts(subject)
    k1 <- average_mlp_input(model, subject, prompts, layer)

    # Step 2: Optimize target value
    v1 <- optimize_value(model, layer, k1, new_fact)

    # Step 3: Apply rank-one update
    W <- model.layers[layer].mlp.W_proj
    delta <- outer_product(v1 - W @ k1, k1) / dot(k1, k1)
    model.layers[layer].mlp.W_proj <- W + delta

    return model

Evaluation

A good model edit should satisfy three criteria:

• Efficacy — the model now produces the new fact when asked directly
• Paraphrase generalization — the edit holds across different phrasings of the same question
• Specificity — unrelated facts remain unchanged

Meng et al. report that ROME achieves near-perfect efficacy and strong generalization while maintaining specificity, outperforming fine-tuning baselines that tend to damage neighboring facts.

Limitations

There are several important caveats to keep in mind:

1. ROME edits a single layer, but factual recall may involve distributed computation across multiple layers and attention heads.
2. The rank-one assumption limits each edit to one fact at a time. Scaling to many simultaneous edits requires the more complex MEMIT algorithm, which distributes updates across layers.⁷
3. Causal tracing identifies where information flows, not necessarily where it is stored. The distinction between mediation and storage remains an open question.
4. Current evaluations test surface-level recall. Whether edits propagate to downstream reasoning (e.g., answering “What country is the Eiffel Tower in?” after editing its city) is less well-studied.
   • Recent work by Cohen et al. (2024) suggests that edited models often fail at multi-hop inference.
   • This points to a deeper issue: factual knowledge may not be as modular as the editing framework assumes.

Connections to AI Safety

Auditing model beliefs

If we can reliably trace which components encode specific facts, we can build tools for knowledge auditing: systematically checking what a model believes and flagging inconsistencies or dangerous knowledge.

Targeted model correction

Rather than retraining a model from scratch when it produces harmful outputs, localized editing offers a surgical alternative. This is especially valuable for deployed models where full retraining is expensive.

Mechanistic anomaly detection

Understanding normal patterns of factual recall gives us a baseline against which to detect anomalous internal computation. If a model’s internal reasoning path for a safety-relevant query looks different from the expected pattern, that could serve as an early warning signal.

The long-term vision of mechanistic interpretability is not just to understand models, but to build the verification infrastructure that makes advanced AI systems trustworthy.

---

Open Questions

How distributed is factual storage really?

Causal tracing reveals dominant sites of mediation, but this does not rule out a long tail of smaller contributions from other components. Developing better tools for characterizing the full distribution of causal effects remains important.

Can we move beyond rank-one edits?

The linear associative memory model is a useful approximation, but real MLP computations involve nonlinearities that may encode facts in more complex ways. Understanding the geometry of how facts are superposed in MLP weight matrices is an active area of research.

What is the relationship between factual recall and in-context learning?

Models can also acquire new facts from their context window. Whether in-context facts and parametric facts share the same retrieval mechanisms is an open and fascinating question.

Footnotes

1. The broader context is the Circuits research agenda pioneered at Anthropic, which aims to reverse-engineer neural networks into understandable computational graphs. Factual recall is one of the most concrete sub-problems within this agenda, since the input-output behavior is easy to specify and test. ↩

2. In some architectures like LLaMA, the MLP uses a gated structure with three matrices instead of two: $\mathbf{m} = W_{\text{down}} \left(\sigma(W_{\text{gate}} \mathbf{h}) \odot W_{\text{up}} \mathbf{h}\right)$. The key-value memory interpretation extends naturally to this setting, with the gating mechanism providing sharper key selection. ↩

3. The corruption is typically Gaussian noise added to the embedding vectors of the subject tokens, with a standard deviation calibrated to three times the empirical standard deviation of the embedding space. This ensures the corruption is strong enough to destroy the subject identity without pushing activations out of distribution entirely. ↩

4. Why the last subject token? In autoregressive models with causal attention, the last token of a multi-token subject is the only position that has attended to all preceding subject tokens. It therefore serves as a natural aggregation point for the full subject representation. This is consistent with findings from probing studies showing that entity representations are most complete at the final token. ↩

5. The connection to the Sherman-Morrison formula becomes clearer when you consider the constraint that $\hat{W}$ should satisfy $\hat{W}\mathbf{k}_1 = \mathbf{v}_1$ while minimizing $\|\hat{W} - W\|_F$. The solution is exactly a rank-one perturbation. In practice, ROME also uses a second-moment matrix $C = \mathbb{E}[\mathbf{k}\mathbf{k}^\top]$ estimated over a sample of Wikipedia text, replacing the denominator with $C\mathbf{k}_1$ to better preserve existing associations under the empirical key distribution. ↩

6. The regularization term $\lambda\|\mathbf{v}\|^2$ prevents the optimized value from drifting too far from the manifold of typical value vectors, which helps maintain model coherence on unrelated inputs. ↩

7. MEMIT (Mass-Editing Memory in a Transformer) generalizes ROME by distributing the rank-one updates across a range of critical layers rather than concentrating the edit at a single layer. This reduces the per-layer perturbation magnitude and improves specificity when editing many facts simultaneously. ↩