PaperlayerView exact revision on arXiv

Source-linked AI summary

Aging with GRACE: Lifelong Model Editing with Discrete Key-Value Adaptors

Thomas Hartvigsen, Swami Sankaranarayanan, Hamid Palangi, Yoon Kim, Marzyeh Ghassemi

arXiv:2211.11031v5cs.LG

TL;DR

Deployed models develop sequential errors, but retraining is expensive and existing editors degrade under repeated edits. GRACE stores localized transformations in a latent-space codebook without changing model weights. Experiments report strong sequential-editing performance across T5, BERT, and GPT, with generalization to new inputs and a small fixed inference cost.

  • Problem

    Lifelong model editing seeks to correct streaming deployment errors repeatedly while retaining prior fixes and pretrained behavior without expensive retraining.

  • Method

    GRACE adds a detachable, similarity-retrieved codebook adaptor that stores latent-space keys and learned values while leaving pretrained weights unchanged.

  • Results

    GRACE outperforms seven alternatives on sequential edits to T5, BERT, and GPT models across question answering, classification, and generation, while generalizing to new inputs.

  • Takeaways & Limitations

    GRACE supports thousands of sequential edits using streaming errors and preserves unrelated behavior through localized, inspectable codebook changes.

  • Takeaways & Limitations

    GRACE adds a one-time 1.32x average inference slowdown, does not consider multi-layer edits, and may require future scaling beyond approximately 5k edits.

Abstract

from arXiv · show

Deployed language models decay over time due to shifting inputs, changing user needs, or emergent world-knowledge gaps. When such problems are identified, we want to make targeted edits while avoiding expensive retraining. However, current model editors, which modify such behaviors of pre-trained models, degrade model performance quickly across multiple, sequential edits. We propose GRACE, a lifelong model editing method, which implements spot-fixes on streaming errors of a deployed model, ensuring minimal impact on unrelated inputs. GRACE writes new mappings into a pre-trained model's latent space, creating a discrete, local codebook of edits without altering model weights. This is the first method enabling thousands of sequential edits using only streaming errors. Our experiments on T5, BERT, and GPT models show GRACE's state-of-the-art performance in making and retaining edits, while generalizing to unseen inputs. Our code is available at https://www.github.com/thartvigsen/grace}.

1 Introduction

Deployed models make sequential, costly-to-correct errors, while existing editing approaches can forget prior edits or degrade pretrained behavior. GRACE addresses this with weight-preserving, targeted codebook edits that generalize to similar inputs and support long edit sequences.

  • Deployed models can hallucinate, perpetuate bias, or factually decay, creating errors that may require immediate correction.
  • Existing continual-learning and model-editing approaches can overfit, forget previous edits or pretraining data, require impractical auxiliary data, and degrade after few edits.
  • Sequential editing must correct flagged errors, generalize to similar future inputs, and preserve unrelated model behavior using only singular streaming inputs.
  • GRACE adds a layer adaptor that caches error embeddings and learned output-decoding values in a codebook without changing model weights.
  • GRACE uses latent-space ε-balls to apply edits near cached keys, encouraging generalization while limiting interference with correct behaviors.
  • GRACE outperforms seven alternatives across sequential edits to T5, BERT, and GPT models on question answering, classification, and generation.Edits generalize to new inputs without memorizing and incur only a small, one-time inference cost over long edit sequences.

2 Methods: Lifelong Model Editing with GRACE

Lifelong model editing repeatedly updates a deployed model while retaining prior fixes and upstream performance. GRACE uses a discrete latent-space codebook with similarity-based deferral, adaptive radii, and learned values to make localized edits without changing weights.

  • Problem formulation: Lifelong model editing updates the same model hundreds to thousands of times without forgetting upstream performance or previous fixes.
  • GRACE adaptor: GRACE wraps a chosen layer with an adaptor containing a codebook and a deferral mechanism that decides whether to apply a cached transformation.
  • GRACE codebook: Each codebook entry stores a cached activation key, a learned value, and a deferral radius controlling similarity-based activation.
  • Deferral mechanism: GRACE retrieves the value associated with the closest key when the query lies within that key’s radius; otherwise, the pretrained layer runs unchanged.
  • Codebook maintenance: New edits create entries, same-label overlaps expand an existing radius, and different-label overlaps split radii before adding a new entry.
  • Codebook maintenance: The initial radius εinit trades broader edit generalization against greater interference with unrelated inputs.
  • Value training: GRACE trains replacement values through backpropagation on the model’s prediction loss while leaving model weights untouched.
  • Illustrative example: In a synthetic local-label-shift experiment, one GRACE key corrected flipped-label errors while barely influencing other inputs, unlike finetuning on those errors alone.

3 Experiments

GRACE is evaluated against continual-learning and model-editing baselines on sequential edits across question answering, classification, and generation. It maintains a stronger balance between editing success and retention, generalizes through compact codebooks, and adds limited inference overhead.

  • Experimental Setup: GRACE is evaluated on authentic streaming errors across T5, BERT, and GPT models, contrasting continual finetuning, replay, EWC, MEND, Defer, and ROME.The tasks cover question answering, document classification, and language generation, with metrics including ES, TRR, and ERR.
  • Comparisons to Existing Methods: 19% and 9%: GRACE’s averaged TRR and ERR outperform the closest competitors on zsRE and SCOTUS, respectively.Across all methods, the corresponding improvements are 86% and 129% because alternatives struggle to balance retention and editing.
  • Comparisons to Existing Methods: GRACE outperforms comparisons on Hallucination while finetuning methods trade competitive ERR for poor TRR and degraded retention on already-accurate sentences.ROME and Memory remain competitive, while GRACE makes edits twice as fast as finetuning.
  • Comparisons to Existing Methods: 1000 edits use only 137 keys on T5 zsRE, demonstrating that GRACE can compress many edits into a small codebook.On zsRE, the codebook contains 210,569 scalar values, of which 70,281 are learnable, or 0.35% of T5’s parameters.
  • Model Analysis: Memorization vs. Generalization: Interior layers and larger ϵinit values improve generalization, while blocks two and four achieve high TRR, high ERR, and strong holdout performance.Small ϵinit values favor TRR by creating more keys, whereas blocks 0 and 6 can generalize worse; increasing ϵinit also produces smaller codebooks.
  • Inference Time: 1.32x: GRACE-edited T5 inference is slower than unedited inference on average, but this one-time cost remains fixed as the codebook grows until memory is exhausted.The fixed cost follows from vectorized search over the keys.

4 Related Work

Existing approaches to lifelong model editing struggle with sequential updates because they overfit, forget prior knowledge, or depend on privileged data. GRACE addresses this setting using singular streaming inputs and caching mechanisms that support longer-term memory.

  • Model editing: Most model-editing methods use regularized finetuning with auxiliary data, such as pre-training examples or semantically equivalent versions of edits.
  • Model editing: Static editors generally make one edit, while multiple-edit methods often operate simultaneously or experience rapid performance decay over sequential edits.
  • Continual learning: Continual learning is a plausible approach, but sequential editing can cause overfitting and forgetting of previous edits and pre-training data.
  • Continual learning: GRACE-related caching approaches retain values to keep inputs in-distribution for downstream encoders and support longer-term memory when resources permit.

5 Limitations and Ethical Considerations

GRACE has practical limitations: similarity search slows inference, the method does not consider multi-layer edits, and scaling beyond approximately 5k edits remains open. Its isolated edits also do not update related knowledge automatically.

  • Method limitations: Similarity search added to model layers slows inference, although accelerating GRACE is identified as a future direction.
  • Method limitations: GRACE has not been evaluated for multi-layer edits, leaving that setting for future work.
  • Method limitations: Although GRACE scales to approximately 5k edits, larger real-world editing scales remain an open direction.
  • Implications: GRACE-style edits update behaviors in isolation, so changing one item of knowledge does not update related items automatically.
  • Ethical considerations: Model editing can be used harmfully, including by increasing hate speech, a limitation shared by GRACE and other model editors.
  • Ethical considerations: Unlike weight-updating editors, GRACE’s inspectable codebook can help trace which predictions arise from particular stored values.

6 Conclusions

The paper studies lifelong model editing for transformer language models, using authentic streaming errors and thousands of sequential edits. It proposes GRACE, a plug-in adaptor that leaves trained weights untouched while retaining original functionality and editing predictions.

  • Study scope: The study evaluates transformer models for natural-language-processing tasks under simulated deployments with thousands of sequential streaming edits.
  • Study scope: Edits use singular authentic errors rather than synthetic edits or privileged sets of training edits, semantically equivalent examples, or pre-training data.
  • GRACE: GRACE is a plug-in adaptor attached to a chosen model layer that leaves trained weights untouched.
  • GRACE: GRACE adaptors retain original model functionality while successfully editing model predictions without forgetting previous behavior.

A Implementation Details

The experiments use sequential single-example updates across several model architectures, with shared editing layers for comparisons. GRACE’s key hyperparameter controls codebook-entry initialization, while runtime measurements report only per-edit time.

  • Training details: All methods use Adam with batch size 1 because edits arrive as singular sequential updates.
  • Hyperparameters: Comparisons use learning rates from 1.0 to 1e−5; Finetuning, Memory, and MEND work best at 1e−2, while GRACE and Defer require 1.0.
  • Layer selection: Each editor modifies the same layer within each model: T5’s final encoder dense layer, BERT’s second-to-last layer, and GPT2-XL’s fully connected component in layer 36.
  • GRACE hyperparameters: GRACE’s unique hyperparameter, ϵinit, sets the initial ϵ-ball size around new codebook entries.
  • Runtime: Runtime results on GPT2-XL report only the average time required to make one edit, excluding preprocessing, precomputation, and pretraining.

B Additional Dataset and Model Descriptions

The experiments use T5, BERT, and GPT2-XL across question answering, document classification, and hallucination language modeling, comparing GRACE with finetuning, regularization, retraining, hypernetwork, deferral, memory, and ROME-based editors.

  • Datasets and models: The evaluation covers zsRE question answering, SCOTUS document classification, and Hallucination language modeling with T5, BERT, and GPT2-XL.The datasets include sequential edits to question–answer pairs, relabeled court-ruling topics, and corrections to hallucinated biography sentences.
  • Datasets and models: The zsRE edit set contains 200 question–answer pairs with at least five rephrasings each, yielding 1000 possible edits.For generalization and long-sequence experiments, the procedure uses 1000 pairs with 10 rephrasings, split into 5000 edits and 5000 holdouts.
  • Datasets and models: The SCOTUS task predicts one of 14 topics, with selected labels relabeled and edits performed on 1991–2009 test documents.Training and validation use different time ranges, and TRR is computed on 914 validation documents from 1982–1991.
  • Datasets and models: The Hallucination dataset contains 1392 potential edits and corrects GPT-3-generated biography sentences using corresponding Wikipedia sentences.These edits use preceding generated text as prompts and the correct next sentence as labels, making them longer and more authentic than many prior editing tasks.
  • Compared editors: Baselines include ordinary finetuning, EWC-regularized finetuning, retraining, MEND, ROME, Defer, and Memory.MEND uses a hypernetwork trained on representative edits and related inputs; Defer selects between the pretrained prediction and a learned replacement, while Memory retrieves learned activations.

D Interpreting GRACE codebooks

GRACE codebooks generalize accurately when unseen holdout edits fall within a key’s similarity radius, but larger initial radii trade larger coverage for smaller codebooks.

  • Generalization: Holdout edits that fall inside a key’s ϵ ball achieve highly accurate F1 scores, while holdouts without associated keys remain uncovered.The figure defines F1 on unseen holdout edits landing inside any key’s ϵ ball and compares coverage with the true average number of rephrasings.
  • Radius and coverage: An initial radius of ϵinit = 0.1 eventually captures 60% of holdouts but grows the codebook to 547 entries.Larger ϵinit values capture more holdouts per key and produce smaller codebooks; smaller values capture fewer holdouts and produce larger codebooks.

E Parameter Efficiency

GRACE stores edits in detachable codebook entries rather than model weights, making memory grow with the number of edits and enabling substantial parameter savings.

  • Memory requirements: Each new GRACE edit requires |h_l−1| + |h_l| + 1 parameters, with the key frozen and only the value and radius-related parameters learnable.The entry stores a cached key activation, a value, and a scalar deferral radius.
  • Memory requirements: 500 edits require 768,500 stored parameters and 256,500 learnable parameters for the 60-million-parameter T5 model.The stored parameters equal 500 × (1024 + 512 + 1), while learnable parameters equal 500 × (512 + 1).
  • Codebook entries: The codebook’s entries correspond to inputs fixed by the same key and the true label associated with that key’s learned value.The table identifies the codebook entry index, the corrected inputs, and the shared true label.

F Ablation Study

Ablations show that GRACE’s radius, layer, value initialization, and token-replacement choices govern the trade-off between codebook size, retention, generalization, and edit success.

  • Codebook size and retention: Up to 1000 edits often compress into about 200 keys, whereas small ϵinit values can produce about 600 keys.The compression reflects increasing ϵ values for new edits and exposes a trade-off between codebook size and memorization.
  • Codebook size and retention: With ϵinit = 0.1, TRR, ERR, and ES remain high throughout editing, while larger-radius settings reduce TRR and ERR as codebooks trade size for generalizability.The same trend holds when editing Block 7, although that later layer is harder to edit successfully.
  • Replacement token: Replacing only the last input token appears superior for Block 2, especially for TRR, while remaining comparable for ERR and ES.For Block 7, the opposite replacement pattern appears, but low overall ES makes the difference unclear.
  • Value initialization: Cold value initialization appears worse than warm initialization for Block 7 despite higher ES, while Block 2 shows negligible differences.The comparison evaluates warm and cold starts across the ablation settings.
  • Edit compression: GRACE can compress multiple edits into one key, but at least one key is required for each unique desired label.The codebook examples show several inputs leading to the same key and learned value.

G Extended Main Results

Extended experiments show that GRACE maintains strong performance across sequential editing comparisons and that its core trade-offs persist over much longer edit streams. Fine-grained analyses also examine how initialization and layer choice affect retention and generalization.

  • GRACE achieves strong performance compared to the comparison methods across editing tasks and over time.The reported trends match the main-paper results in both barplots and finer-grained per-edit plots.
  • GRACE’s main claims remain true in experiments extending to 5,000 sequential edits on the QA editing task.The extended analysis evaluates hyperparameter and layer effects throughout a substantially longer edit stream.
  • Figures 12–16 compare editors over time on BERT SCOTUS, GPT2-XL Hallucination, zsRE, SCOTUS, and Hallucination.These plots provide task-specific views of editing trajectories across models and benchmarks.
  • Editing blocks 0 and 6 use more keys and achieve higher TRR, but produce lower ERR and worse Holdout generalization after 3,000 edits.The Holdout set contains unseen rephrasings of all 3,000 edits, while TRR and ERR measure editing behavior.
Loading 2211.11031v5…