Dissecting Recall of Factual Associations in Auto-Regressive Language Models

In plain words

A language model that answers "Beats Music is owned by" with "Apple" must combine two things: which thing the question is about (Beats Music) and what kind of fact is asked for (the owner). This paper asks how the model does that inside its layers, not where the fact is stored.

The authors switch off one information route at a time. They block the last word of the question from looking at the subject words, or at the relation words, in a range of layers, then check how much the right answer's probability drops. This shows that the relation words reach the last word first, and the subject words reach it later.

Next they look at what the subject's own representation holds. Early in the network, the model fills the representation of the subject's last word with many facts about that subject. Later attention heads read the answer out of that bundle, and some of those heads seem to store particular subject-to-answer pairs in their weights. The authors suggest this view can guide later work on finding and editing stored facts.

TL;DR

• Analyses GPT-2 XL §3 and GPT-J §3 on CounterFact queries whose top prediction is correct §3.
• Attention knockouts locate two information flows into the last position: relation first, then subject §5.
• Early MLP sublayers drive the enrichment of the subject representation §6.3.
• Upper MHSA sublayers extract the attribute for 68.2% of GPT-2 XL queries tab. 2, through heads that often store subject-attribute mappings §7.4.

Problem

Models store facts in their parameters, and earlier work asked where facts are stored, for example in mid-layer MLPs §1. The question here is how the model pulls the fact out during inference §1. The setting is a subject-relation query whose next token is the attribute, with the model's top prediction correct §1 §3. Only the top-1 prediction is analysed, so the findings describe queries the model already answers correctly §3.

Method

Notation. The residual stream at token i and layer ℓ is the sum of the previous state, the MHSA output and the MLP output eq. 1. The output distribution is the softmax of the prediction head applied to the final state, which is the tied embedding matrix or a trained linear layer eq. 2 §2. Each MHSA sublayer is a sum over heads, with a causal mask §2. Each MLP sublayer is a two-layer network with activation σ eq. 6 §2.

Attention knockout (§5). For a source position r and target position c, the method blocks the attention edge from r to c in a chosen set of layers by adding −∞ to the mask in all heads eq. 7 §5. Blocking the last position from the subject tokens, the relation tokens (all non-subject tokens) or itself, over a window of k layers around layer ℓ, gives the change in p(answer) §5. The window is k = 9 for GPT-2 XL and k = 5 for GPT-J §5. The paper contrasts this with causal tracing, which asks where information is located, while knockout asks where it propagates §5.

Subject representation (§6). The representation at the last subject token is projected to the vocabulary with the same projection as the output head §6.1. Its top-50 tokens are scored with the attribute rate, which is the fraction that appear in a candidate set A_s built from Wikipedia paragraphs about the subject §6.1 app. C.1. The same rate is computed for the first subject token, the token after the subject, and the last input token, for comparison §6.1. Static subject token embeddings give much lower rates §6.2.

Sublayer knockout (§6.3). For a window of 10 consecutive layers, the MHSA or MLP update at the last subject position is set to zero. The effect is measured by the attribute rate at layer 40 of GPT-2 XL §6.3. Zeroing early MLP updates lowers that rate by about 88% on average; zeroing MHSA updates lowers it by less than 30% §6.3.

Extraction events (§7.1). At every layer ℓ, the MHSA output and the MLP output at the last position are projected to the vocabulary. An extraction event is a layer where the top-1 token equals the model's answer §7.1. The per-query extraction rate counts queries with at least one event in any layer tab. 2. The paper also reports the rank of the answer in the subject representation at extraction events, which averages 999.5 §7.2.

Patching (§7.3). Early-layer representations at the subject, relation or last positions are patched into the model, and the extraction rate is measured again §7.3. Patching the subject positions lowers the rate by up to 50% §7.3.

Knowledge heads (§7.4). For a head, the matrix G = Eᵀ W_VO E is built eq. 8. For each extraction event, the paper checks whether the answer is among the top 10 tokens of the row of G for any subject token §7.4.

# Attention knockout over a window of layers  [§5] [eq. 7] [app. A.4]
# Input:  model with num_layers layers, token positions, the set of
#         blocked source positions, the layer under test, window size k
# Output: relative change in the answer probability (percent)

def knockout_change(model, blocked, layer, window, base_prob):
    # Window of layers: from layer - k//2 up to layer + (k+1)//2, exclusive
    # end, clipped to the model's layers  [code: notebook#L597-L606]
    first = max(0, layer - window // 2)
    last = min(model.num_layers, layer + (window + 1) // 2)
    # In each layer of the window, block the last position's attention
    # to every blocked position, in all heads  [eq. 7]
    edges = {
        other: [(last_position, target) for target in blocked]
        for other in range(first, last)
    }
    new_prob = answer_probability(model, edges)
    return (new_prob - base_prob) * 100.0 / base_prob

# Extraction events for one query  [§7.1] [tab. 2]
# Input:  one update vector per layer (MHSA or MLP) at the last position,
#         the embedding matrix E, the model's answer token
# Output: per-layer booleans, True where the update predicts the answer

def extraction_events(updates, embedding, answer_token):
    events = []
    for update in updates:
        # Score every vocabulary token by projecting the update  [§7.1]
        scores = embedding @ update
        # An event: the top-1 token is the model's answer  [§7.1]
        events.append(argmax(scores) == answer_token)
    return events

# Attribute rate of one representation  [§6.1] [app. C.1]
# Input:  hidden state (d,), embedding E, candidate set A_s for the subject
# Output: fraction of the top-50 tokens that are in A_s

def attribute_rate(hidden, embedding, candidates, top_m=50):
    # Vocabulary scores of the representation  [§6.1] [eq. 2]
    scores = embedding @ hidden
    # Keep the top_m tokens after filtering  [code: notebook#L1420-L1446]
    top_tokens = top_k_indices(scores, top_m)
    # Count how many are candidate tokens of the subject  [§6.1]
    hits = [token for token in top_tokens if token in candidates]
    return len(hits) / top_m

Hyperparameters

Results

Author-reported. Setup: CounterFact queries, GPT-2 XL §3 and GPT-J §3, inference only.

The table below gives the extraction statistics of GPT-2 XL tab. 2. Each MHSA row blocks the attention from the last position to the named positions; "all but X" blocks every position except X. The last row is the MLP update instead. The extraction rate is the fraction of queries with at least one extraction event; the layer count averages over those queries.

Source: tab. 2. The row labels are restated in the card's own words.

Other headline numbers. The introduction gives an extraction rate of about 70% across predictions §1, which matches the 68.2% above. The extraction layers coincide with the layers where subject information propagates to the last position §7.1 fig. 2 fig. 5. Patching early layers at non-subject positions raises the extraction rate from 0.05 to 0.59 in the first layer, so the first layer does most of the extra work for those positions §7.3.

Paper vs code

The paper's code link (abstract footnote) points to the dissecting_factual_predictions directory of the google-research repository. It is one Jupyter notebook and one helper module, at commit 3ef52111c1fabc301880818f5a547570e926a76b. The code is GPT-2 XL only, and the helper hooks state this [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L126] factual_associations_dissection.ipynb.

Differences, each checked by reading the code:

• Dataset: the paper uses 1,209 GPT-2 XL queries with a correct prediction §3. The loop runs over all 1,000 entries of the CounterFact "known" file, with no check that the prediction is correct [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L83-L84] factual_associations_dissection.ipynb.
• Knockout indexing: the paper writes the knockout as an update to mask M^{ℓ+1} eq. 7. The code masks the attention of the listed layer itself factual_associations_dissection.ipynb. The window contents match, but the paper's index convention differs by one.
• MLP sublayer knockout: the paper sets the MLP update at the last subject position to zero §6.3. The code sets the pre-activation of the MLP to zero at every position [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L256-L271] and runs it with all positions [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L1337]. With a GELU activation, that makes the MLP output equal to its output bias, not zero, and it changes all positions, not one.
• MHSA sublayer knockout: the code uses an attention mask with no allowed edges in those layers [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L1342] factual_associations_dissection.ipynb. Read literally, every logit gets the same additive value, so the softmax becomes uniform over all positions, including later ones. That is not the zero update of §6.3. Not verified by running the code.
• Patching: the paper feeds patched representations to MHSA at a later layer ℓ' > ℓ §7.3. The code patches the block input of one layer at a time and reads the attention output of that same layer, for every layer, including earlier ones [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L1040-L1060].
• Candidate paragraphs: the code expects a CSV of paragraphs and has no path set factual_associations_dissection.ipynb. The BM25 retrieval of the paper is not in the code.
• Attribute rate filter: the code filters the top-500 tokens for stopwords and short tokens, then keeps the first 50 [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L1446]. The paper does not state the order.
• Position prefix: for the token-embedding projection, the code prepends an end-of-text token to fix the first-position effect [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L1265]. The paper does not mention this step.
• Precision and device: the model loads in default precision and moves to CUDA utils.py. The paper gives no hardware or precision.
• GPT-J: the paper reports GPT-J results §3 app. E. The code has GPT-2 hooks only, so GPT-J results cannot be reproduced from it.

Pitfalls

• The attribute is the model's first output token, which can be a word fragment such as "Wash" for "Washington" §3. Match on tokens, not on whole words.
• Only queries whose top prediction is correct are analysed §3. Rates are conditional on that, so they do not describe all queries.
• Window size is a confound. Effects grow with k app. A.4. With k = 1, blocking the relation positions in GPT-2 XL shows no significant change, which the paper reads as relation information being processed over several layers app. A.4.
• Edge clipping: near layer 0 or layer L−1 the window has fewer layers than k factual_associations_dissection.ipynb. Plots at the edges therefore use a smaller window.
• First-position bias: blocking attention to the first token alone changes predictions a lot. Blocking the subject positions drops the probability by up to 40% when the first position is open, against 60% when it is blocked app. A.2 fig. 2. The main results include the blocked first position. Reproduce both settings if you want the subject effect in isolation.
• Subject-relation order matters: the relation-first subset shows a more spread-out effect app. A.1. The authors attribute this to the first-position bias app. A.1.
• Conflict inside the paper: the conclusion says the MLP sublayers extract the attribute "eminently" §9. Section 7, the abstract and the introduction say the upper MHSA sublayers extract it §7 §1. This card follows §7 and tab. 2.
• Attribute rate is a proxy. It counts how many top tokens appear in Wikipedia-derived candidate sets, so an unrelated frequent word that appears in those paragraphs also counts app. C.1. Subjects with no matching paragraph are dropped factual_associations_dissection.ipynb.
• Extraction is strict: it needs an exact top-1 match with the answer §7.1. The per-layer rates and the per-query rate are different numbers. Use the one the statement names.
• Knockouts leak across layers. The paper notes that blocking one layer does not stop information passed on earlier, so it blocks a range of layers instead [Limitations, unnumbered section].
• Projections to the vocabulary approximate what early layers encode [Limitations, unnumbered section]. Treat attribute rates in early layers as a lower-quality signal.
• Cost: for each prompt, the main knockout cell runs 3 blocked groups × 48 layers = 144 forward passes on GPT-2 XL [code: dissecting_factual_predictions/factual_associations_dissection.ipynb#L577-L606]. The per-position map runs one forward pass for each position and layer factual_associations_dissection.ipynb.
• Not stated: hardware, batch size, precision, random seed and the exact CounterFact subset used in the paper.

Lineage

• Builds on: causal tracing (Meng et al., cited as the source of the CounterFact data and the locating view) §1 §3 §5. The paper's knockout method asks where information propagates, not where it is located §5. It extends the MLP key-value view of stored facts and the knowledge-neuron work §6.3 §8.
• Compared against: no baseline method is compared; the paper compares sublayers and positions within the same models §6.3 §7.1.
• Superseded by: not stated in this wiki.