Summing Up the Facts: Additive Mechanisms Behind Factual Recall in LLMs

Bilal Chughtai; Alan Cooney; Neel Nanda

Source-linked AI summary

Summing Up the Facts: Additive Mechanisms Behind Factual Recall in LLMs

Bilal Chughtai, Alan Cooney, Neel Nanda

arXiv:2402.07321v1Feb 11, 2024cs.LGcs.CL

TL;DR

The paper asks how LLMs store and retrieve factual knowledge in explicit recall prompts. It analyzes the mechanisms behind recall and extends direct logit attribution to source tokens, finding that distinct mechanisms additively combine to reinforce correct attributes. The study also identifies limitations of narrow circuit analysis and notes constraints of direct logit attribution.

Problem
It remains unclear how transformer-based LLMs mechanistically retrieve factual knowledge, especially when analyses consider only narrow parts of the input or circuit.
Method
The paper uses mechanistic interpretability and extends direct logit attribution to decompose attention-head outputs by source token, including subject and relation contributions.
Results
Factual recall comprises several distinct, independent, qualitatively different mechanisms whose additive contributions constructively interfere on the correct attribute.
Takeaways & Limitations
The additive motif suggests that comprehensively understanding model behavior requires analyzing multiple information sources and contributing mechanisms together.
Takeaways & Limitations
The analysis is conducted on a small dataset, and direct logit attribution can be limited; source-token attribution may also mislead in longer contexts.

Abstract

from arXiv · show

How do transformer-based large language models (LLMs) store and retrieve knowledge? We focus on the most basic form of this task -- factual recall, where the model is tasked with explicitly surfacing stored facts in prompts of form `Fact: The Colosseum is in the country of'. We find that the mechanistic story behind factual recall is more complex than previously thought. It comprises several distinct, independent, and qualitatively different mechanisms that additively combine, constructively interfering on the correct attribute. We term this generic phenomena the additive motif: models compute through summing up multiple independent contributions. Each mechanism's contribution may be insufficient alone, but summing results in constructive interfere on the correct answer. In addition, we extend the method of direct logit attribution to attribute an attention head's output to individual source tokens. We use this technique to unpack what we call `mixed heads' -- which are themselves a pair of two separate additive updates from different source tokens.

1. Introduction

The paper investigates how LLMs retrieve factual attributes and finds that recall is produced by several distinct mechanisms whose additive contributions constructively reinforce the correct answer.

1. Introduction: The mechanisms are qualitatively different and each independently contributes positively to the correct mean-centred logit, even when individual contributions cannot solve the task.Their output-logit distributions differ meaningfully, and the correct answer need not be the individual component’s argmax.
1. Introduction: Factual recall uses four independent mechanisms—subject heads, relation heads, mixed heads, and MLPs—that act on two information sources: the subject and relation.Each mechanism boosts the correct answer, while mixed heads combine separate subject- and relation-source contributions.
1. Introduction: Additively combining the mechanisms produces constructive interference on correct attributes and is more robust than relying on individual mechanisms alone.Most individual mechanisms are incapable of performing the task alone, whereas their sum elicits the correct answer.
1. Introduction: The analysis cautions that narrow circuit studies can leave important input sources unexplained, because factual recall depends on both relation and subject information.The paper contrasts this broader analysis with prior work that isolated only one component of a task.
1. Introduction: The paper extends direct logit attribution to decompose an attention head’s output into contributions from individual source tokens.This source-token attribution helps disentangle mixed heads into separate subject- and relation-based updates.

2. Methods

The methods define factual recall over subject–relation–attribute tuples and use direct logit attribution to trace how model components and source-token groups affect output logits.

2. Methods: The task represents factual information as tuples (s, r, a), prompts the model with the subject and relation, and evaluates whether the correct attribute appears among the top ten logits.The study primarily examines Pythia-2.8b and also checks whether similar mechanisms appear in other models.
2. Methods: The attribute sets S and R capture facts relevant to the subject and relation, respectively, with the correct attribute belonging to their intersection.These counterfactual attribute sets support analysis of subject- and relation-based output behavior.
2. Methods: The prompts are organized into PREFIX, SUBJECT, RELATION, and END token groups, with factual information required to reach the END position for answer generation.The setup uses prompts prefixed with Fact: and distinguishes source positions for later attribution.
2. Methods: Direct logit attribution converts component outputs into output-logit effects by treating the residual stream as an accumulated sum of model-component contributions.The method applies the approximately linear residual-stream-to-logit mapping to individual attention heads and MLPs.
2. Methods: DLA by source-token group further decomposes attention-head effects into weighted contributions from distinct source positions, enabling analysis of mixed heads.The attribution is most appropriate for the short prompts used here and may be misleading in longer contexts where information moves through intermediate tokens.

3. Results

Factual recall is implemented by four distinct mechanisms that form two subject- and relation-related additive update clusters, constructively combining on the correct attribute. Subject, relation, mixed heads, and MLPs contribute differently, with mixed heads combining source-token contributions.

Four mechanisms—subject heads, relation heads, mixed heads, and MLPs—form two additive update clusters tied to the subject or relation and act at END.Their contributions constructively interfere on the correct attribute, although individual mechanisms may also boost incorrect attributes.
3.1. Subject Heads: Subject heads read enriched SUBJECT representations and extract subject attributes, sometimes misfiring when the prompted relationship does not match their category.They attend primarily to SUBJECT and can extract the same attribute for a subject across different relationships.
3.2. Relation Heads: Relation heads attend primarily to RELATION and extract multiple relation attributes rather than preferentially selecting the subject-specific correct attribute.Their direct effects are largely consistent across subjects, and causal patching does not reduce average performance.
3.3. Mixed Heads: Mixed heads attend to both SUBJECT and RELATION, combining a subject-derived correct-attribute contribution with relation-derived attributes in a single head output.The paper attributes these two updates to different source positions using direct logit attribution by source.
3.4. MLPs: MLP outputs often boost many relation attributes, including the correct one, without preferentially favoring the correct attribute or depending directly on the subject.Their direct-effect category is broader than that of individual heads.

4. Discussion

The paper argues that additive mechanisms can explain factual recall and offers a partial mechanistic account of the reversal curse. Its toy example illustrates how qualitatively different contributions can jointly solve a task even when individual components are insufficient.

Additivity: An additive solution combines distinct components whose outputs differ and constructively reinforce the correct answer, even when no component solves the task alone.The paper links this behavior to additive residual-stream contributions and an approximately linear map from residual representations to logits.
Reversal Curse: Factual recall may rely on asymmetric, unidirectional mechanisms, so training on “A is B” need not strengthen the inverse “B is A” mapping.The proposed circuit enriches A and extracts B, while the reverse direction would require separate mechanisms; input and output representations are also distinct.

5. Related Work

Related work has localized factual knowledge, studied retrieval and editing, and proposed multi-step recall circuits. This paper builds on those lines by examining how individual components move and combine subject- and relation-relevant information.

Interpreting Factual Knowledge: Research on model editing has shown both that MLP layers can store or manipulate factual information and that editing success may not reliably localize the underlying fact.Subsequent work suggests editing can introduce a “loud” fact, while layer-level editing performance may be an unreliable localization signal.
Interpreting Factual Knowledge: Prior studies describe factual recall as subject enrichment, relation propagation to the END token, and selective extraction by later attention heads.The paper presents its analysis as a deeper examination of the individual mechanisms implementing this circuitry.
Extracting Knowledge from LMs: Black-box prompting evaluates what language models know and whether knowledge generalizes across paraphrases, while this work investigates mechanisms that may explain such behavior.The related work also connects factual recall to the reversal curse, where models fail to generalize from “A is B” to “B is A”.
Mechanistic Interpretability: Mechanistic interpretability combines feature analysis, mathematical frameworks, circuit discovery, component inspection, and causal interventions to reverse-engineer learned computations.This paper contributes to that broader effort by analyzing the circuitry used for factual recall.

6. Conclusion

The paper analyzes factual-recall circuitry and finds several distinct mechanisms that combine additively to produce correct answers. It presents this additive motif as a promising direction for understanding neural-network behavior while acknowledging the study’s limited scope.

6. Conclusion: Several distinct mechanisms interact additively in factual recall, constructively interfering to produce the correct answer.The paper describes each mechanism as insufficient alone, while their summed contributions are more robust.
6. Conclusion: The study expands mechanistic analysis beyond narrow circuits and uncovers factual-recall mechanisms that prior work had missed.The authors frame comprehensive study of multiple input sources as important for understanding neural networks.
6. Conclusion: The paper’s broader goal is to advance AI interpretability by shedding light on how black-box machine-learning systems function.The authors connect this understanding to the safe and beneficial development of such systems.

A. Limitations

The study provides evidence for a range of factual-recall mechanisms but does not claim to explain them all. Its analysis is centered on one model, a small dataset, and mostly high-frequency attributes, with mixed-head boundaries remaining fuzzy.

Scope and interpretation: The investigation does not explain all factual-recall mechanisms and does not quantify the importance of each identified mechanism.The Subject, Relation, and Mixed head categories are useful but not fully clean, with the Mixed boundary described as somewhat arbitrary.
Experimental scope: The experiments primarily use Pythia-2.8b and a fairly small dataset, limiting the demonstrated scope of the conclusions.The paper notes additional dataset-curation limitations in Appendix C.
Attribute frequency: Main-text plots emphasize high-unigram-frequency attributes, although the additive and constructive-interference picture also appears for less common categories with weaker individual logit-lens significance.The frequency choice simplifies analysis because polysemantic heads tend to write common attributes with higher norm.

B. Future Work

The paper identifies open questions about correlated features, MLP-neuron computation, model editing, prompting, and multi-step factual recall.

Understanding Correlation: Constructive interference may arise from computations over correlated features stored in superposition, but whether the four mechanisms suffice remains open.The authors note that linear methods may not disentangle compressed features and that more complex mechanisms may be involved.
Understanding MLP neurons: MLP analysis remains preliminary: the paper finds MLPs generally boost many relation attributes, but not how individual neurons or inputs implement this.Future work could examine individual neurons and whether relation information is used explicitly or composed with relation-head outputs.
ROME: Future work could analyze ROME end to end using the mechanisms identified for factual recall, given imprecise localization and evidence that some facts are unusually loud.The paper connects this direction to limitations reported for ROME.
Prompting Set Up: Prompting studies could test how few-shot examples, prompt injections, and paraphrases alter factual-recall performance and internal mechanisms.The paper proposes comparing mechanisms across prompting setups.
Multi-Step Factual Recall: Multi-step factual recall remains an open test of whether additivity lets models solve sequential-looking tasks in one forward pass.The paper uses prompts requiring two apparent inference steps as a motivating example.

C. Dataset

The dataset is manually constructed to isolate categorical factual recall, avoid ambiguity and copying, and support analyses holding subjects or relations constant.

C. Dataset: The dataset restricts examples to categorical facts whose correct attribute is uniquely determined by the tokenized subject and relation.It excludes ambiguous answers such as “Paris” versus “France” for “The Eiffel Tower is in,” and synonym-confounded attributes.
C. Dataset: Examples involving straightforward copying are excluded because they may rely mainly on induction heads rather than the general factual-recall mechanism.The dataset therefore avoids prompts such as “The Sydney Opera House is in the city of Sydney.”
C. Dataset: The dataset contains 106 prompts across 10 relations, and the primary model achieves high accuracy on it.The passage points to Table 2 for the relations and Figure 7 for accuracy validation.
C. Dataset: GPT-4 generated candidate attribute lists, which were then manually filtered, including removal of attributes beginning with “the.”This procedure was used to generate the subject- and relation-associated sets.

C.1. Example Datapoints

The paper provides complete examples of factual tuples together with the associated subject and relation attribute sets used for analysis.

C.1. Example Datapoints: Example datapoints instantiate each tuple as a subject, relation, attribute, and associated subject- and relation-attribute sets.The examples are presented as full (s, r, a, S, R) records.

D. Further Methods

The methods interpret factual recall through logit-based decompositions, including source-token-level attribution of attention heads, and compare head types across several models and relations.

D. Further Methods: The transformer’s output probabilities are obtained by applying softmax to logits produced from the residual stream, with LayerNorm, unembedding, attention, and MLP operations defining the computation.The notation leaves model-parameter dependence implicit and treats the residual stream as the sequence of residual activations.
D. Further Methods: The logit lens interprets intermediate residual activations by converting them into vocabulary distributions at successive layers.It relies on the residual stream’s additive structure and an approximately linear map to logits.
D. Further Methods: Direct Logit Attribution measures how individual model components directly affect output logits by decomposing accumulated residual-stream contributions.The method decomposes attention layers into heads and MLP layers into neurons.
D. Further Methods: DLA by source token further decomposes an attention head’s output according to the positions whose values contribute through attention weighting.This enables direct attribution of head effects to groups of source tokens.
D. Further Methods: The dataset analysis examines ranks of attributes from subject and relation sets, generally finding higher ranks for the strongest relation attributes.The table also highlights correlations such as France–Paris and decorrelated counterfactual attributes.
D. Further Methods: Additional model analyses compare subject, relation, and mixed attention heads for sports and country relations across GPT2-XL, GPT-J, and Pythia-6.9b.The figures use source-token-split DLA and counterfactual attributes to distinguish head types.

E.3. Relative Mechanism Importance

The paper finds that subject, relation, mixed heads, and MLPs all materially contribute to factual recall, while mixed-head behavior reflects subject–relation information propagation and model-dependent effects.

Relative mechanism importance: Subject, relation, mixed heads, and MLP layers contribute 18%, 24%, 27%, and 30% of the final mean-centred logit, respectively.The analysis omits several negatively suppressive components.
Relative mechanism importance: Edge-patching any individual mechanism significantly harms loss, supporting complementary contributions from all four mechanisms.The experiment avoids naive-ablation confounds from self-repair, known as the hydra effect.
Subject heads: Subject heads often attend strongly to SUBJECT and extract attributes independently of the requested relation, including irrelevant attributes.Some heads attend uniformly, whereas others vary their subject attention with the prompt through query composition.
Relation heads: Relation heads generally attend more to RELATION than SUBJECT, but their subject-independence and attribute selection vary across model size and relation.In Pythia-6.9B, some relation heads extract the correct attribute for different subjects, unlike the earlier Pythia-2.8B observation.
Mixed heads: Mixed heads can reflect subject information already propagated into relation tokens, producing near-top relation-token attribution and top SUBJECT attribution.This propagation is isolated through attention knockout and is especially relevant when a head’s specialization overlaps the relation category.
MLPs: MLPs on END mostly lack significant subject-dependent indirect effects, although some relations show such effects that the paper does not explain.The analysis therefore focuses on the MLPs’ direct effects for those cases.

E.8. Category Identification

The paper examines how attention-head output categories relate to factual-recall mechanisms and finds that heads are polysemantic, with overlapping and sometimes misleading specializations. This motivates interpreting head categories rather than equating them directly with subject or relation attribute sets.

Category identification: Heads are polysemantic: their categories are not fully aligned with the subject- or relation-relevant attribute sets S and R.The category summary uses 10,000 randomly selected CounterFact prompts and top DLA tokens from subject, relation, and mixed heads.
Category identification: L13H31 responds to sports and countries, plausibly because its specialization in locations, positions, and places overlaps both categories.The example illustrates that lexical or semantic category overlap can produce cross-relation head responses.
Category identification: L18H25 misfires on factual recall because transport and consumables overlap with sports words such as Golf, swimming, and track.It is nevertheless the eighth most important mixed head across the two studied relations by DLA.
Category identification: The paper frames category analysis as relevant to understanding attention-head superposition and whether heads are the correct fundamental units of study.A theoretical example shows multiple contextual computations compressed into fewer heads through additive combinations.