PaperlayerView exact revision on arXiv

Source-linked AI summary

Mass-Editing Memory in a Transformer

Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, David Bau

arXiv:2210.07229v2cs.CLcs.LG

TL;DR

Large language models can lack specialized or current knowledge, while existing editing methods have been limited in the number of associations they can update. MEMIT directly inserts many memories through explicit multi-layer parameter updates, and experiments show scaling to thousands of memories while maintaining key editing properties. The paper also notes scope limitations and warns that editing can be misused to insert false or damaging information.

  • Problem

    Large language models may lack specialized or current knowledge, while prior knowledge-editing work has mainly addressed single associations or struggled to scale to thousands of edits.

  • Method

    MEMIT uses explicitly calculated updates across transformer layers identified as causal mediators of factual recall to insert multiple memories.

  • Results

    MEMIT scales to much larger edit sets than prior approaches while maintaining generalization, specificity, and fluency across experiments on GPT-J and GPT-NeoX.

  • Takeaways & Limitations

    The results suggest that explicit analysis of internal computations can support more transparent and practical large-scale model updates.

  • Takeaways & Limitations

    The studied knowledge representation is limited to directional subject–relation–object relations and does not cover spatial or temporal reasoning, mathematical knowledge, or linguistic knowledge.

Abstract

from arXiv · show

Recent work has shown exciting promise in updating large language models with new memories, so as to replace obsolete information or add specialized knowledge. However, this line of work is predominantly limited to updating single associations. We develop MEMIT, a method for directly updating a language model with many memories, demonstrating experimentally that it can scale up to thousands of associations for GPT-J (6B) and GPT-NeoX (20B), exceeding prior work by orders of magnitude. Our code and data are at https://memit.baulab.info.

1 INTRODUCTION

MEMIT addresses the need to update specialized or obsolete knowledge directly, rather than retraining large language models. It introduces a scalable approach for inserting many memories while preserving generalization, specificity, and fluency.

  • Large language models may lack specialized knowledge or retain obsolete facts, motivating direct knowledge updates for applications such as search and content generation.Retraining a large model can be prohibitive.
  • MEMIT treats language-model knowledge as memorized subject–relation–object tuples that can be edited in the model’s weights.The example changes Michael Jordan’s sport from basketball to baseball.
  • MEMIT is designed to update hundreds or thousands of facts simultaneously, whereas naive sequential application of existing editing methods does not scale.
  • Experiments evaluate whether MEMIT maintains generalization, specificity, and fluency while scaling across GPT-J and GPT-NeoX.The comparisons include rank-one, hypernetwork, and fine-tuning baselines.

2 RELATED WORK

Prior work treats knowledge bases as explicit, queryable structures or as implicit knowledge in language models, while editing methods have struggled to scale. MEMIT builds on efforts to localize factual knowledge within transformer computations.

  • Structured knowledge bases can be precisely queried, measured, and updated, but their coverage of uncatalogued knowledge is limited.
  • Language models offer broad, schema-free knowledge access, but their knowledge is implicit, prompt-sensitive, and difficult to catalog, add, or update.
  • Hypernetwork knowledge editors generate weight updates for new facts, but prior studies report that they fail to scale beyond a few edits or to 10 beliefs.
  • Direct model-editing work localizes factual recall to transformer mechanisms, including MLP layers, and uses causal mediation analysis to identify relevant layers.MEMIT narrows its focus to these mechanisms.

3 PRELIMINARIES: LANGUAGE MODELING AND MEMORY EDITING

The paper models language-model memories as directional subject–relation–object facts elicited by prompts and evaluates edits through efficacy, generalization, specificity, and fluency. The central challenge is achieving these properties across thousands of edits without conflicting requests.

  • Language modeling: A transformer decoder generates text by sampling conditional token distributions from hidden representations computed recursively across layers.The model is parameterized by a D-layer transformer decoder.
  • Memory representation: The studied memories are directional tuples (s, r, o), recalled by prompting a subject–relation pair and predicting the object tokens.An example prompt asks what sport Michael Jordan plays.
  • Memory representation: The edit-request set excludes conflicting requests that assign different objects to the same subject and relation.For example, one subject cannot be assigned both baseball and professional soccer for the same relation.
  • Evaluating edits: A good edit must increase the target fact’s probability while generalizing across paraphrases, preserving unrelated facts, and maintaining fluent generation.
  • Evaluating edits: Achieving effective, generalized, specific, and fluent edits remains challenging even for a few edits, motivating evaluation at the scale of thousands.

4 METHOD

MEMIT identifies the MLP layers that causally mediate factual recall and distributes explicitly calculated updates across them to insert many memories. Its procedure computes target hidden vectors, spreads residual changes across critical layers, and applies batch updates while accounting for downstream activation changes.

  • Identifying the critical path of MLP layers: Early attention gathers subject information at the last subject token, after which critical MLPs add memories to the residual stream for later attention to read.The figure presents this as the causal path MEMIT edits.
  • Identifying the critical path of MLP layers: MEMIT targets the range of MLP layers identified as causal mediators of factual recall rather than editing a single layer.For GPT-J, the critical range is R = {3, 4, 5, 6, 7, 8}.
  • Batch update for a single linear associative memory: MEMIT formulates each layer as a linear associative memory that maps input keys encoding memorized properties to desired memory vectors with minimal squared error.The batch objective preserves previously stored associations while inserting a large batch of new memories.
  • Updating multiple layers: MEMIT spreads each memory’s residual change approximately evenly across the critical layers, updating them iteratively in ascending order and recollecting downstream activations.Iterative updates account for the fact that changing one layer affects activations in subsequent layers.
  • Updating multiple layers: For each requested edit, MEMIT computes a target vector z_i that would encode the new association at the target layer and token.The target vector is optimized to maximize prediction of the desired object across prompts with random prefixes.
  • Updating multiple layers: The resulting algorithm takes requested edits, a generator, target layers, and covariance statistics, then returns a modified generator containing the edits.Layer updates use covariance-informed batch solutions and update the MLP weights directly.

5 EXPERIMENTS

MEMIT is evaluated on GPT-J and GPT-NeoX against fine-tuning, MEND, and sequential ROME across large-scale factual and counterfactual editing tasks. The experiments examine efficacy, generalization, specificity, fluency, scaling, category effects, mixed edits, and runtime.

  • Experimental setup: MEMIT is evaluated on GPT-J (6B) and GPT-NeoX (20B) against FT-W, MEND, and sequential ROME.The experiments compare direct multi-fact editing with fine-tuning, a hypernetwork editor, and sequential single-fact editing.
  • Editing 10K memories in zsRE: 10,000 zsRE edits test whether MEMIT can add correct information while preserving paraphrase generalization and specificity.Efficacy measures recall of the new object, Paraphrase measures performance on rephrasings, and Specificity measures unrelated facts that should remain unchanged.
  • Editing different categories of facts: MEMIT achieves higher overall scores than FT and MEND across all 27 tested fact categories, including harder relations such as athlete sport edits.The category analysis uses relations with at least 300 COUNTERFACT cases each.
  • Editing different categories of facts: MEMIT generally improves both specificity and generalization, although trade-offs remain for relations including P127 and P641.MEND shows a clear generalization-specificity trade-off, while FT consistently fails to achieve good specificity.
  • Mixed edits: Mixed-relation edits produce performance close to the average of editing each relation separately, while scaling to 700 facts.This behavior is reported across four combinations of relations with similar or different subject and object classes.

6 DISCUSSION AND CONCLUSION

MEMIT enables large-scale factual-memory updates through explicit analysis of model internals, while its representation remains limited to directional subject–relation–object associations. The authors suggest interpretability-based editing may support more transparent model updates, control, and auditing.

  • Discussion and Conclusion: MEMIT edits much larger sets of factual memories than prior approaches while maintaining specificity, generalization, and fluency.The method directly manipulates specific layer parameters and scales to 100x larger edit sets.
  • Discussion and Conclusion: The studied knowledge representation covers directional (s, r, o) relations but excludes spatial, temporal, mathematical, linguistic, procedural, and symmetric knowledge.Symmetric associations must be processed separately in opposite directions.
  • Discussion and Conclusion: Large-scale model updates can be constructed using explicit analysis of internal computations.The authors connect this result to more transparent and practical ways to edit, control, and audit models.

7 ETHICAL CONSIDERATIONS

The paper cautions that language models should not be treated as authoritative sources of facts. Memory-editing methods may reduce the cost and energy of correcting errors, but they could also insert false or damaging information.

  • Ethical Considerations: Language models are not reliable knowledge sources and should not be used as authoritative sources of facts.This caution applies despite the paper testing their ability to serve as knowledge bases.
  • Ethical Considerations: Memory-editing methods might enable malicious actors to insert false or damaging information absent from a model’s original training data.The same methods may also reduce the cost and energy needed to fix model errors.

9 REPRODUCIBILITY

The experiments use NVIDIA A6000 workstations with HuggingFace Transformers and PyTorch. GPT-NeoX’s size imposes a multi-GPU memory requirement that prevents execution on lower-memory GPUs.

  • Reproducibility: Experiments run on workstations with NVIDIA A6000 GPUs using HuggingFace Transformers and PyTorch.PyTorch executes the model-editing algorithms on GPUs.
  • Reproducibility: GPT-J experiments fit on one 48GB A6000, whereas GPT-NeoX requires at least two GPUs for model execution and editing.The GPT-NeoX setup uses one 48GB GPU for float16 inference and another slightly smaller GPU for the editing method.
  • Reproducibility: The experiments will not run on GPUs with less memory because of the models’ size.

A CAUSAL TRACING

Causal tracing identifies transformer states that mediate factual recall, guiding MEMIT’s choice of editable layers. GPT-J largely matches previously reported causal structure, with an early-attention effect that path-dependent analysis does not identify as an important factual-recall mediator.

  • Causal Tracing: MLPs at the last subject token and attention modules at the last token are identified as important causal mediators of factual-association expression.
  • Causal Tracing: Causal tracing measures how restoring hidden states after subject-token noise improves predicted probabilities for 501 correctly recalled facts.The procedure averages results over 10 noise samples and also restores runs of 10 MLP or attention layers.
  • Causal Tracing: GPT-J exhibits a causal structure similar to that reported for GPT2-XL, while early attention at the last subject token shows a distinct strong effect.
  • Causal Tracing: Path-dependent experiments suggest attention at the last subject token is not an important mediator of factual recall about the subject.The early attention effect likely reflects concentrated computation when GPT-J recognizes and chunks the subject name.

B IMPLEMENTATION DETAILS

The implementation compares MEMIT with fine-tuning, MEND, and ROME, emphasizing tuned baselines, layer-specific settings, and the computational benefits of parallel mass editing.

  • Baselines: 5×10−4 selects the optimal tradeoff between generalization and specificity for fine-tuning on 10,000 edits.The regularization strength depends strongly on the number of edits, so it is tuned specifically for n = 10,000.
  • Baselines: FT-W achieves near-perfect efficacy with one edited layer but has low specificity, indicating sufficient edit capacity.The baseline updates layer 21 of GPT-J and uses soft weight decay rather than a hard L∞-norm constraint.
  • Baselines: 98.25 seconds is MEND’s execution time for 10,000 GPT-J updates, making it the fastest evaluated method.MEND accumulates gradients from all edit examples and passes them through a hypernetwork together.
  • Baselines: 44,248.26 sec ≈12.29 hr is ROME’s time for 10,000 GPT-J edits, or approximately 4 seconds per edit.ROME performs updates iteratively using prefix sampling and covariance statistics collected from Wikitext.
  • MEMIT: MEMIT uses R = {3, 4, 5, 6, 7, 8} and λ = 15,000 on GPT-J, while GPT-NeoX uses R = {6, 7, 8, 9, 10} and λ = 20,000.The optimization uses different Wikitext covariance sample sizes and iteration counts for GPT-J and GPT-NeoX.
  • MEMIT: 3,226.35 sec ≈0.90 hr is MEMIT’s time for 10,000 GPT-J updates when z_i vectors are pre-computed and cached.The most expensive step is inverting a large square matrix; computing all z_i vectors takes 23,546.65 sec ≈6.54 hr but is embarrassingly parallel.

C EVALUATION METRICS

The evaluation measures whether edits succeed, generalize across paraphrases, preserve unrelated facts, and maintain fluent generations, while also testing performance across diverse relation mixtures.

  • Probability tests: Efficacy measures top-1 recall of the edited object on the original prompt, while Paraphrase measures accuracy on rephrasings.The prompts match what the editing method sees at runtime.
  • Probability tests: Specificity measures accuracy on unrelated or neighborhood prompts whose facts should remain unchanged.In COUNTERFACT, neighborhood prompts share the same correct answer object.
  • Aggregated metrics: The aggregated Score is the harmonic mean of Efficacy, Paraphrase, and Specificity.This combines editing success, generalization, and preservation of unrelated knowledge.
  • COUNTERFACT metrics: Efficacy Success, Paraphrase Success, and Neighborhood Success compare the edited object’s probability with the original or correct object across corresponding prompts.Editing Score is the harmonic mean of ES, PS, and NS.
  • Generation tests: Reference Score measures cosine similarity between TF-IDF vectors for free-form generations and reference Wikipedia text about the edited object.Higher similarity reflects more consistent phrasing and vocabulary with the reference text.
  • Generation tests: Generation Entropy tests excessive repetition through the entropy of n-gram frequency distributions.The metric relies on the n-gram frequency distribution f_n(·).
  • Diversity analysis: MEMIT’s effectiveness on mixed relation pairs closely follows the average performance of the individual relation splits.The comparison uses edit sets with equal numbers of examples from each relation and four levels of diversity.

E DEMONSTRATIONS

The demonstrations apply MEMIT to refreshed election knowledge and specialized astronomy facts, then examine how layer selection, attention editing, and covariance adjustment affect performance.

  • Knowledge freshness: 100% efficacy (ES) and 94% generalization (PS) were achieved when MEMIT incorporated 2022 United States election results into GPT-J.The edits represented congressional, gubernatorial, and senatorial outcomes as subject-relation-object tuples.
  • Specialized knowledge: 86% accuracy after MEMIT exceeded the 53% accuracy of unmodified GPT-J on 289 star–constellation tuples.The specialized astronomy domain covered stars from 18 constellations.
  • Varying edited layers: Using more critical MLP layers increases efficacy, generalization, and specificity, whereas editing late-layer MLPs performs considerably worse.The tested layer configurations include regions with high and low causal effect, supporting the role of causal analysis in selecting layers.
  • Varying edited layers: Attention edits perform considerably worse than edits to the corresponding MLP modules.The comparison tests both early- and late-layer attention modules.
  • Covariance adjustment: Around λ ≈ 10^4, the aggregated score reaches a maximum as specificity and fluency trade off against efficacy and generalization.Specificity and fluency increase monotonically with λ, while efficacy and generalization decrease as λ increases.
Loading 2210.07229v2…