Fast Model Editing at Scale

In plain words

A large language model sometimes gives a wrong answer, and you want to fix that one answer without retraining the model or breaking its other behaviour. The obvious fix is to fine-tune the model on the corrected example, but a single example usually makes the model overfit: the answer changes, yet rephrasings of the question and unrelated questions can change too.

The paper's idea is to train a few small helper networks ahead of time. Each helper watches the gradient that a fine-tuning step would produce for one layer, and outputs a more targeted change to that layer's weights. A gradient for one layer is a big grid of numbers, but it can be built from two short vectors, so the helpers only have to work with those short vectors. Once the helpers are trained, a new correction takes one step: compute the gradient for the wrong input and the desired output, pass it through the helpers, and apply the result. The paper reports this works on models with billions of parameters.

TL;DR

MEND (Model Editor Networks using Gradient Decomposition) is a small set of learned networks, one per shape of edited weight matrix. They map each layer's fine-tuning gradient to a low-rank weight edit, so a single wrong input and desired output can be fixed in one step. The paper reports the highest edit success of the methods it ran at each of the four large sizes (GPT-Neo, GPT-J, T5-XL, T5-XXL). On drawdown it is not always lowest: FT + KL is lower on GPT-Neo. These numbers are author-reported. Read the Pitfalls section before reimplementing: several details are missing from the paper, and the code differs in places.

Problem

Fine-tuning on one example overfits, so it damages locality (unrelated outputs change) and generality (rephrasings of the edited question do not change) §1. Fine-tuning with a KL penalty on unrelated data helps locality but still generalises poorly, and it needs the training set at edit time §1. Learned editors such as Editable Neural Networks (ENN) need to train and keep a copy of the base model, which is too expensive at 10B+ parameters app. B. KnowledgeEditor (KE) is cheap but failed in the paper's large-model experiments §4.1.

The paper scores editors on edit success, the average accuracy on the edit input plus its equivalence neighbourhood eq. 1, §2, and on drawdown, the loss of accuracy or increase of perplexity on unrelated data §4.

Method

Rank-1 structure of the gradient

For a linear layer with weight W, input u and pre-activation output gradient δ, the loss gradient for one example is the outer product δ uᵀ. For a batch or a sequence, the gradient is the sum of such outer products over the elements §3.1, app. D. MEND works on the factors (u, δ), not on the full weight gradient. This needs only 2d inputs per layer rather than d² §3.1.

Editor network

For each layer, a network g takes z = concat(u, δ) and returns pseudo-activations ũ and pseudo-deltas δ̃ of the same sizes §3.1. The edit is the sum of their outer products over the elements eq. 2. The final weights are W̃ = W − α ∇̃, where α is a learned per-layer scalar step size §3.1. The network has one hidden block with a low-rank residual form, and two residual blocks, both initialised to the identity eq. 3a, 3b, fig. 2, §3.1, app. A. Parameters are shared across layers of the same weight shape, with a learned per-layer scale and offset in each block, similar to FiLM §3.1.

# MEND editor network g for one layer  [eq. 3a, 3b] [§3.1]
# z: concat of normalised u and delta, shape (B, m + n)
# Block 1 and block 2 each use low-rank weights U_j V_j and a ReLU.

def mend_network(z, scale_1, offset_1, scale_2, offset_2):
    # First block: residual plus a ReLU of the low-rank map  [eq. 3a]
    # The bias b appears in this block only, as written in the paper.
    hidden = z + relu(
        scale_1 * (U_1 @ (V_1 @ z) + bias) + offset_1
    )
    # Second block: residual plus a ReLU of the low-rank map  [eq. 3b]
    return hidden + relu(
        scale_2 * (U_2 @ (V_2 @ hidden)) + offset_2
    )

The output is split into ũ and δ̃, and the layer edit is their sum of outer products over the elements eq. 2, §3.1. Input normalisation and the identity initialisation are needed for training to work; the ablation shows both matter app. A, §4.4.

Editing

The edit procedure runs one forward pass on the edit pair, caches each edited layer's input, computes the negative log-likelihood, takes each layer's gradient with respect to its output, runs it through the editor network, and forms the new weights alg. 2. The procedure is a single update with no iteration alg. 2.

# MEND edit procedure  [alg. 2]
# Input: model, editable weights W, editor params phi, edit pair (x_e, y_e)
# Output: edited weights

def edit(model, editable_weights, phi, x_e, y_e):
    # Forward pass on the edit pair; cache each layer's input u  [alg. 2]
    loss = negative_log_likelihood(model, x_e, y_e)  # [§2]
    for layer in editable_weights:
        # Gradient of the loss with respect to the layer's output  [alg. 2]
        delta = grad_wrt_output(loss, layer)
        # Map the factors (u, delta) through the editor  [alg. 2] [§3.1]
        u_tilde, delta_tilde = editor_g(phi, layer, u, delta)
        # Sum of outer products over elements gives the layer edit  [eq. 2]
        edit_matrix = sum_outer(delta_tilde, u_tilde)  # (m, n)
        # The paper subtracts the edit; the code adds it. See Pitfalls.
        layer.weight = layer.weight - edit_matrix  # [alg. 2]
    return edited_weights

Training

The editor parameters are trained on an edit dataset of tuples (x_e, y_e, x_e', y_e', x_loc). Each step samples one tuple, performs the edit, and computes two losses. The edit-success loss is the negative log-likelihood of the equivalence example's target under the edited model eq. 4a. The locality loss is the KL divergence between the unedited and the edited model on the unrelated locality input eq. 4b. The total is c_e L_e + L_loc, with c_e = 0.1 in all experiments §3.2, alg. 1. The editor parameters are updated with Adam; the base model is not updated §3.2, alg. 1. No higher-order gradients are computed, because the base parameters are not optimised §3.2.

# MEND training loop  [alg. 1]
# Input: pre-trained model, editable weights W, editor phi_0, edit data D_tr
# Output: trained editor parameters phi

def train_mend(model, editable_weights, phi, edit_data, c_edit):
    # Repeat for each training step t  [alg. 1]
    for step in count(1):
        # Sample edit, equivalence and locality examples  [alg. 1]
        x_e, y_e, x_eq, y_eq, x_loc = sample(edit_data)
        # Apply the edit with the current editor parameters  [alg. 1] [alg. 2]
        edited = edit(model, editable_weights, phi, x_e, y_e)
        # Edit-success loss on the equivalence example  [eq. 4a]
        loss_edit = -log_prob(edited, y_eq, given=x_eq)
        # Locality loss: KL between unedited and edited outputs  [eq. 4b]
        loss_loc = kl(model(x_loc), edited(x_loc))
        # Combined objective with weight c_e on the edit loss  [§3.2]
        loss = c_edit * loss_edit + loss_loc
        # Adam update of the editor parameters only  [alg. 1] [§3.2]
        phi = adam_step(phi, grad(loss, phi))

The locality KL is approximated on samples from the dataset rather than on the model's own outputs. For seq2seq question answering it is averaged over the answer tokens; for Wikitext it is averaged over all tokens; for FEVER it is the exact Bernoulli KL app. C.2.

Hyperparameters

Values the paper states are cited to the paper. Values marked code are in the official repository (commit in the frontmatter) and are not stated in the paper.

Results

Author-reported throughout. The main comparison is on large models (GPT-Neo 2.7B, GPT-J 6B, T5-XL 2.8B, T5-XXL 11B). ES is edit success (higher is better); DD is drawdown (lower is better) §4.1. ENN could not be run at this size because of its memory use §4.1, large-model table.

Small models, one edit at a time §4.2, small-model table:

Batched edits on zsRE with BART-base, MEND summed over edits §4.3, batched-edits table:

Ablations on distilGPT-2 Wikitext and BART-base zsRE §4.4, ablation table. Removing input normalisation (ES 0.02) or identity initialisation (ES 0.27) makes Wikitext editing fail. Sharing parameters across editors costs little on zsRE. Editing only the smaller of ũ and δ̃ keeps ES 0.80 on Wikitext and 0.98 on zsRE, which the authors suggest matters for scaling. The no-sharing variant has ES 0.86 but drawdown 0.195 against 0.225 for MEND, so it is slightly better on drawdown. Editing attention matrices instead of MLPs lowers ES on all four large models app. E.

Caching the edit example's hidden state is competitive on zsRE and FEVER but fails on Wikitext, and its threshold needs tuning per task app. G, caching table.

Paper vs code

Official code: eric-mitchell/mend (MIT), commit e04fdb9cc784188906feffeb171025872933a5a8. The repository is linked from the paper's project page; its owner is an author. The code matches the paper on: the loss weights c_e = 0.1 and c_loc = 1 §3.2 config.yaml; Adam trainer.py; the edit loss and locality loss trainer.py; the early-stopping patience and validation size config.yaml, config.yaml; the identity initialisation of U nn.py; the ReLU activation and single hidden layer mend.yaml, mend.yaml; and editing the MLP blocks listed in the config files distilgpt2.yaml, bart-base.yaml, gptj.yaml, t5xxl.yaml.

Differences:

• Update sign. The paper writes W̃ = W − α∇̃ §3.1 and W̃ = W − δ̃ũᵀ alg. 2. The code computes the loss as a positive negative log-likelihood, takes its gradient, and sets the new parameter to p + lr·g mend.py, mend.py, mend.py, mend.py. Unless the learned step sizes become negative, this moves against the gradient's descent direction. Check the sign on a small case before using it.
• Rank. The paper does not state the rank. The code uses rank 1920, capped by the factor dimension mend.yaml, nn.py. The paper says MEND keeps its parameters O(d) §3.1 but its ablation table gives O((m+n)²) §4.4, ablation table. These statements are inconsistent with each other; the code's rank cap determines the real size.
• V initialisation. The paper says Xavier uniform app. A. The code uses a standard normal for V in the identity case nn.py.
• Bias. Equation 3b has no bias term eq. 3b. The code adds a learned, zero-initialised bias to every low-rank layer, including the second block nn.py, nn.py.
• Step size. The paper uses one learned scalar per layer §3.1. The code learns one step size per edited matrix, initialised to 1e-4 mend.py.
• Normalisation statistics. The paper uses the edit training set app. A. The code keeps running statistics updated during every training forward pass, freezes them in evaluation mode, skips batch elements whose gradient is zero, and needs at least two samples mend.py, mend.py.
• Batch size. The paper says a batch size of 10 with gradient accumulation app. C. The code default is batch_size 2 with accumulate_bs 10 config.yaml, config.yaml. It is unclear whether the paper's 10 means accumulated edits per step or per forward pass.
• Maximum steps. The paper says 500,000 app. C; the code default is 1,000,000 config.yaml.
• T5 matrices. The code edits DenseReluDense.wi_0 and wo in each T5 block, not wi_1 t5xxl.yaml. The paper says "MLP weight matrices" app. C and does not name them.
• In-place edit. edit makes the base model's parameters in place and returns a wrapper around the same model, unless detach_history is set mend.py, mend.py. A caller who keeps the original reference will see the edited weights.

Pitfalls

• Sign of the update: see Paper vs code. This is the first thing to check before reproducing the edits.
• Rank and size: the paper does not give the rank. Use the code's cap (1920) if you want to match the reported runs, and count the parameters yourself, since the paper's complexity statements disagree.
• Cost per edit: the edit runs a forward and a backward pass through the base model mend.py. It is one step, but not a table lookup, and it needs the base model in memory with gradient support.
• Normalisation: the first training forward passes must contain non-zero gradients, or the statistics are never set mend.py. Without normalisation and identity initialisation, the paper's ablations show edits fail §4.4.
• Locality is approximated from dataset samples, not from the model's own outputs. The paper calls this less principled than a Monte Carlo estimate from model samples, but says it is cheaper and works in practice app. C.2.
• Generality is tested only with paraphrases from backtranslation, and with a Wikitext prefix-truncation heuristic. The paper says these do not test whether an edit implies other facts §5.
• Overgeneralisation: the paper reports that edits sometimes change outputs for related but different inputs; the locality examples are not hard enough to prevent it §5, app. F.
• The "only approach that edits models above 10B" claim §1, §5 is weaker in the table than in the abstract. FT reaches 0.87 edit success on T5-XXL with drawdown under 0.001 §4.1, large-model table, against MEND's 0.89. The clearer separation is on GPT-Neo and GPT-J, where KE fails.
• Batched-edit numbers disagree. The text reports 96% edit success and under 1% accuracy loss at 25 edits §4.3. The batched-edit table reports 0.89 and 0.011 §4.3, batched-edits table. Treat the table as the reported result and the text as an error until checked.
• ENN is a re-implementation with a single inner step, not the original app. C.1, and is absent from the large-model comparison because of memory use §4.1.
• VRAM for T5-11B training is estimated by doubling bfloat16 usage, not measured in full precision app. C.3.
• Equations 3a and 3b are written with ⊙ for elementwise scaling; the paper does not say whether the offset o and the scale s are per output dimension. The code applies them per output dimension and per parameter index nn.py, nn.py.

Lineage

• Builds on: Editable Neural Networks (ENN; Sinitsin et al., 2020), a meta-learned editor that changes the base model's weights in training. MEND keeps the editor separate from the base model. Also builds on the KL-constrained fine-tuning baseline from KnowledgeEditor (KE; De Cao et al., 2021), and on the low-rank model-edit observation from LoRA (Hu et al., 2021) §5, §3.1.
• Compared against: FT, FT + KL, ENN and KE, plus a caching editor (app. G).
• Superseded by: not stated in the wiki.