PaperlayerView exact revision on arXiv

Source-linked AI summary

PMET: Precise Model Editing in a Transformer

Xiaopeng Li, Shasha Li, Shezheng Song, Jing Yang, Jun Ma, Jie Yu

arXiv:2308.08742v6cs.CLcs.AIcs.LG

TL;DR

Existing model-editing methods optimize mixed Transformer-layer hidden states for FFN updates, although those states also contain MHSA and residual information. PMET optimizes MHSA and FFN component states but updates only FFN weights, and reports state-of-the-art results on COUNTERFACT and zsRE. The paper further analyzes MHSA’s extraction patterns and factual-knowledge storage, while noting that benchmark success does not necessarily mean edited knowledge is internally reasoned over.

  • Problem

    Existing methods use Transformer-layer hidden states containing MHSA, FFN, and residual information as FFN target representations, limiting the precision of model editing.

  • Method

    PMET simultaneously optimizes MHSA and FFN Transformer-component hidden states, then uses only optimized FFN states to update FFN weights precisely.

  • Results

    PMET achieves state-of-the-art comprehensive performance on zsRE and COUNTERFACT, with 3.3% average reliability improvement on COUNTERFACT and 0.4% on zsRE over the state-of-the-art method.

  • Takeaways & Limitations

    MHSA encodes general knowledge-extraction patterns and stores a small amount of factual knowledge, so PMET supplements MHSA-state optimization without updating MHSA weights.

  • Takeaways & Limitations

    Editing performance does not necessarily indicate true internalization: PMET- and MEMIT-edited models may fail to reason using edited knowledge.

Abstract

from arXiv · show

Model editing techniques modify a minor proportion of knowledge in Large Language Models (LLMs) at a relatively low cost, which have demonstrated notable success. Existing methods assume Transformer Layer (TL) hidden states are values of key-value memories of the Feed-Forward Network (FFN). They usually optimize the TL hidden states to memorize target knowledge and use it to update the weights of the FFN in LLMs. However, the information flow of TL hidden states comes from three parts: Multi-Head Self-Attention (MHSA), FFN, and residual connections. Existing methods neglect the fact that the TL hidden states contains information not specifically required for FFN. Consequently, the performance of model editing decreases. To achieve more precise model editing, we analyze hidden states of MHSA and FFN, finding that MHSA encodes certain general knowledge extraction patterns. This implies that MHSA weights do not require updating when new knowledge is introduced. Based on above findings, we introduce PMET, which simultaneously optimizes Transformer Component (TC, namely MHSA and FFN) hidden states, while only using the optimized TC hidden states of FFN to precisely update FFN weights. Our experiments demonstrate that PMET exhibits state-of-the-art performance on both the COUNTERFACT and zsRE datasets. Our ablation experiments substantiate the effectiveness of our enhancements, further reinforcing the finding that the MHSA encodes certain general knowledge extraction patterns and indicating its storage of a small amount of factual knowledge. Our code is available at https://github.com/xpq-tech/PMET.

Introduction

Model editing modifies limited LLM knowledge at lower cost than retraining, but existing optimization-based methods use mixed Transformer-layer states for FFN updates. PMET separates component roles by optimizing MHSA and FFN states while updating only FFN weights, achieving strong benchmark performance.

  • Motivation: Model editing changes a minor proportion of LLM knowledge without vanilla training or fine-tuning, with efficacy, generalization, specificity, fluency, and consistency measuring outcomes.Weight-modified methods directly edit model weights and avoid the additional content often required by weight-preserved methods.
  • Problem: Existing optimization-based methods use optimized Transformer-layer hidden states to update FFN weights, even though those states combine MHSA, FFN, and residual information.This produces imprecise updates from non-accurate target knowledge representations.
  • Analysis: MHSA hidden states change frequently during knowledge extraction, while FFN intermediate states stabilize after the 15th GPT-J layer, suggesting different component roles.The authors interpret MHSA as encoding general knowledge-extraction patterns and FFN as more stable in later layers.
  • Method: PMET simultaneously optimizes MHSA and FFN Transformer-component hidden states but uses only optimized FFN states as target representations for precise FFN-weight updates.The supplementary MHSA optimization addresses occasional bottlenecks when FFN states cannot align with target knowledge.
  • Results: 3.3% average reliability improvement on COUNTERFACT and 0.4% on zsRE are reported over the state-of-the-art method.The experiments cover GPT-J (6B) and GPT-NeoX (20B), while ablations report a balance among reliability, specificity, fluency, and consistency.
  • Contributions: MHSA is reported to act as a knowledge extractor that stores general extraction patterns and a small amount of factual knowledge.PMET therefore avoids using optimized MHSA states to update MHSA weights.

Related Work

Related work classifies model editing by whether original weights are modified and examines Transformer components to clarify how knowledge is represented and recalled. This literature motivates PMET’s focus on separating MHSA extraction patterns from FFN factual storage.

  • Model Editing: Weight-preserved methods protect non-target knowledge using external models, in-context learning, or representation changes, but additional content grows as modifications increase.Weight-modified methods directly alter model weights and avoid this increasing-content issue.
  • Model Editing: Earlier weight-modified approaches include multi-loss and constrained fine-tuning, followed by meta-learning and optimization-based methods addressing overfitting and editing limitations.ROME incrementally updates FFN weights, while MEMIT enables mass editing in one operation.
  • Post-Hoc Explanation of Transformers: Prior Transformer analyses generally identify FFN as a main carrier of factual knowledge and MHSA as responsible for token associations, content interactions, and subject-attribute extraction.Other studies report redundant information in MHSA, suggesting possible general patterns for knowledge extraction.
  • Post-Hoc Explanation of Transformers: Existing findings suggest MHSA may store general extraction patterns, but they do not fully clarify whether it stores factual knowledge.The paper addresses this question by analyzing MHSA and FFN hidden states.

Methodology

PMET reframes editing around subject-associated knowledge and separates MHSA and FFN hidden-state roles. It jointly optimizes component states but uses only optimized FFN states to update FFN weights.

  • PMET defines editing from a subject-centric perspective, associating multiple knowledge instances with a subject so edited models can reason about that subject.
  • Investigating the Role of MHSA and FFN in LLMs’ Knowledge Recall: The analysis compares hidden-state changes before and after MHSA and FFN in both vector and vocabulary spaces across GPT-J layers.It uses 1209 factual statements and top-50 vocabulary mappings for the analysis.
  • Investigating the Role of MHSA and FFN in LLMs’ Knowledge Recall: After layer 15, FFN hidden states change more slowly and stabilize, whereas MHSA states continue changing frequently with uncertain directions during knowledge extraction.
  • Investigating the Role of MHSA and FFN in LLMs’ Knowledge Recall: These findings motivate treating MHSA as encoding general knowledge-extraction patterns, so PMET does not update MHSA weights when introducing new knowledge.
  • PMET Method: PMET jointly optimizes MHSA and FFN component hidden states, then retains only optimized FFN states as target representations for updating FFN weights.The method computes target representations in the last critical FFN layers and updates only FFN weights there.
  • PMET Method: For multi-layer editing, PMET spreads the residual across critical layers using a square-root schedule rather than MEMIT’s even spread.Keys represent hidden states before the specific edited weights, and the FFN update follows MEMIT’s algorithmic steps afterward.

Experiments

PMET is evaluated against established editing baselines on COUNTERFACT and zsRE, including large-scale edits and ablations. It generally improves reliability while trading off some specificity, with the balance depending on weight updates and residual spreading.

  • Editing performance across edit counts: PMET outperforms all baselines except MEMIT on specificity across 17 GPT-J editing settings, while leading on the other reported metrics.The settings vary with the number of edits sampled on a log-scale curve.
  • COUNTERFACT results: PMET leads existing methods in score, efficacy, fluency, and consistency on 10K COUNTERFACT edits, but trails MEMIT and MEND in specificity.PMET and MEMIT prioritize reliability, whereas MEND favors specificity.
  • COUNTERFACT results: On 10K GPT-NeoX COUNTERFACT edits, PMET exceeds MEMIT in reliability and consistency but lags in specificity.The authors associate the specificity gap with larger model changes from square root propagation.
  • zsRE results: On 10K zsRE edits, PMET outperforms existing methods in efficacy, generalization, and specificity.The original GPT-J model has specificity 27.0, limiting the edited models’ specificity relative to that reference.
  • Ablation study: Ablations show that jointly optimizing MHSA and FFN hidden states improves reliability, while updating MHSA weights marginally improves generalization but harms specificity.The experiments were conducted on COUNTERFACT with GPT-J (6B).
  • Ablation study: Square root spreading improves reliability but causes larger model changes and reduced specificity, whereas even spreading preserves specificity and fluency at lower reliability.The authors interpret even spreading as losing update information and square root spreading as mitigating that loss with more side effects.

Conclusion

The paper identifies MHSA as a knowledge extractor that encodes general extraction patterns and some factual knowledge. PMET uses this insight to optimize MHSA and FFN hidden states while updating only FFN weights, achieving state-of-the-art editing performance and balancing multiple metrics.

  • Conclusion: MHSA encodes general knowledge extraction patterns and stores a small amount of factual knowledge, clarifying its distinct role from FFN in model editing.The findings provide additional insight into the roles of MHSA and FFN.
  • Conclusion: PMET jointly optimizes MHSA and FFN hidden states but uses only optimized FFN hidden states to update FFN weights for precise editing.This design avoids unnecessary MHSA weight updates.
  • Conclusion: Experiments on zsRE and COUNTERFACT demonstrate state-of-the-art performance, while ablations support the method’s enhancements and balance across metrics.The conclusion also reports that MHSA stores a small amount of factual knowledge.

Limitations

The paper reports that PMET and MEMIT’s evaluation performance does not necessarily indicate true internalization of edited knowledge. It also identifies benchmark-definition mismatches and calls for more sophisticated methods and metrics.

  • Edited models may fail to reason using newly edited knowledge despite promising evaluation results.The paper gives an example where an edited fact does not support related reasoning.
  • The subject-centric definition of knowledge editing was not strictly followed in benchmark and dataset construction.The benchmarks were adapted to existing evaluation methods rather than fully adhering to the paper’s definition.
  • The authors propose developing more sophisticated editing methods and evaluation metrics in future work.They cite MQuAKE and RIPPLEEDITS as examples of potentially useful evaluation directions.

Ethical Statement

Model editing is intended to correct erroneous and outdated knowledge, but it could also be misused to generate false, toxic, and harmful content. The authors therefore caution against excessive trust in generated content until LLMs are well regulated.

  • Model editing can be misused to make LLMs generate false, toxic, and harmful content.
  • The authors emphasize not placing excessive trust in generated content until LLMs are well regulated.
  • The intended use of model editing is correcting errors and outdated knowledge in LLMs.

A. Analysis of Incremental Weight

The ablation analysis links incremental-weight norms to editing behavior. PMET achieves the best overall performance among comparable cases while producing the smallest incremental-weight norm.

  • The ablation study examines how incremental-weight norms relate to editing performance and model changes.Figure 4 reports norm changes for the ablation cases.
  • Smaller norms in the even-spread case result from greater residual reduction and loss of update information.This information loss compromises reliability but favors preservation of the original model.
  • MHSA-weight updates have smaller norms than other cases because GPT-J’s MHSA weights are 0.25 times its FFN weights.
  • PMET, MHSA-only updating, and FFN-only updating have similar incremental-weight norms and editing performance.The three cases show closely matched norms and comparable results in Table 3.
  • PMET has the smallest incremental-weight norm among the comparable cases while achieving the best overall performance.The authors interpret this as the incremental weight more accurately capturing the updates needed for editing.

B. Metrics of Model Editing Problem

The model-editing problem is evaluated through reliability and specificity, with reliability separated into efficacy and generalization. The paper adapts subject-centric knowledge into explicit and implicit sets and also considers fluency and consistency.

  • Reliability measures success on knowledge related to the target, whereas specificity measures success on unrelated knowledge.Specificity is also called locality.
  • Reliability is divided into efficacy on edit sequences and generalization on semantically consistent paraphrase sequences.The two sequence types differ in syntax while preserving semantic content.
  • Knowledge related to a subject is divided into explicit and implicit knowledge sets to align the problem definition with existing metrics and datasets.
  • Explicit knowledge pairs are directly stated, while implicit knowledge contains information derived from corresponding explicit knowledge.The paper illustrates this distinction using the Forbidden City’s location and related derived knowledge.
  • Fluency and consistency are also used to assess the generation capability of edited models.

C. Datasets Detail

The evaluation uses zsRE and COUNTERFACT to assess whether edited models correct explicit and implicit knowledge while preserving unrelated knowledge. COUNTERFACT additionally evaluates generation quality through fluency and consistency.

  • zsRE: zsRE evaluates editing by changing a factual answer and testing explicit questions, rephrased implicit questions, and unrelated questions.The example changes Watts Humphrey’s university from Trinity College to University of Michigan while preserving the answer to an unrelated question.
  • Evaluation metrics: Efficacy measures success on explicit questions, generalization measures success on implicit questions, and specificity measures preservation on unrelated questions.These three metrics correspond to the first three evaluation targets in zsRE and COUNTERFACT.
  • COUNTERFACT: COUNTERFACT uses similar efficacy, generalization, and specificity benchmarks but includes multiple paraphrases and neighborhood prompts.Neighborhood prompts serve as COUNTERFACT’s counterpart to zsRE’s unrelated loc questions.
  • COUNTERFACT: COUNTERFACT also contains generation prompts for evaluating the edited model’s fluency and consistency.These prompts contain subjects and assess generation capability.

D. Experimental Detail

Experiments update identified critical layers of GPT-J and GPT-NeoX using settings for hidden-state optimization, covariance estimation, and norm constraints. PMET’s algorithmic steps are fundamentally similar to MEMIT, yielding almost equivalent time consumption.

  • Model and layer settings: GPT-J uses critical layers R = {3, 4, 5, 6, 7, 8}, while GPT-NeoX uses R = {6, 7, 8, 9, 10}.The experiments mainly update these identified critical layers, with baseline settings following MEMIT.
  • Hidden-state optimization: φ starts at 1 and µ ranges from 0 to 1; increasing µ retains more original knowledge, while increasing φ has the opposite effect.The optimization later sets φ = 0.1 and stops when DKL < 0.01.
  • Covariance estimation: GPT-J covariance estimation samples 10K Wikitext instances with λ = 6000, or λ = 4500 for zsRE, while GPT-NeoX uses 5K samples and λ = 15000.GPT-J uses fp32 sampling; GPT-NeoX samples in fp16 and stores the covariance matrix in fp32.
  • Optimization constraints: PMET optimizes hidden states for 30 steps with model-specific learning rates and constrains their norms relative to the original intermediate states.GPT-J uses learning rate 0.2, or 0.5 in zsRE; GPT-NeoX uses 0.5, with constraints of less than 3⁄4 and 4⁄5 of original norms, respectively.
  • Runtime: PMET and MEMIT have almost equivalent time consumption because their algorithmic steps are fundamentally similar.The comparison concerns the methods’ overall algorithmic cost.
Loading 2308.08742v6…