Locating and Editing Factual Associations in GPT

Kevin Meng; David Bau; Alex Andonian; Yonatan Belinkov

Source-linked AI summary

Locating and Editing Factual Associations in GPT

Kevin Meng, David Bau, Alex Andonian, Yonatan Belinkov

arXiv:2202.05262v5Feb 10, 2022cs.CLcs.LG

TL;DR

The paper asks how factual associations are stored in autoregressive transformers. It uses causal tracing to locate decisive computations and ROME to edit their feed-forward weights. The results support localized factual storage in midlayer MLPs and show that ROME can preserve generalization and specificity together on counterfactual assertions.

Problem
The paper investigates how factual associations are stored in GPT-like autoregressive transformers, an under-explored question relative to masked models.
Method
The paper combines causal tracing of hidden-state effects with ROME, which directly updates feed-forward weights associated with factual recall.
Results
ROME achieves good generalization and specificity simultaneously on counterfactual assertions, while previous methods sacrifice one or the other.
Takeaways & Limitations
The experiments provide evidence that factual associations are localized in midlayer MLP computations and demonstrate the feasibility of directly manipulating computational mechanisms.
Takeaways & Limitations
ROME edits only one fact at a time, and the study does not investigate logical, spatial, or numerical knowledge.

Abstract

from arXiv · show

We analyze the storage and recall of factual associations in autoregressive transformer language models, finding evidence that these associations correspond to localized, directly-editable computations. We first develop a causal intervention for identifying neuron activations that are decisive in a model's factual predictions. This reveals a distinct set of steps in middle-layer feed-forward modules that mediate factual predictions while processing subject tokens. To test our hypothesis that these computations correspond to factual association recall, we modify feed-forward weights to update specific factual associations using Rank-One Model Editing (ROME). We find that ROME is effective on a standard zero-shot relation extraction (zsRE) model-editing task, comparable to existing methods. To perform a more sensitive evaluation, we also evaluate ROME on a new dataset of counterfactual assertions, on which it simultaneously maintains both specificity and generalization, whereas other methods sacrifice one or another. Our results confirm an important role for mid-layer feed-forward modules in storing factual associations and suggest that direct manipulation of computational mechanisms may be a feasible approach for model editing. The code, dataset, visualizations, and an interactive demo notebook are available at https://rome.baulab.info/

1 Introduction

The paper investigates how GPT-like autoregressive transformers store factual associations and finds evidence that these associations correspond to localized computations that can be directly edited.

The paper studies how factual associations are stored within GPT-like autoregressive transformer models, whose knowledge storage remains under-explored.
Causal mediation analysis identifies feed-forward MLPs at middle layers as decisive for recalling facts while processing the last subject token.
ROME alters parameters in a decisive feed-forward layer to update specific factual associations.
ROME is similarly effective to existing model-editing approaches on a standard zero-shot relation extraction benchmark.
On counterfactual assertions, ROME maintains generalization and specificity simultaneously, whereas previous methods sacrifice one or the other.

2 Interventions on Activations for Tracing Information Flow

The paper traces factual-prediction information flow by corrupting subject representations and restoring selected hidden states, identifying middle-layer MLP computations as especially causal during subject processing.

Facts are represented as tuples containing a subject, relation, and object, and elicited by prompting GPT with the subject and relation before examining its object prediction.
The transformer computation is organized as hidden states across layers and token positions, with attention and MLP contributions computed from previous states.
Causal Tracing compares clean, corrupted, and corrupted-with-restoration runs to measure whether restoring a hidden state recovers the original factual prediction.
Causal Tracing Results: 18.6% ATE and 8.7% AIE at layer 15 show substantial causal effects, including a newly identified early site at the last subject token.
Causal Tracing Results: 6.6% MLP AIE versus 1.6% attention AIE at the last subject token suggests a decisive early-site role for MLP modules.
Causal Tracing Results: Severing MLP computations removes causal effects from lower layers, while severing attention shows no analogous transition, confirming an essential role for middle-layer MLPs.
The proposed mechanism has midlayer MLPs process subject inputs to recall memorized properties, which high-layer attention copies toward the final token.

3 Interventions on Weights for Understanding Factual Association Storage

ROME tests whether factual associations are stored in mid-layer MLPs as editable key–value computations, using targeted weight changes and evaluations designed to separate generalization from specificity.

Interventions on Weights: ROME models MLP projections as linear associative memories that map subject-derived keys to stored factual values.The method treats the feed-forward projection as a key–value store and modifies it with a rank-one update.
Interventions on Weights: The rank-one update inserts a new key–value pair directly by constraining the targeted matrix to map the subject key to a new value.The covariance-based closed-form update uses the original matrix, key covariance, and residual error for the new pair.
Targeting Decisive States: ROME selects the MLP input at the final subject token as the lookup key, following causal evidence that this location is decisive for factual recall.Causal tracing found strong effects at the last subject token, while ROME edits are most successful there and in middle layers.
Evaluation: On zsRE, ROME is competitive with existing editing methods, while COUNTERFACT provides more sensitive tests of efficacy, paraphrase generalization, and neighborhood specificity.COUNTERFACT contains 21,919 records and distinguishes robust factual changes from superficial target-word regurgitation.
Evaluation: On COUNTERFACT, ROME achieves generalization and specificity simultaneously, whereas other methods exhibit overfitting, underfitting, poor specificity, or ineffective edits.The comparison covers GPT-2 XL and GPT-J test sets and identifies failure patterns including bleedover and regurgitation.
Limitations: ROME edits are intended for mechanistic understanding rather than large-scale training, and each edit changes one directional factual association at a time.Editing reversed subject–relation directions requires separate edits, and the paper does not investigate logical, spatial, or numerical knowledge.

4 Related Work

Prior work studies model knowledge through probing, causal interventions, extraction benchmarks, and localization or editing methods, with limitations across these approaches.

Probing classifiers assess whether properties are encoded in representations but can be dissociated from the network’s behavior.
Causal mediation analysis identifies hidden states or neurons that contribute to model behaviors, including gender assumptions and syntactic agreement.
Knowledge extraction commonly uses fill-in-the-blank prompts, with later improvements from prompt diversification and fine-tuning on open-domain facts.
Localization studies characterize transformer MLPs as key–value memories and edit BERT facts by writing object embeddings into selected MLP rows.

5 Conclusion

The paper clarifies information flow during factual recall and uses that understanding to develop ROME, a model editor for modifying stored knowledge. Experiments support direct manipulation of computational mechanisms in large pretrained models, while the methods edit one fact at a time.

ROME is presented as a simple, principled model editor grounded in the paper’s analysis of factual knowledge storage.
The experiments provide insight into how facts are stored and demonstrate the feasibility of directly manipulating computational mechanisms in large pretrained models.
The methods apply only to editing a single fact at once.
Code, notebooks, datasets, benchmarks, and visualizations are open-sourced online.

6 Ethical Considerations

The paper discusses transparency and error correction as potential benefits of directly editing large models, while warning that the capability could also enable misuse. It therefore cautions against using language models as authoritative factual sources in critical settings.

Direct editing may improve transparency of large autoregressive transformers and reduce energy consumed to correct their errors.
The ability to edit large models could be abused to add malicious misinformation, bias, or adversarial data.
The authors stress that large language models should not be used as authoritative factual knowledge sources in critical settings.

Appendices

The appendix derives ROME’s update through least-squares optimization with an equality constraint, using a rank-one update formulation. It states the assumptions and linear-system steps used to solve for the edited matrix and auxiliary variables.

The appendix frames the result as classical least-squares with equality constraints combined with a rank-one update rule.
The derivation assumes W is the optimal least-squares solution mapping previous keys K to values V.
The Frobenius norm expresses total squared error because the optimized variable is a matrix W rather than a vector.
The edited matrix Ŵ is required to solve the same least-squares problem subject to an additional equality constraint.
The constrained optimization is formulated with a Lagrangian, where Λ minimizes the stated objective.
Subtracting the original normal-equation solution from the constrained solution yields the update rule, using C = KK^T and symmetry of C.
Equations 14 and 7 form a block linear system that allows Ŵ and Λ to be solved simultaneously.

B.1 Experimental Settings

The experiments use causal tracing to identify hidden states that mediate factual predictions, testing robustness across models, layers, tokens, and corruption choices. Effects concentrate in middle-layer MLPs processing the subject’s last token, while causal tracing provides insights that integrated gradients do not.

Causal tracing setup: Causal tracing compares clean, corrupted, and corrupted-with-restoration runs to measure how individual hidden states contribute to factual predictions.Subject embeddings are corrupted with Gaussian noise, then selected hidden states are restored while subsequent computation proceeds normally.
Causal tracing results: 23.6% was the average maximum restored score for runs of MLP lookups, with the strongest effect centered on the last subject token around layer 17.The experiment restores ten consecutive MLP values over a layer interval and finds the largest average effect at the last entity token.
Cross-model robustness: Across GPT-2, GPT-NeoX, GPT-J, and smaller GPT-2 models, localized causal effects recur despite differences in model size and architecture.Early-site MLP modules retain high indirect causal effects at the last subject token, although peak layers vary between models.
Corruption choices: Noise must be sufficiently large for causal tracing: smaller variance makes clean and corrupted runs too similar, while alternative distributions yield smaller effects.Multivariate Gaussian and uniform noise still produce measurable causal traces, but effects are smaller than with the main 3σt spherical Gaussian.
Comparison with salience methods: Integrated gradients produce scattered saliency maps and do not reveal the last subject token or midlayer MLP role identified by causal tracing.Causal tracing is also reported to remain informative under different noise configurations.

C Details on the zsRE Evaluation Task

The zsRE evaluation uses factual statements with paraphrase and neighborhood prompts to assess model-editing performance. Its established splits and task structure support comparisons among editing methods.

Dataset structure: zsRE contains factual statements, paraphrase prompts, and neighborhood prompts, using the same train/test splits as Mitchell et al. (2021).Nonhypernetwork methods such as ROME do not require training, so the corresponding training split is discarded for those methods.
Baseline evaluation: KE-zsRE and MEND-zsRE are custom-trained on the zsRE training split to provide fair comparisons with ROME and other nonhypernetwork methods.The original KE and MEND hypernetworks were trained on WikiText rather than zsRE.

D Details on the COUNTERFACT Dataset

COUNTERFACT evaluates whether edits change underlying factual knowledge rather than only superficial word choices, while preserving generalization and specificity. The dataset combines counterfactual rewrites with paraphrase, neighborhood, generation, and reference-text tests.

Dataset purpose: COUNTERFACT distinguishes superficial word-choice changes from specific and generalized changes in underlying factual knowledge.Its records are designed to support more sensitive evaluation of knowledge editing.
Record composition: Each record contains a subject, relation, original and requested objects, one rewriting prompt, paraphrases, neighborhood prompts, generation prompts, and reference texts.The rewrite and paraphrase prompts derive from relation-specific templates, while neighborhood entities share the subject’s predicate.
Evaluation dimensions: COUNTERFACT’s evaluation tests whether an edit changes the requested association, generalizes across paraphrases and generations, and remains specific to neighboring facts.The benchmark provides five rewrite-success criteria through its combined prompt and reference-text structure.
Baselines: The editing baselines include unconstrained and constrained fine-tuning, knowledge-neuron editing, KE, and MEND, with hyperparameters selected through searches or adopted defaults.Fine-tuning changes localized MLP weights, while KE and MEND use learned rank-1 updates.
ROME intervention: ROME intervenes at layer 18, the center of causal effect in MLP layers and the approximate transition from MLP keys to values.The update uses key selection, value optimization, and value insertion.
Model-scale evaluation: ROME outperforms the next-best baseline on GPT-2 Medium and GPT-2 Large evaluations, as measured against GPT-2 XL’s FT+L baseline.The evaluations cover both COUNTERFACT and zsRE.

G.1 GPT-2 XL (1.5B) Generation Examples

Generation examples compare how editing methods affect generalization, fluency, subject essence, and compositionality on GPT-2 XL and GPT-J. ROME often generalizes strongly while preserving essence and fluency, though ambiguous relations and occasional hallucinations remain difficult.

Garth Knox: ROME produces a straightforward and correct generation for Garth Knox’s birthplace, whereas MEND, KE, and FT+L fail to generalize well.FT suggests an indirect parental connection rather than explicitly stating the edited birthplace.
Ambiguous relations: The ambiguous “pastoral” relation remains difficult: GPT-2 XL edits focus on landscapes rather than pastoral music, and ROME can repeat the word.The prompts do not disambiguate the intended musical meaning.
GPT-J examples: On GPT-J, ROME generalizes deeply while maintaining subject essence and fluency across the displayed examples.For Sonic Drift 2, ROME retains that the subject is a game after changing its creator to Microsoft.
Liberty Island: ROME and FT generalize the Liberty Island rewrite to Scotland across direct and indirect prompts, while other methods are inconsistent or fail.ROME’s generation connects Liberty Island’s edited location with Loch Lomond, demonstrating compositional use of the new association.
Remaining failures: ROME can still produce occasional generation errors, including a non-German hallucinated name in the Garth Knox example.The GPT-J example otherwise reports that FT+L and ROME work well.

I Are Attention Weight Interventions Effective?

The section tests whether editing attention weights can alter factual generation, comparing AttnEdit with ROME. ROME generalizes to paraphrases and prompts where AttnEdit fails, supporting a distinction between factual recall and word prediction.

Interpretation: The results support the hypothesis that middle-layer MLPs mediate factual recall while late-layer attention predicts specific word sequences.The hypothesis assigns factual recall to middle-layer MLPs processing subject tokens and word prediction to late-layer attention modules.
Intervention design: The experiment contrasts ROME's feed-forward intervention with AttnEdit's constrained attention-weight fine-tuning.ROME implements the MLP operation, while AttnEdit modifies W Q_i weights across heads at a selected layer.
Generation results: Both methods reproduce the memorized fact under the original rewriting prompt.The generation comparison reports success for both ROME and AttnEdit on the original prompt.
Generation results: ROME succeeds on paraphrases and generalization prompts, whereas AttnEdit fails on those inputs.This qualitative comparison supports the distinction between factual recall and word prediction.

J Human Evaluation

The human evaluation compares text generated by unmodified GPT-2 XL, ROME, and FT+L after counterfactual editing. ROME produces more counterfactual-consistent text than FT+L, but human raters detect reduced fluency.

Study design: Fifteen volunteers evaluated 50 counterfactual text samples, producing 150 consistency rankings.The study compared unmodified GPT-2 XL, ROME-modified text, and FT+L-modified text.
Study conditions: The evaluation used unpaid remote volunteers who worked for less than 30 minutes and could opt out at any time.These were the stated participation conditions for the human study.
Consistency: ROME is much more successful than FT+L at generating text consistent with the counterfactual.The finding aligns with Table 4, where ROME generalizes better than FT+L.
Fluency: Human raters detect reduced fluency under ROME relative to the comparison baselines.The entropy measure does not discern this fluency reduction, which appears in examples from Figures 27 and 28.