PaperlayerView exact revision on arXiv

Source-linked AI summary

Dissecting Recall of Factual Associations in Auto-Regressive Language Models

Mor Geva, Jasmijn Bastings, Katja Filippova, Amir Globerson

arXiv:2304.14767v3cs.CL

TL;DR

The paper investigates how factual associations stored in autoregressive language models are retrieved during inference. Using information-flow analysis and attention-edge interventions, it identifies a three-step mechanism: early MLP-driven subject enrichment, relation propagation, and attribute extraction, typically by upper attention heads. The findings characterize how factual associations are stored and extracted and motivate research on knowledge localization and model editing.

  • Problem

    It remains unclear how factual knowledge encoded in language-model parameters is extracted during inference.

  • Method

    The authors analyze correctly answered subject-relation queries using information-flow tracing, attention knockouts, representation projections, and MHSA/MLP interventions.

  • Results

    The identified mechanism enriches subject representations in early layers, passes relation information to the prediction, and extracts attributes through upper MHSA attention heads in ∼70% of predictions.

  • Takeaways & Limitations

    Factual associations are stored in lower network layers as attribute-rich subject representations and are extracted prominently through MHSA heads that often encode subject-attribute mappings.

  • Takeaways & Limitations

    Vocabulary projection only approximates information encoded in intermediate representations, especially in early layers.

Abstract

from arXiv · show

Transformer-based language models (LMs) are known to capture factual knowledge in their parameters. While previous work looked into where factual associations are stored, only little is known about how they are retrieved internally during inference. We investigate this question through the lens of information flow. Given a subject-relation query, we study how the model aggregates information about the subject and relation to predict the correct attribute. With interventions on attention edges, we first identify two critical points where information propagates to the prediction: one from the relation positions followed by another from the subject positions. Next, by analyzing the information at these points, we unveil a three-step internal mechanism for attribute extraction. First, the representation at the last-subject position goes through an enrichment process, driven by the early MLP sublayers, to encode many subject-related attributes. Second, information from the relation propagates to the prediction. Third, the prediction representation "queries" the enriched subject to extract the attribute. Perhaps surprisingly, this extraction is typically done via attention heads, which often encode subject-attribute mappings in their parameters. Overall, our findings introduce a comprehensive view of how factual associations are stored and extracted internally in LMs, facilitating future research on knowledge localization and editing.

1 Introduction

The paper asks how factual knowledge encoded in language-model parameters is extracted during inference. It identifies a three-step mechanism in which subject information is enriched, relation information propagates, and attention extracts the attribute.

  • The study addresses how factual knowledge is extracted from model parameters during inference, a question left unclear by prior localization work.
  • The authors analyze factual queries where a subject and relation prompt the model to predict the correct attribute.They restrict analysis to cases where the model predicts the correct next-token attribute.
  • Attention-edge interventions identify two consecutive critical information-flow points: relation information reaches the prediction before subject information.
  • The observed extraction behavior occurs in ∼70% of predictions.
  • Subject information is enriched at the last subject token across early layers, forming a representation that encodes many subject-related attributes.The enrichment is driven by early MLP sublayers.
  • The last token uses relation information to extract the corresponding attribute from the enriched subject representation via attention-head parameters.The paper reports that attention heads often encode subject-attribute mappings.

2 Background and Notation

The paper formalizes autoregressive decoder-only transformers as layered token representations updated by global MHSA and local MLP sublayers. MHSA can be decomposed into per-head attention-weighted updates, while MLPs can store factual knowledge as key-value memories.

  • An input sequence of N tokens is embedded into vectors and transformed through L layers, each containing MHSA and MLP sublayers.
  • Each MHSA sublayer performs global updates by aggregating information from representations at the previous layer.
  • The MHSA output can be represented as a sum of matrices, each induced by one attention head.
  • For head j at layer ℓ, Aℓ,j stores attention weights, Mℓ,j masks them, and autoregressive masking restricts attention to preceding positions.
  • Each MLP sublayer computes a local update for each token representation.
  • Prior work casts transformer MLP sublayers as key-value memories that store factual knowledge.

3 Experimental Setup

The experiments study factual open-domain queries in which models predict an attribute for a subject-relation pair. They use correctly answered COUNTERFACT queries and analyze GPT-2 and GPT-J.

  • The task requires predicting an attribute a from a factual subject-relation pair (s, r, a).Queries appear in question-answering or fill-in-the-blank formats.
  • The study focuses on how attributes are extracted internally, despite language models often predicting the correct attribute.
  • Queries are represented as token sequences, with subject tokens marking the subject and non-subject tokens expressing the relation.
  • The dataset is sampled from COUNTERFACT, retaining queries for which each model predicts the correct attribute.The predicted attribute may be a sub-word prefix of the attribute name.
  • The models are GPT-2 with 48 layers and 1.5B parameters and GPT-J with 28 layers and 6B parameters.Both use vocabularies of approximately 50K tokens; GPT-J uses parallel MHSA and MLP sublayers.
  • The samples contain 1,209 queries for GPT-2 and 1,199 for GPT-J.

4 Overview: Experiments & Findings

The paper combines attention blocking with analyses of subject representations and extraction to trace factual prediction formation. Its overview identifies relation-then-subject information flow, early MLP-driven subject enrichment, and upper-MHSA extraction.

  • Attention blocking identifies critical information-flow points by measuring prediction effects when selected computation paths are knocked out.
  • Critical information propagates to the prediction first from relation positions and then from subject positions.
  • Early MLP sublayers enrich the last subject-position representation with many subject-related attributes.
  • Upper MHSA sublayers typically extract the relevant attribute through attention heads that often encode subject-attribute mappings in their parameters.
  • Figure 2 measures relative prediction-probability changes after interventions on attention edges to the last position, using 9-layer GPT-2 and 5-layer GPT-J windows.

5 Localizing Information Flow via Attention Knockout

The authors use attention knockout to localize where critical information reaches the prediction. Information from relation positions propagates first, followed by information from subject positions, at distinct stages.

  • Method: Attention Knockout: Attention knockout blocks selected attention edges to test whether critical information propagates between input positions during inference.The intervention targets MHSA because it is the module that communicates information between positions.
  • Results: Blocking attention from the last position to subject tokens causes prediction-probability decreases of up to 60% in middle-upper layers.This identifies direct propagation from subject positions to the prediction at those layers.
  • Results: Critical information from non-subject positions reaches the last position before critical information from subject positions.This ordering holds across different subject-relation orders.
  • Results: The interventions reveal disjointed computational stages with peaks of critical information propagating from different positions to the prediction.The next analysis examines the critical information transferred from subject positions.

6 Intermediate Subject Representations

Intermediate representations at the last subject position become rich in subject-related attributes across layers. Early MLP sublayers contribute most to this enrichment, while upper attention sublayers later support extraction from the enriched representation.

  • 6.1 Inspection of Subject Representations: Attributes rate measures the portion of top-ranked tokens that appear in a subject-specific candidate-attribute set built from Wikipedia contexts.The candidate set is formed by retrieving subject paragraphs with BM25, tokenizing them, and removing common words and sub-words.
  • 6.1 Inspection of Subject Representations: The experiment compares the last-subject position with the first subject position, the position after the subject, and the prediction position.The last-subject position is examined because it attends to all subject positions.
  • 6.1 Inspection of Subject Representations: The last-subject representation’s attributes rate increases across layers and reaches close to 50% in intermediate-upper layers for GPT-2 and GPT-J.Its rate is substantially higher than at other examined positions in that layer range.
  • 6.2 Static Subject Embeddings: The highest attributes rate across individual subject-token embeddings averages 19.3 in GPT-2 and 28.6 in GPT-J, while mean subject embeddings average 4.1 and 11.5.These rates are lower than those of higher-layer subject representations, indicating that other components contribute to attribute extraction.
  • 6.3 Subject Representation Enrichment: Canceling early MLP updates decreases the attributes rate by ∼88% on average in GPT-2, whereas canceling MHSA updates decreases it by less than 30%.Similar trends appear for other layers and GPT-J, and MLP updates promote subject-related concepts.
  • 6.3 Subject Representation Enrichment: The findings place subject-representation enrichment primarily in early MLP sublayers, extending the view of factual associations as distributed across early MLPs.This is consistent with interpreting MLP sublayers as key-value memories.

7 Attribute Extraction via Attention

Attribute extraction is typically performed by upper MHSA sublayers that use information from the enriched subject representation and the last position. Interventions show that subject attention is critical, while MLPs contribute less frequently.

  • Attribute extraction mechanism: Upper MHSA sublayers promote the attribute to the prediction, with 68.2% of GPT-2 examples showing extraction events.Extraction layers coincide with those where critical subject information propagates to the last position.
  • Intervention analysis: Blocking attention to subject or non-subject positions suppresses extraction to 30.2% or 31.5%, respectively, while allowing self- and subject-attention raises it to 44.4%.These interventions suggest that critical information is centered on the last position and subject representation.
  • Attribute extraction mechanism: MLP sublayers show a 31.3% extraction rate, substantially below MHSA, indicating that both contribute but MHSA is the prominent mechanism.MLP extraction follows MHSA extraction in 17.4% of examples and occurs without MHSA extraction in another 10.2%.
  • Attribute extraction mechanism: MHSA extraction is non-trivial: the attribute has average rank 999.5 in the subject representation before the MHSA update promotes it to rank 1.This rules out simple propagation of an already top-ranked attribute.
  • Importance of subject enrichment: Patching early subject representations decreases extraction by up to 50%, whereas patching non-subject representations has a weaker effect.The result stresses the importance of subject enrichment for attribute recall and suggests non-subject positions become ready early.
  • Knowledge attention heads: Attention-head parameters encode the subject-attribute mapping in 30.2% of GPT-2 and 39.3% of GPT-J extraction events.The mappings span 150 GPT-2 heads, mostly in upper layers 24–45, with frequent heads encoding hundreds of mappings.

8 Related Work

Related work has localized factual information and studied prediction processes, internal representations, and mechanistic interventions. This paper is motivated by the remaining uncertainty about how factual predictions are constructed during recall.

  • Knowledge tracing and editing: Prior knowledge-tracing work located factual information in layers and neurons, often informing knowledge-editing approaches.Editing studies include methods for changing or modifying encoded knowledge.
  • Knowledge tracing and editing: Evidence that encoded facts can change after editing locations different from those identified by localization methods highlights limited understanding of factual prediction construction.This discrepancy motivates the paper’s focus on the recall process of factual associations.
  • Prediction and representation studies: Related prediction studies identify early MLP sublayers as crucial for memorized predictions and propose editing knowledge encoded in hidden representations.The paper also connects to broader studies of how language models organize information internally.
  • Mechanistic interpretability: Mechanistic-interpretability work uses vocabulary projections and interventions in transformer computation to study language-model inner workings.A concurrent study zeroed MHSA values, producing an effect equivalent to the blocking intervention used here.

9 Conclusion

The paper analyzes factual-association recall as an internal attribute-extraction process in autoregressive transformers. It identifies enrichment in lower layers, followed by prominent MHSA-based extraction using attention heads that often encode mappings.

  • Conclusion: The analysis unveils a core attribute-extraction mechanism for factual associations in autoregressive transformer-based language models.The mechanism is studied as part of the models’ internal recall process.
  • Conclusion: Factual associations are stored in lower network layers and extracted prominently by MLP sublayers to form attribute-rich subject representations.These representations contain many subject-related associations before the final extraction step.
  • Conclusion: Given a subject-relation query, upper MHSA sublayers prominently extract the correct attribute from the enriched subject representation.The relevant attention heads often encode subject-attribute mappings in their parameters.

Limitations

The appendix identifies interpretive and intervention-related limitations while testing whether information-flow findings remain consistent across settings and window sizes. It also reports positional effects and supports the main subject-enrichment observations with gradient-based analyses.

  • Interpretation limits: Intermediate-representation projections to vocabulary space approximate encoded information, especially in early layers.The authors note that the increasing attributes rate is unlikely to be explained solely by this approximation.
  • Intervention limits: Attention knockouts cannot fully prevent information leakage across positions from earlier layers.Blocking ranges of layers alleviates, but does not eliminate, this issue; the critical edges remain relevant for the study’s primary goal.
  • Subject-relation order: Subject information reaches the last position across subject-relation orders, but relation-position interventions have greater and more distributed effects when the relation precedes the subject.The reported pattern is observed for both subject orders, while the first-position effect complicates interpretation.
  • Information localization: Blocking the last subject position damages predictions by 50%–100%, whereas leaving it unblocked usually preserves prediction probability.The analysis attributes critical subject information primarily to the last subject position in GPT-2.
  • Robustness: The two critical information-flow sites remain across window sizes, with relation information appearing in early layers and subject information in upper layers.A single-layer GPT-2 relation knockout is an exception, suggesting relation information may be processed across multiple layers.
  • Supporting analysis: Gradient-times-input analysis finds subject positions relevant deep into computation, while relation inputs matter mainly in the first few layers.This supports the reported distinction between subject enrichment and earlier relation processing.

C.2 Additional Sublayer Knockout Results

Additional knockout experiments compare MLP and MHSA contributions to subject-representation enrichment. Across successive layers, canceling MLP updates substantially reduces attribute evidence, whereas canceling MHSA updates has a smaller effect.

  • Layerwise comparison: Across successive layers, canceling MLP updates reduces attributes rate dramatically, while canceling MHSA updates has a more benign effect.The additional analysis evaluates the effect at each successive layer rather than only at one upper layer.
  • MHSA knockouts: Canceling MHSA sublayers causes a much smaller attributes-rate decrease of <30%.This contrasts with the destructive effect of canceling early MLP updates.

E Additional Results for GPT-J

Additional GPT-J analyses reproduce the paper’s main patterns and extend them with layerwise extraction, representation-patching, and attention-parameter analyses. They also examine examples and subject-enrichment updates in greater detail.

  • Additional analyses: GPT-J appendix analyses measure MHSA and MLP effects on attributes rate, extraction rates across layers, and representation replacement at different positions.Figures 17–19 provide these additional analyses.
  • Examples and statistics: The appendix includes per-example extraction statistics, subject-representation token projections, and intervention examples across GPT-2 and GPT-J.These materials cover named subjects and factual queries used to illustrate the analyses.
  • Overall consistency: GPT-J results are consistent with the GPT-2 findings described throughout the paper.The appendix summarizes this consistency across the additional analyses.
  • MLP sub-updates: Manual inspection identifies subject-related concepts among dominant early-MLP sub-updates, while quantifying this process remains non-trivial.The sub-updates are examined through projections of selected parameter vectors to the vocabulary.
  • Attention mappings: Specific attention heads often encode mappings between input subjects and predicted attributes in their parameters.The appendix provides examples of these mappings and examines heads that repeatedly perform extraction.
  • Knowledge hubs: Heads extracting attributes for at least 10% of queries act as knowledge hubs whose top mappings encode hundreds of factual associations across relations.The analysis focuses on seven such GPT-2 heads.
Loading 2304.14767v3…