PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Model Editing at Scale

Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, Christopher D. Manning

arXiv:2110.11309v2cs.LGcs.AIcs.CL

TL;DR

Post-hoc correction of isolated errors in very large neural networks is difficult because fine-tuning can overfit and existing editing methods may not scale. MEND transforms fine-tuning gradients into targeted updates and is reported to edit the largest GPT and T5 models effectively, including many simultaneous edits.

  • Problem

    Localized post-deployment correction of large neural networks remains difficult because their distributed representations make targeted edits non-trivial, while fine-tuning can overfit and other methods may not scale.

  • Method

    MEND trains lightweight editor networks to transform fine-tuning gradients into targeted parameter updates using their low-rank structure.

  • Results

    MEND consistently outperformed comparison methods on very large Transformer editing tasks and achieved 96% edit success with less than 1% accuracy degradation for 25 simultaneous edits.

  • Takeaways & Limitations

    MEND is reported as the only method that successfully edits the largest publicly available GPT and T5 Transformer models.

  • Takeaways & Limitations

    MEND can over-generalize edits to distinct but related inputs because its locality examples may be insufficiently challenging.

Abstract

from arXiv · show

While large pre-trained models have enabled impressive results on a variety of downstream tasks, the largest existing models still make errors, and even accurate predictions may become outdated over time. Because detecting all such failures at training time is impossible, enabling both developers and end users of such models to correct inaccurate outputs while leaving the model otherwise intact is desirable. However, the distributed, black-box nature of the representations learned by large neural networks makes producing such targeted edits difficult. If presented with only a single problematic input and new desired output, fine-tuning approaches tend to overfit; other editing algorithms are either computationally infeasible or simply ineffective when applied to very large models. To enable easy post-hoc editing at scale, we propose Model Editor Networks using Gradient Decomposition (MEND), a collection of small auxiliary editing networks that use a single desired input-output pair to make fast, local edits to a pre-trained model's behavior. MEND learns to transform the gradient obtained by standard fine-tuning, using a low-rank decomposition of the gradient to make the parameterization of this transformation tractable. MEND can be trained on a single GPU in less than a day even for 10 billion+ parameter models; once trained MEND enables rapid application of new edits to the pre-trained model. Our experiments with T5, GPT, BERT, and BART models show that MEND is the only approach to model editing that effectively edits the behavior of models with more than 10 billion parameters. Code and data available at https://sites.google.com/view/mend-editing.

1 INTRODUCTION

MEND addresses the difficulty of making targeted post-deployment corrections to large neural networks without harming unrelated behavior. It uses lightweight editor networks and low-rank gradient transformations to enable scalable, fast edits to very large language models.

  • Motivation: Deploying and maintaining large models requires localized post-deployment updates that correct undesirable outputs for individual inputs.The introduction frames targeted model patching as a key challenge despite large models’ improved performance across vision and language tasks.
  • Editing objectives: Useful edits should provide reliability on the problematic input, locality on unrelated inputs, and generality across related inputs.The example contrasts changing the answer to a UK prime-minister question while preserving an unrelated sports-team answer.
  • Method: MEND trains lightweight editor networks that transform standard fine-tuning gradients into weight updates, using low-rank structure to avoid costly direct gradient-to-update parameterization.The method treats the fine-tuning gradient as an information-rich input while addressing the enormous cost of mapping high-dimensional gradients directly.
  • Contribution: The work contributes a scalable algorithm for fast editing of very large pre-trained language models and evaluates it across language tasks and transformer architectures.The stated contribution emphasizes editing the largest GPT-style and T5 language models consistently.

2 THE MODEL EDITING PROBLEM

The model editing problem is to use one input–output pair to change a model’s behavior on that example and its equivalence neighborhood while preserving unrelated behavior. Editors are evaluated by reliability, locality, generalization, and efficiency, with edit success combining reliability and generality.

  • Problem formulation: A model editor uses a single edit pair (xe, ye) to produce new parameters θe that change the base model’s output for xe and related examples.The editor maps the edit input, desired label, loss function, base parameters θ, and optional editor parameters φ to θe.
  • Problem formulation: The desired edit should generalize across the equivalence neighborhood N(xe, ye), such as paraphrased questions with the same answer.For example, “Who is the prime minister of the UK?” and “Who is the UK PM?” can share the label “Boris Johnson.”
  • Evaluation criteria: A reliable editor predicts ye on xe, while a local editor causes small pre- versus post-edit disagreement on unrelated inputs.Locality is quantified using the expected KL divergence between the original and edited model predictions on randomly sampled locality inputs.
  • Evaluation criteria: An editor generalizes when it predicts y′e on alternative inputs from N(xe, ye), and is efficient when computing φ and evaluating E require little time and memory.Edit success summarizes reliability and generality through average accuracy on the edit input and uniformly sampled equivalence-neighborhood inputs.

3 MODEL EDITOR NETWORKS WITH GRADIENT DECOMPOSITION

MEND learns auxiliary editor networks that transform single-example fine-tuning gradients into targeted, one-step parameter updates for local model edits. It makes this mapping tractable through rank-1 gradient decomposition, low-rank editor parameterization, parameter sharing, and edit-locality training objectives.

  • Core method: MEND learns editor networks gℓ that transform a raw fine-tuning gradient into a targeted parameter update for a single-step model edit.Each network edits the weights of its associated layer using a new edit pair (xe, ye) at test time.
  • Gradient decomposition: A d × d weight matrix would otherwise require 2d^3 parameters for a linear gradient-to-edit mapping, motivating MEND’s low-rank decomposition.For MLPs, each per-example weight gradient is rank 1, allowing the editor to condition on 2d inputs rather than d^2.
  • Parameter sharing: MEND shares editor parameters across layers with matching weight shapes, requiring only two parameter sets when editing all MLP layers in a transformer architecture.The editor’s total parameters remain O(d) through low-rank factorization.
  • Editor architecture: MEND parameterizes each gℓ as a single-hidden-layer MLP with low-rank weight matrices and residual connections, producing pseudoactivations that form the edited gradient.The final update is ˜W = Wℓ−αℓ ˜∇Wℓ, with a learned per-layer scalar step size αℓ.
  • Training objective: Training optimizes edit success and locality losses through the edited model, using an independently sampled locality example and Adam updates to the editor parameters.The locality loss is the KL divergence between pre-edit and post-edit predictions on unrelated inputs, and experiments use ce = 0.1.

4 RELATED WORK

Prior model-editing methods constrain parameter updates, but such constraints may not ensure localized changes in model behavior. MEND instead distinguishes itself from related gradient-transform, probing, and multi-example editing approaches by using factorized gradients and a single error example for model editing.

  • Parameter-space editing: Parameter-space constraints used by prior fine-tuning methods do not always translate into useful function-space constraints for neural networks.These methods include reducing parameter-space distance or finding a minimum-L2-norm update that edits the output.
  • Model-editor comparison: MEND provides a unique combination of useful attributes among the conceptual model-editor comparisons in Table 1.The table compares whether editors preserve the model, use only the edit pair at test time, support batched edits, and scale to 10B parameters.
  • Gradient transformations: Prior meta-learning methods transform gradients for few-shot learning but do not leverage factorized gradients, limiting their transformations to simpler forms.The passage describes these transformations as typically linear and notes that they may affect the model’s forward-pass function.
  • Knowledge localization: Knowledge-probing studies test where linguistic structures or world facts reside in language models but typically do not intervene on the models’ knowledge.The cited exceptions to this pattern are Dai et al. (2021) and Wang et al. (2020).
  • Knowledge localization: Unlike exceptions that assume many datapoints representing the knowledge to be edited, this work considers model editing from a single example illustrating the model’s error.The comparison concerns the amount of evidence available for editing, not the broader probing methods themselves.

5 EXPERIMENTS

Experiments evaluate MEND across editing datasets, model scales, simultaneous edits, and design ablations using edit success and locality-based drawdown metrics. MEND is presented as effective on very large models, consistently strong across scales, efficient for multi-edit settings, and sensitive to key architectural choices.

  • Experimental setup: Experiments use editing datasets pairing each edit input with a plausible target label, including question answering, text generation, and fact-checking settings.Targets mimic practical edits such as changing a QA answer or steering a generative continuation.
  • Experimental setup: Edit success measures reliability and generality, while drawdown measures locality through perplexity increases or accuracy decreases relative to the base model.The drawdown metric depends on the task’s underlying performance measure.
  • Very large models: MEND provides the most successful edits on GPT-Neo (2.7B), GPT-J (6B), T5-XL (2.8B), and T5-XXL (11B) across the evaluated tasks.The large-model experiments compare MEND with existing editors on Wikitext generative editing and zsRE question answering.
  • Small-scale models: At small scale, ENN and MEND perform most consistently across FEVER, zsRE, and Wikitext, with ENN exceeding MEND on Wikitext.The evaluated BERT-base, BART-base, and distilGPT-2 models are 1–2 orders of magnitude smaller than the very large models.
  • Multiple simultaneous edits: MEND consistently provides significantly more effective edits than ENN when applying 1, 5, 25, 75, or 125 simultaneous zsRE edits.MEND forms simultaneous edits by summing separately computed parameter edits and applies 25 edits in a single model update.
  • Ablations: Removing identity initialization or input normalization significantly lowers editing performance, while lower-dimensional MEND variants retain competitive performance.The ablations evaluate parameter sharing, initialization, normalization, and variants that compute pseudoactivations, pseudodeltas, or whichever is lower-dimensional.

6 DISCUSSION

MEND is presented as an efficient approach for editing very large neural networks, successfully editing the largest publicly available GPT and T5 Transformer models. The discussion identifies insufficiently challenging locality examples as a limitation and suggests alternative locality losses or harder negative mining.

  • Conclusion: MEND efficiently edits very large neural networks with 10 billion+ parameters.The method is called Model Editor Networks with Gradient Decomposition.
  • Conclusion: MEND successfully edits the largest publicly available Transformer models from the GPT and T5 families.The conclusion describes MEND as the only method achieving this result.
  • Limitations & Future Work: Existing model editors, including MEND, may over-generalize because locality examples are not challenging enough.This can cause changed outputs for distinct but related inputs.
  • Limitations & Future Work: Alternative locality losses or harder negative mining may help address over-generalization.These are proposed as ways to improve the enforcement of edit locality.

ETHICS STATEMENT

The work acknowledges that internet-scraped training data may encode or produce content conflicting with ICLR community values. It presents model editing as a possible mitigation for undesirable behaviors but also notes potential misuse.

  • ETHICS STATEMENT: Internet-scraped training corpora may contain or produce content counter to community values, while model editing could help maintainers change undesirable behaviors.The passage also flags that model editors could potentially be misused.

REPRODUCIBILITY … B.1 EDITABLE NEURAL NETWORKS (ENN)

The paper supports reproducibility through detailed algorithmic and experimental documentation, fixed random seeds, and access to pre-existing datasets and models. It explains MEND’s identity initialization and normalization, contrasts model editing with continual learning and ENN, and discusses ENN’s memory costs and intrinsic editability.

  • REPRODUCIBILITY: Reproducibility is supported through detailed algorithm, setup, hyperparameter, and implementation descriptions, fixed random seeds, and instructions for obtaining datasets and models.The descriptions appear in Sections 3 and C, while datasets and models come from De Cao et al. (2021).
  • A EFFECTIVE INITIALIZATION AND NORMALIZATION FOR MEND: MEND uses a residual connection and partially random, partially zero initialization to implement the identity function.This initialization preserves the raw fine-tuning gradient as a useful editing starting point and was motivated by ablations showing random initialization produced less effective edits.
  • A EFFECTIVE INITIALIZATION AND NORMALIZATION FOR MEND: Edit inputs and gradient updates are normalized to zero mean and unit variance using statistics computed over the edit training set and sequence index.The passage defines normalized ū_l and δ̄_l+1 in terms of these statistics.
  • B EXTENDED DISCUSSION OF RELATED WORK: Model editing updates behavior while retaining old information and unrelated behaviors, whereas continual learning typically introduces wholly new behaviors or datasets.Both areas address avoiding catastrophic forgetting, but their expected learning objectives differ.
  • B.1 EDITABLE NEURAL NETWORKS (ENN): ENN jointly optimizes base-task performance, edit reliability on edit examples, and locality relative to the pre-edit model on unrelated data.Its objective is L_ENN(θ, D_base, D_edit, D_loc) = L_base(D_base, θ) + c_edit · L_edit(D_edit, θ′) + c_loc · L_loc(D_loc, θ, θ′).
  • B.1 EDITABLE NEURAL NETWORKS (ENN): ENN encodes editability intrinsically in the model parameters, whereas MEND supplies extrinsic editability through learned parameters independent of the edited model.ENN adds no editability parameters, while extrinsic editors can potentially be applied across multiple base models.
  • B.1 EDITABLE NEURAL NETWORKS (ENN): ENN requires an extra copy of the original model and additional activations and gradients, producing substantial VRAM costs for very large models.Figure 4 reports that ENN memory grows much more rapidly than MEND and KE, preventing single-GPU estimation for T5-11B and GPT-J.

B.2 KNOWLEDGEEDITOR (KE)

KNOWLEDGEEDITOR (KE) is a hypernetwork-based extrinsic editor that generates structured gradient updates from edit examples. It is contrasted with MEND through its direct use of the raw edit example and rank-1 update components, and evaluated on FEVER and zsRE editing examples.

  • KE approach: KE is an RNN hypernetwork that conditions on the input, incorrect output, and desired label to generate masks, offsets, and gradient scaling factors.For transformer weight matrix W_i, it outputs m_i, b_i, and α, with the update θ′ = θ − α(m_i ⊙ ∇W_i) + b_i.
  • Comparing KE and MEND: Unlike MEND, KE maps the raw edit example to a single rank-1 mask and rank-1 offset over the fine-tuning gradient.The passage frames this distinction as motivating hypothesized challenges that MEND avoids.
  • Editing datasets: FEVER locality loss uses KL divergence between Bernoulli distributions from the pre-edit and post-edit models for a locality example.This describes the fact-checking editing dataset example.
  • Editing datasets: zsRE locality is approximated by computing token-level KL divergence at the labeled answer tokens because evaluating all possible answers is computationally expensive.The example uses the label “EastEnders” for this approximation.

C EXPERIMENTAL DETAILS … C.3 ENVIRONMENT DETAILS

The experiments edit selected MLP weights, use tuned algorithm-specific hyperparameters, approximate locality constraints when exact KL computation is costly, and run on single GPUs. Environment choices include bfloat16 for fitting T5-11B, whose full-precision parameters exhaust the largest GPU’s memory.

  • C EXPERIMENTAL DETAILS: GPT and BERT-style models edit MLP weights in the last 3 transformer blocks, while BART and T5-style models edit encoder-decoder MLPs in the last 2 blocks.This corresponds to 6 weight matrices for GPT/BERT-style models and 8 for BART/T5-style models.
  • C EXPERIMENTAL DETAILS: MLP-layer editing generally performs better across algorithms than attention-layer editing.
  • C.1 HYPERPARAMETERS: Fine-tuning uses model-dependent learning rates: 5e-6 for GPT-Neo and GPT-J, 1e-4 for T5, and 1e-6 for smaller models.Learning rates that are too large decrease locality, whereas rates that are too small slow edits.
  • C.1 HYPERPARAMETERS: ENN starts with an inner-loop learning rate of 1e-2, learns it in the outer loop, fine-tunes all parameters, and uses one inner-loop update step.The learned learning rate improves performance over the fixed version, while the single update is used for computational reasons.
  • C.1 HYPERPARAMETERS: ENN’s second loss term emphasizes reducing the largest incorrect logit because this work uses only a single gradient-step edit.
  • C.2 COMPUTING THE LOCALITY CONSTRAINT: Locality KL-divergence is approximated with dataset samples when model outputs exceed a few tokens; seq2seq evaluation considers answer tokens, while FEVER uses exact Bernoulli KL.
  • C.3 ENVIRONMENT DETAILS: All runs use a single NVIDIA RTX Titan or A40 GPU without gradient checkpointing or memory-reduction optimizations, with bfloat16 used to fit T5-11B.In full precision, T5-11B’s parameters alone use all memory on the largest GPU; MEND and KE VRAM is estimated by doubling bfloat16 usage.

C.4 DATASET CONSTRUCTION & EXAMPLES · D RANK-1 GRADIENT FOR MLPS

The paper constructs realistic edit datasets from input–label pairs, including plausible labels that may intentionally differ from factual answers, and uses locality examples to constrain model behavior. For an MLP with batch size 1, it derives the weight gradient as a rank-1 outer product using standard backpropagation quantities.

  • C.4 DATASET CONSTRUCTION & EXAMPLES: Edit datasets pair an edit input x_e with a plausible edit label y_e, which need not be factually correct.The example assigns Addis Ababa to a question about Ursula K. Le Guin, despite her actual birthplace being Berkeley, California.
  • C.4 DATASET CONSTRUCTION & EXAMPLES: The locality-sample distribution approximates the model’s predictive distribution to reduce training computation and implementation complexity.The passage notes that generally low drawdown for most models suggests this approximation provides good locality control.
  • C.4 DATASET CONSTRUCTION & EXAMPLES: Training examples include edit and locality inputs, with locality samples constraining pre- and post-edit predictive distributions to remain similar at every token.Table 8 identifies x_loc as the locality example and y_e and y′_e as edit labels.
  • D RANK-1 GRADIENT FOR MLPS: For an MLP with batch size 1, the loss gradient with respect to layer ℓ’s weight matrix W_ℓ is represented using the layer input u_ℓ and next-layer pre-activation gradient δ_ℓ+1.The pre-activations are defined as z_ℓ+1 = W_ℓu_ℓ, with δ_ℓ+1 assumed pre-computed by standard backpropagation.
  • D RANK-1 GRADIENT FOR MLPS: The derivative of the loss with respect to an individual weight W_ij is obtained by applying the chain rule.The derivation relates the weight derivative to the derivative with respect to the next-layer pre-activation.
  • D RANK-1 GRADIENT FOR MLPS: The derivation’s second equality uses the relationship between layer ℓ weights and next-layer pre-activations, yielding the stated rank-1 gradient identity.The final step concludes the identity introduced for the gradient with respect to W_ℓ.

E EDITING ATTENTION PARAMETERS · F ADDITIONAL QUALITATIVE EXAMPLES OF MEND · G EDITING THROUGH CACHING

MEND edits transformer MLP parameters more effectively than attention parameters, while qualitative examples reveal undergeneralization and overgeneralization failures. A caching-based editor works on simpler tasks but is threshold-sensitive and fails on longer, more complex Wikitext contexts.

  • E EDITING ATTENTION PARAMETERS: Editing attention query, key, value, and output matrices generally reduces performance compared with editing MLP matrices.The comparison uses the same transformer blocks as the main experiment but replaces the two MLP matrices with attention matrices.
  • E EDITING ATTENTION PARAMETERS: Attention-parameter editing consistently lowers generative-model drawdown and edit success, as well as T5 sequence-to-sequence edit success.
  • F ADDITIONAL QUALITATIVE EXAMPLES OF MEND: MEND qualitative examples on a 770M-parameter T5-large model include correct generalization, undergeneralization, and overgeneralization.The examples use a model fine-tuned on Natural Questions; one example is counted correct despite an incorrectly generated token.
  • F ADDITIONAL QUALITATIVE EXAMPLES OF MEND: Undergeneralization occurs when the edit example changes correctly but other examples in its equivalence neighborhood remain unchanged.
  • F ADDITIONAL QUALITATIVE EXAMPLES OF MEND: Overgeneralization occurs when superficially similar but semantically different inputs receive the edited output.The paper highlights this as a failure case in the additional qualitative examples.
  • G EDITING THROUGH CACHING: Caching performs slightly better on zsRE, slightly worse on FEVER, and completely fails on Wikitext editing.The comparison selects task-specific thresholds producing drawdown similar to MEND; Wikitext failure is attributed to longer, more complex contexts.
  • G EDITING THROUGH CACHING: A caching editor replaces normal predictions with the edit label when a new input’s hidden state lies within threshold ϵ of the edit example’s hidden state.The cached hidden state is the sequence-length-averaged final-layer representation of the edit input and corresponding label tokens.
  • G EDITING THROUGH CACHING: Caching is sensitive to its threshold, and a value effective for one task may perform poorly on another.Representations are averaged over question tokens for zsRE, fact-statement tokens for FEVER, and previously seen tokens at each time step for generative modeling.
Loading 2110.11309v2…