Learn AI Create the Future — with Pranay
Main site

Looking Inside a Trained Model

Where a model stores a fact

Facts a model has memorised, like a country's capital, tend to live in the feedforward sublayers of the middle layers, not in the attention sublayers.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. Where you have already seen it
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A model's memorised facts, like a country's capital, tend to live in one specific kind of layer. They are not spread evenly everywhere.

Recalling your own phone number is instant. You are not working it out, you are pulling it from memory. Following turn-by-turn directions to a new address is different: you are actively combining information as you go.

A transformer has both kinds of machinery. One part, the feedforward layer, works something like memory lookup. Another part, attention, works something like following directions between words. Stored facts lean heavily on the first.

Why it exists

Earlier lessons showed you can test which layer matters, using activation patching. That answers "which layer", but a transformer layer is actually built from two different pieces stacked together.

Researchers wanted a sharper answer: not only which layer, but which piece inside that layer. Is a fact more like a memory lookup, or more like following a chain of connections between words?

Separating those two pieces, and testing each on its own, gives a real answer.

How it works

  Prompt: "The capital of France is"

  Remove the ATTENTION piece at one layer -> answer barely changes
  Remove the FEEDFORWARD piece at that same layer -> answer gets much worse

  This particular fact leans on the feedforward piece,
  more than the attention piece, at that layer.

Repeat this test, layer by layer and piece by piece. It builds a picture of where a specific memorised fact actually lives.

Where you have already seen it

  • Model editing tools. Techniques that let researchers directly change a stored fact, like updating a company's CEO, target this exact kind of layer.
  • Explaining a wrong fact. When a model confidently states something false, tracing which layer produced it helps explain why.
  • AI safety research on "unlearning". Attempts to make a model forget a specific fact often target these same feedforward layers directly.

Remember this

  • A transformer layer has two pieces: attention, and feedforward. They do different jobs.
  • Memorised facts tend to lean on the feedforward piece more than the attention piece.
  • Testing each piece separately gives a sharper answer than testing a whole layer at once.

What to learn next

Developer — Code and libraries.

This removes the attention piece and the feedforward piece of each layer separately, one at a time, and compares which removal hurts the "Paris" prediction more.

Setup

bash
pip install transformers torch

The first run downloads gpt2, roughly 500 MB.

Testing attention versus feedforward, layer by layer

mlp_fact_demo.py
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

tok = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.eval()

prompt = tok("The capital of France is", return_tensors="pt")
paris_id = tok(" Paris")["input_ids"][0]

def prob_of_paris(zero_mlp_layer=None, zero_attn_layer=None):
    handles = []
    if zero_mlp_layer is not None:
        def kill_mlp(module, inputs, output):
            return torch.zeros_like(output)
        handles.append(model.transformer.h[zero_mlp_layer].mlp.register_forward_hook(kill_mlp))
    if zero_attn_layer is not None:
        def kill_attn(module, inputs, output):
            return (torch.zeros_like(output[0]),) + output[1:]
        handles.append(model.transformer.h[zero_attn_layer].attn.register_forward_hook(kill_attn))
    with torch.no_grad():
        logits = model(**prompt).logits[0, -1]
    for h in handles: h.remove()
    return torch.softmax(logits, dim=-1)[paris_id].item()

baseline = prob_of_paris()
print(f"baseline, nothing removed: P(Paris) = {baseline:.4f}\n")

for layer in range(12):
    p_no_mlp = prob_of_paris(zero_mlp_layer=layer)
    p_no_attn = prob_of_paris(zero_attn_layer=layer)
    print(f"layer {layer:2}: remove feedforward -> {p_no_mlp:.4f}   remove attention -> {p_no_attn:.4f}")
Output
baseline, nothing removed: P(Paris) = 0.0322

layer  0: remove feedforward -> 0.0009   remove attention -> 0.0007
layer  1: remove feedforward -> 0.0133   remove attention -> 0.0256
layer  2: remove feedforward -> 0.0089   remove attention -> 0.1046
layer  3: remove feedforward -> 0.0037   remove attention -> 0.0203
layer  4: remove feedforward -> 0.1059   remove attention -> 0.0373
layer  5: remove feedforward -> 0.0312   remove attention -> 0.0418
layer  6: remove feedforward -> 0.0984   remove attention -> 0.0190
layer  7: remove feedforward -> 0.0533   remove attention -> 0.0185
layer  8: remove feedforward -> 0.0164   remove attention -> 0.0245
layer  9: remove feedforward -> 0.0155   remove attention -> 0.0174
layer 10: remove feedforward -> 0.0216   remove attention -> 0.0322
layer 11: remove feedforward -> 0.0589   remove attention -> 0.0205

Read this carefully, because it is genuinely messy, not a clean story. Layer 0 breaks almost everything, removing either piece, since early layers are foundational to the whole computation. Layers 8 and 9 show removing the feedforward piece hurts noticeably more than removing attention, matching the "facts live in mid-layer feedforward sublayers" idea from the research literature. But layer 2 shows removing attention actually raises P(Paris) well above baseline, a genuinely surprising result this single example does not fully explain.

Line by line

kill_mlp and kill_attn return all-zero tensors in place of that sublayer's real output, removing its contribution to the residual stream entirely for that forward pass.

kill_attn handles a tuple, kill_mlp does not. GPT-2's attention sublayer returns (output, present_key_value_or_None) in this library version, while the feedforward sublayer returns a plain tensor. Mismatching these breaks the model with a confusing error, worth checking for in your own code.

Only one prompt is tested here. Research papers on this topic average results over hundreds or thousands of factual prompts, specifically to smooth out noisy single-example results like layer 2's surprising attention finding above.

Common mistakes

Trusting a single-prompt result as a general finding. As shown above, one prompt can show a confusing, non-monotonic pattern that a large, averaged study would smooth out. Treat this script as a demonstration of method, not a conclusion.

Zeroing out a sublayer instead of patching it. Zeroing removes information entirely, a harsher intervention than patching in a different, still-plausible value. Both are used in the literature; they answer subtly different questions, and are worth telling apart.

Ignoring layer 0's results. Early layers tend to matter for almost everything, so a huge drop there is not evidence of anything specific to this fact. Middle layers are the more informative signal.

Try it yourself

Change the prompt to "The capital of Japan is" and rerun. Compare the pattern of which layers matter most against the France example.

Expect the general shape, mid-layer feedforward sublayers mattering more, to roughly hold, while the exact layer numbers and magnitudes shift. That prompt-to-prompt variation is itself the reason real studies average over many examples.

What to learn next

Researcher — Mathematics and papers.

The key-value memory view of feedforward layers

Geva et al. (2021), Transformer Feed-Forward Layers Are Key-Value Memories, propose reading a feedforward sublayer as an associative memory:

text
FFN(x) = f(x @ W_in) @ W_out
  • W_in's columns act as keys: each hidden unit activates when the input matches a learned pattern.
  • W_out's rows act as values: each activated unit adds its corresponding row into the residual stream, effectively voting for tokens in the vocabulary.
  • f is the nonlinearity (GELU in GPT-2), producing sparse-ish activation over the large intermediate dimension.

Under this reading, a "memorised fact" corresponds to a small set of key-value pairs, keys that fire specifically on the fact's subject, values that push the residual stream toward the correct answer token.

Causal tracing localises facts to mid-layer feedforward sublayers

Meng et al. (2022), introduced already in the activation-patching lesson, ran causal tracing across the CounterFact dataset of factual-recall prompts, patching clean activations back into a corrupted run at every layer, sublayer and position.

Averaged over many prompts, the strongest restoration effect concentrates at the feedforward sublayer, at the last subject token position, in a narrow band of middle layers (roughly layers 4 through 8 of GPT-2-XL's 48, proportionally similar in smaller models). This is a genuine statistical regularity across many facts, sharper than the noisy single-prompt result reproducible in the developer block.

From localisation to editing: ROME and MEMIT

Meng et al. (2022) used this localisation directly to build ROME (Rank-One Model Editing). Treat the identified feedforward layer as a linear associative memory. Then compute a targeted rank-one weight update that changes one key-value association, for example redirecting "The Eiffel Tower is located in" from Paris to Rome, while leaving unrelated facts and fluency intact.

Meng et al. (2023) extended this to MEMIT, editing thousands of facts at once. It distributes the update across several middle layers rather than one, addressing interference that arises when too many facts share a single layer's limited capacity.

Known limitations of the localisation claim

Hase et al. (2023), Does Localization Inform Editing?, present a striking finding: the layer where causal tracing shows the strongest effect is not always the most effective layer to edit at, for actually changing the model's stated belief. Localisation (where information is causally used for this specific forward pass) and editability (where a weight change most reliably changes future behaviour) turn out to be related but not identical questions.

This is an important caveat on the entire section: causal tracing tells you where information flows for a given input, not necessarily the single best lever for a downstream engineering goal like editing.

Complexity

Causal tracing at the sublayer level costs O(L * 2 * n) forward passes, patching attention and feedforward separately at every layer and token position, for sequence length n and L layers. ROME-style editing, once a target layer is identified, is a single closed-form weight update, cheap relative to the search that located it.

Key references

  • Geva, M. et al. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. arXiv:2012.14913
  • Meng, K. et al. (2022). Locating and Editing Factual Associations in GPT. arXiv:2202.05262
  • Meng, K. et al. (2023). Mass-Editing Memory in a Transformer. arXiv:2210.07229
  • Hase, P. et al. (2023). Does Localization Inform Editing? Surprising Differences in LLMs' Factual Knowledge. arXiv:2301.04213

Current state and open problems

Mid-layer feedforward localisation of facts is one of the most replicated findings in mechanistic interpretability, and directly enabled a practical line of model-editing tools built on top of it.

The Hase et al. finding is the open problem worth taking seriously. Localisation and editability can diverge. "Where a fact lives" is not a single, settled question with one answer. It depends on what operation you intend to perform: reading it out, or changing it.

What to learn next