Memory-Based Model Editing at Scale

In plain words

A trained language model is hard to correct after training. Fixing one wrong answer usually means changing its weights, and that can quietly break answers that used to be right. This method, called SERAC, leaves the big model alone. It keeps a list of corrections as plain sentences next to the model.

For each new question, a small helper asks whether the question falls under one of the stored corrections. If none applies, the original model answers as before. If one does, a second small model reads the correction and the question together and gives the answer the correction implies. Think of a sticky note on the side of a textbook, plus a helper who decides when the note is relevant.

The appeal is that corrections do not pile up inside the weights, so new ones can be added without retraining the big model. The authors also introduce test sets where a correction must carry over to related questions, not only to the exact wording. The reported results are the authors' own.

TL;DR

SERAC is a memory-based model editor. It stores each edit as text in a cache and wraps a frozen base model with a scope classifier and a counterfactual model. A classifier decides whether an input falls under an edit's scope; if so, the counterfactual model answers, otherwise the base model does §3.1 eq. 1. The authors report the highest edit success on all four benchmarks with 10 simultaneous edits (5 on ConvSent), with a drawdown cost on FC tab. 3. The numbers are author-reported on the paper's own new benchmarks.

Problem

Deployed models are static, and a correction usually needs retraining or a parameter update §1. Existing learned model editors predict weight updates from gradients. The authors argue that gradients are a weak signal for deciding which inputs an edit should affect, and that these editors degrade as edits accumulate and must be refit for each base model §1 §3 §5.2. The editing problem is defined as changing behaviour on a set of in-scope inputs while leaving out-of-scope inputs unchanged, with "hard" cases on both sides: in-scope inputs that need reasoning, and out-of-scope inputs that look similar to the edit §2.

Method

Overview. The wrapper has three pieces: a store of edit descriptors, a relevance classifier, and a second model that answers in place of the base model whenever an edit applies §3.1. The base model's parameters are not changed during training or editing §3. An edit descriptor is either an input-output pair or a free-text description such as a topic and a sentiment §2.

Step 1. Store the edit. Each edit is appended to the cache as its descriptor. Nothing is written into the base model §3.1. The authors say the memory costs about 3 KB per edit §5.2. Applying k edits in sequence or as one batch gives the same model, because both append to the same memory §5.1.

Step 2. Score scope. The scope classifier gives, for the input and each cached edit, a probability that the input lies in that edit's scope §3.1. The default ("Embed") encoder embeds the edit and the input separately and uses the negative squared Euclidean distance between the embeddings as the log likelihood §3.1. A full cross-attention variant ("Cross") encodes each edit-input pair jointly; the authors use it for FC §5.1 §5.2.

Step 3. Route. Take the best-scoring edit, with score β. If β is at least 0.5, the counterfactual model answers from the edit and the input. Otherwise the base model answers. An empty cache always uses the base model eq. 1 §3.1. Retrieving only the single best edit limits interference between edits §3.1.

# SERAC inference: route one input to base or counterfactual model  [eq. 1]
# Input:  input x; edit cache edits (list of edit descriptors)
# Output: prediction for x

def serac_predict(x, edits, base, scope, counterfactual):
    # With no stored edits the base model answers  [§3.1]
    if len(edits) == 0:
        return base(x)
    # Scope probability of x under each cached edit  [§3.1]
    scores = [scope(z_e, x) for z_e in edits]
    # Retrieve only the most relevant edit  [§3.1] [eq. 1]
    best = argmax(scores)
    # Override the base model when the best score is at least 0.5  [eq. 1]
    if scores[best] >= 0.5:
        return counterfactual(edits[best], x)
    return base(x)

Step 4. Score scope with embeddings (default). Pseudocode for the default scope function, as described in the text §3.1. The paper names distilBERT as the encoder in its runs app. C.

# Embedding-based scope probability  [§3.1] [app. C]
# Input:  edit descriptor z_e (text); input x (text)
# Output: scalar in (0, 1]

def scope_score(z_e, x):
    # Fixed-length embedding of the edit descriptor  [§3.1]
    edit_vector = embed(z_e)              # (D,), D not stated
    # Fixed-length embedding of the input, same encoder  [§3.1]
    input_vector = embed(x)               # (D,)
    # Negative squared distance is the predicted log likelihood  [§3.1]
    log_likelihood = -sum((edit_vector - input_vector) ** 2)
    return exp(log_likelihood)

Step 5. Train the scope classifier. Each training edit has one in-scope input (label 1) and one out-of-scope input (label 0). The loss is the binary cross-entropy averaged over the edit dataset §3.2 eq. 2. In-scope inputs come from the sampler I and out-of-scope inputs from O. For FC, the out-of-scope input is a claim from another page, or the claim itself when the evidence contradicts it §4.

\ell(\phi) = -\,\mathbb{E}\big[\log g_\phi(z_e, x_{in}) + \log(1 - g_\phi(z_e, x_{out}))\big] eq. 2

Step 6. Train the counterfactual model. The counterfactual model receives the concatenated edit and input and is trained by maximum likelihood on the in-scope label sequence §3.2 §3.1 eq. 3. The two components are trained separately §3.2.

\ell(\psi) = -\,\mathbb{E}\big[\log p_\psi(y_{in} \mid z_e, x_{in})\big] eq. 3

# Training losses for the two learned components  [§3.2] [eq. 2] [eq. 3]
# Input:  one edit z_e with in-scope pair (x_in, y_in) and out-of-scope x_out

def scope_loss(z_e, x_in, x_out):
    # In-scope pair has label 1  [§3.2] [eq. 2]
    positive = -log(scope(z_e, x_in))
    # Out-of-scope pair has label 0  [§3.2] [eq. 2]
    negative = -log(1 - scope(z_e, x_out))
    return positive + negative

def counterfactual_loss(z_e, x_in, y_in):
    # Negative log likelihood of the in-scope label sequence  [§3.2] [eq. 3]
    return -counterfactual.log_prob(y_in, z_e, x_in)

Evaluation metrics §4. For classification-style tasks, edit success (ES) is the exact-match rate on in-scope inputs and drawdown (DD) is the rate at which out-of-scope outputs change from the base model eq. 4 eq. 5. For ConvSent, ES multiplies a sentiment term and an on-topic term, both built from average per-token log likelihoods of pre-generated responses §4. DD becomes a KL divergence between base and edited output distributions eq. 7.

ES_{ex}(z_e) = \mathbb{E}_{x_{in}}\,\mathbb{1}\{f_e(x_{in}) = y_{in}\} eq. 4

DD_{ex}(z_e, O) = \mathbb{E}_{x_{out}}\,\mathbb{1}\{f_e(x_{out}) \neq f_{base}(x_{out})\} eq. 5

ES_{sent} = z_{sent} \cdot z_{topic} eq. 6

DD_{sent}(z_e, O) = \mathbb{E}_{x_{out}}\,\mathrm{KL}\big(p_{base}(\cdot \mid x_{out}) \,\|\, p_e(\cdot \mid x_{out})\big) eq. 7

Hyperparameters

Results

Author-reported. Setup: four settings (QA, QA-hard, FC, ConvSent) with the base models above, 10 simultaneous edits (5 on ConvSent), and five baselines: fine-tuning (FT), cache and lookup (LU), MEND, ENN and retrieve-and-prompt (RP) §5.1 tab. 3. Hardware is not stated.

Table 3: comparison across editing problems tab. 3. Bold as in the paper: best value in a row or within 1% of it.

Table 4: component breakdown on QA-hard and FC tab. 4. Cls acc. is the scope classifier's accuracy; rep acc. is the counterfactual model's accuracy on the inputs routed to it.

FC does not split in-scope inputs by difficulty, and does not label out-of-scope inputs, so those cells are pooled or omitted tab. 4.

Table 5: scope classifier architecture on QA-hard and FC with 10 edits tab. 5. Embed is the embedding variant; Cross is the cross-attention variant; D is distilBERT and B is BERT-base.

Table 6: wall-clock time and memory for one forward pass after 10 edits tab. 6. "out" is an input routed to the base model; "in" is routed to the counterfactual model.

The memory increase comes mostly from the classifier and counterfactual weights, a fixed cost that does not grow with the number of edits; the cache itself is small §5.2.

Other results in the text.

• Many edits (QA, T5-large): SERAC keeps near-constant ES − DD as edits grow to 75, while ENN and MEND reach failure there §5.1 fig. 3.
• Re-use across base models: on QA, one editor gives ES above 0.99 and DD below 0.01 across T5 sizes; on QA-hard, ES above 0.92 and DD below 0.03 §5.2.
• Misuse: the authors note that the ConvSent setting could let an adversary steer a dialogue model's opinions §7 app. A.

Paper vs code

Official code: github.com/eric-mitchell/serac (MIT, copyright Eric Mitchell), linked from the authors' project page, which the abstract names. The repository owner is the first author. The paper's SERAC appears to be the rep algorithm in the code, inferred from its scope classifier plus replacement model, and run with +alg=rep main.sh main.sh main.sh main.sh.

Discrepancies, each with both citations:

• Scope classifier loss is class-balanced in code. The paper averages the binary cross-entropy over pairs eq. 2. Code takes the mean loss over positives and the mean loss over negatives and adds them losses.py trainer.py. Negatives are paired with their own edit only, unless use_all_negatives is set, which defaults to false trainer.py rep.yaml.
• Training weights differ. The paper trains the two components separately §3.2. Code adds them into one objective, 0.1 times the counterfactual loss plus 1.0 times the scope loss, and steps one Adam optimiser over both parameter sets trainer.py config.yaml config.yaml trainer.py. Adam largely cancels a constant scale on each group's gradient, so the effect is probably small; this was not tested.
• Hard negatives are sampled differently. Appendix D.1 says: top 100 neighbours, top 50 ignored, sampled uniformly app. D.1. Code defaults to top 100 with the first 25 excluded config.yaml. It keeps 75 candidates utils.py, weights them by exp(similarity / 0.1) utils.py and normalises the weights utils.py. Hard negatives are drawn with probability 0.5 by default config.yaml utils.py. The QA-hard runs enable hard negatives and do not override the exclusion main.sh.
• Batch sizes match the text for ConvSent, not the Table 3 caption. The scripts use batch size 10 for FC and QA and 5 for ConvSent main.sh main.sh main.sh, which matches k=10 and k=5 in the text §5.1.
• FC cross-attention backbone. The FC script enables cross-attention without changing the backbone, and the default is distilBERT main.sh rep.yaml. The paper's appendix lists BERT-base as the FC counterfactual, not as the scope classifier app. C. The backbone behind Table 3's FC number is not stated.
• Fixed scale in the embedding variant. The code multiplies the squared distance by a scale that is a buffer fixed at 1.0, not a learned parameter rep.py. The paper does not mention a scale.
• Embedding and cross-attention details the paper omits. The embedding is the [CLS] position of the last hidden layer rep.py. The cross-attention input is the edit, a separator token and the input, scored by log-softmax at index 0 rep.py. The paper does not state which class index means "in scope".
• Stopping rule and step count. Not in the paper. Code defaults are max_iters 1,000,000, validation every 5,000 steps, early-stopping patience 40,000 on the validation supervised loss config.yaml config.yaml config.yaml trainer.py. The scripts do not override them.
• Gradient accumulation. Code accumulates over 10 batches by default config.yaml. The paper does not mention accumulation.

Pitfalls

• Table 3 caption and text disagree on k. The caption says all problems use k=10; the text says ConvSent uses k=5 tab. 3 §5.1. The code and scripts use 5 for ConvSent. Implement k=5 for ConvSent.
• Table 3 does not match Table 5 row by row. Table 3's QA-hard ES (0.913) and FC ES (0.877) match neither the Embed-D/Embed-B nor the Cross-D/Cross-B rows in Table 5 tab. 3 tab. 5. The paper does not say which configuration produced Table 3. Table 3 also gives ES 0.986 for QA, while the re-use paragraph says "above 0.99" for QA §5.2.
• Hard-negative exclusion. Paper: 50 ignored; code default: 25. Treat the code as the reference if you reproduce the released experiments; the paper is the reference for the stated method. Sampling is uniform in the paper and similarity-weighted in code.
• Counterfactual model differs per setting. T5-small for QA app. C, bert-base-uncased for FC app. C, BlenderBot-small-90M for ConvSent app. C. Do not reuse one across settings without re-running the paper's setup.
• Base checkpoint naming. The main text says the QA base model is "fine-tuned on Natural Questions" §4. The appendix gives google/t5-large-ssm-nq, described as pretrained on NQ app. C. Use the checkpoint named in code t5large.yaml.
• Generator and classifier names in the ConvSent footnote. A footnote in §4 says BlenderBot 2.0 generated responses; the appendix names facebook/blenderbot-3B app. D.2. The sentiment classifier is described as RoBERTa in the footnote and RoBERTa-large in the appendix §4 app. D.2.
• Unlikelihood term. ConvSent training adds an unlikelihood term on incorrect-sentiment responses, but no equation is given §4. Code enables it by default config.yaml.
• Scope classifier threshold is fixed. Routing uses 0.5 eq. 1 rep.py. The paper does not calibrate it per task.
• Overhead is mostly fixed, and FC is slow. Extra memory is mostly the classifier and counterfactual weights: about 0.5 GB over the base model for QA and about 0.7 GB for FC tab. 6 §5.2. The FC scope pass takes 19 ms against 7 ms for the base model, because of the cross-attention classifier tab. 6 §5.2. The authors say the edit memory has no bound in continual editing and propose periodic distillation as future work §7.
• Figures and sources not in the paper. The source tree contains a multi-model figure marked "placeholder; numbers are made up" that the paper does not include, and an older draft in history/. Neither is used here.
• Appendix letters. Appendix A, B, C and D are inferred from the order of the appendix inputs in the source; the compiled paper was not built.

Lineage

• Builds on: fine-tuning, Editable Neural Networks (Sinitsin et al. 2020), Model Editor Networks with gradient decomposition (MEND, Mitchell et al. 2021), and the cache-and-lookup baseline (LU) from the MEND paper §1 app. B. Dense passage retrieval for the embedding scope classifier §3.1.
• Compared against: FT, LU, MEND, ENN and retrieve-and-prompt (RP, the ablation that keeps the scope classifier but uses the base model to answer) §5 §5.1.
• Related: BeliefBank for consistency with an external memory §6; the editing papers of Cao et al. 2021 and Hase et al. 2021 used for evaluation setup §4.
• Superseded by: not stated in this wiki.