Model Editing Knowledge Editing in Large Language Models
# Model Editing & Knowledge Editing in Large Language Models
## 1. Introduction & Motivation
Large language models memorize an enormous number of facts during pretraining, but that knowledge inevitably becomes stale (a company's CEO changes, a sports record is broken, a piece of legislation is superseded) or is simply wrong (the model learned an inaccurate association from noisy training data). Re-training or even fully fine-tuning a model every time a single fact needs correcting is prohibitively expensive and risks unintended side effects on everything else the model knows. Model editing (also called knowledge editing) is the research area devoted to making small, precise, targeted updates to a trained model's factual associations — changing what it says about one specific fact — while leaving its behavior on everything else essentially untouched.
The field is built on a specific empirical and theoretical premise: factual associations in transformer language models are not smeared uniformly across all parameters, but are substantially localized to identifiable components, particularly the feed-forward (MLP) layers at certain depths, which several lines of interpretability work have characterized as functioning like key-value associative memories. This localization is what makes surgical editing possible at all — if facts were maximally distributed across every parameter, no small, targeted intervention could change one without perturbing everything else. Model editing sits at the intersection of interpretability (understanding where and how knowledge is stored) and practical model maintenance (efficiently keeping a deployed model's knowledge current without the cost of retraining), and its central technical challenge is achieving high efficacy (the edit actually takes effect), generalization (the edit also applies to paraphrases and logically related queries), and specificity (the edit does not corrupt unrelated facts) simultaneously.
## 2. Core Concepts & Theory
The key theoretical foundation, established through causal tracing experiments, is that factual recall in transformer language models tends to be causally concentrated in a small number of mid-layer MLP modules at the token position corresponding to the subject of the fact. Causal tracing works by running the model on a clean prompt (obtaining a correct output), corrupting the subject token's representation (destroying the model's ability to recall the fact), and then selectively "patching" back in individual clean-run activations at different layers to see which restorations recover the correct output — the layers and positions where restoration matters most are identified as causally responsible for storing that fact. This finding motivates a locate-and-edit paradigm: first identify which layer's MLP module is responsible for a given class of factual associations, then directly modify that layer's weights, rather than treating the whole network as an opaque black box to be fine-tuned end-to-end.
The mechanistic view underlying most editing methods treats an MLP layer's weight matrix as an associative (key-value) memory: the first linear layer maps a "key" (a representation encoding the subject and relation of a fact, e.g., "the capital of France is...") to an activation pattern, and the second linear layer maps that activation pattern to a "value" (a representation encoding the answer, "...Paris"). Editing a fact then becomes the problem of modifying this associative memory to map a new or corrected key-value pair correctly, while disturbing the (very large number of) other key-value associations already stored in the same weight matrix as little as possible — a classic constrained-optimization problem: satisfy one new constraint exactly, while minimizing the norm of the resulting perturbation as measured against the existing, known key distribution.
## 3. Mathematical Formulation
Treating a target weight matrix $W \in \mathbb{R}^{d_{ ext{out}} imes d_{ ext{in}}}$ as an associative memory storing key-value pairs $(k_i, v_i)$ such that $Wk_i \approx v_i$ for a large set of existing facts, inserting a new association $(k_*, v_*)$ can be posed as finding a minimal perturbation $\Delta$ such that
$
(W + \Delta) k_* = v_*, \qquad \Delta = \arg\min_{\Delta'} \; \|\Delta'\|_F \;\; ext{subject to} \;\; (W + \Delta') k_* = v_*
$
The closed-form solution to this constrained minimum-norm problem is a rank-one update:
$
\Delta = \frac{(v_* - Wk_*)\, u^ op}{u^ op k_*}, \qquad u = C^{-1} k_*
$
where $C = \mathbb{E}[k k^ op]$ is the (uncentered) second-moment/covariance matrix of keys drawn from the existing fact distribution, estimated empirically from a large sample of the model's own activations on generic text. Using $u = C^{-1}k_*$ rather than simply $u = k_*$ (the naive minimum-Euclidean-norm solution) directs the update preferentially along directions of the key space that are *rarely used* by other stored facts — since $C^{-1}$ amplifies low-variance (rarely-occupied) directions and suppresses high-variance (heavily-used) ones, the resulting update is far less likely to disturb the many other associations that rely on the commonly-used directions. This is exactly the mechanism used by ROME (Rank-One Model Editing).
For editing many facts simultaneously, MEMIT generalizes this to a batch least-squares formulation: given $m$ new key-value pairs $\{(k_j, v_j)\}_{j=1}^m$, solve jointly for a single update $\Delta$ (often additionally distributed across several layers rather than concentrated in one) satisfying all $m$ constraints together,
$
\Delta = \arg\min_{\Delta'} \; \sum_{j=1}^{m} \big\|(W+\Delta')k_j - v_j\big\|_2^2 + \lambda \|\Delta'\|_F^2
$
rather than applying $m$ separate rank-one updates sequentially, which is important because sequential single-fact edits are each computed with respect to the *already-edited* weight matrix and are oblivious to the other pending edits, causing them to interfere with and partially overwrite one another as the edit count grows.
## 4. Advanced Theory & Extensions
Beyond ROME and MEMIT, several extensions target different points in the efficacy/generalization/specificity design space. MEND (Model Editing Networks with Gradient Decomposition) trains an auxiliary hypernetwork to predict a suitable weight update directly from the gradient of a standard fine-tuning loss on the new fact, amortizing the cost of computing edits so that new edits can be applied via a fast forward pass through the hypernetwork rather than solving a fresh optimization problem each time. SERAC and related memory-based approaches take a fundamentally different, non-parametric strategy: rather than modifying the base model's weights at all, they maintain an external, explicit edit memory and a scope classifier that detects when an incoming query falls within the scope of a stored edit, routing such queries to a small counterfactual model while leaving all other queries untouched by the base model — trading editing precision guarantees for added inference-time system complexity.
A significant line of follow-up work has also stress-tested the locate-and-edit hypothesis itself: subsequent studies found that the *causal tracing* localization of a fact (which layer's activations matter most when patched) does not always coincide with the layer whose *weights*, when edited, most effectively and specifically change the fact — a dissociation between "where information is causally read from" and "where an edit is most effectively applied," which complicates the clean story that causal tracing directly identifies the optimal edit target and remains an active area of investigation. Sequential and lifelong editing research further studies how editing methods degrade as thousands of edits are applied over the model's deployed lifetime, since even well-behaved individual edits can compound into significant model degradation at scale if left unaddressed.
## 5. Computational Considerations
Rank-one and batch closed-form edits (ROME, MEMIT) are computationally cheap relative to any form of fine-tuning: they require estimating a covariance matrix $C$ once (a one-time, amortizable cost, typically computed from a large sample of generic text passed through the model), followed by a closed-form linear-algebra solve per edit or per batch of edits — orders of magnitude faster than gradient-based fine-tuning, and requiring no backpropagation through the full model at edit time. Hypernetwork-based methods (MEND) shift cost to an upfront training phase for the hypernetwork itself, after which individual edits are cheap forward passes. Memory-based methods (SERAC) avoid modifying the base model entirely, trading this for the added inference-time cost and system complexity of a scope classifier and counterfactual model that must run alongside the (unmodified) base model for every query. At scale, storing and estimating a reliable key covariance matrix for very large models can itself be a nontrivial engineering undertaking, since it requires representative activation statistics gathered across a broad sample of the model's actual usage distribution.
## 6. Practical Implementation Strategies
A practical editing pipeline typically starts with causal tracing (or a similar localization method) on a representative sample of facts of the relevant type to identify which layer(s) to target, since the optimal layer can vary somewhat by model architecture and fact category. For single or small numbers of edits, a ROME-style covariance-weighted rank-one update is usually sufficient and fast; for editing many facts at once (e.g., correcting a batch of outdated facts after a knowledge cutoff), a MEMIT-style joint batch solve, often distributing the total required update across several consecutive layers rather than concentrating it in one, better controls interference both among the new edits themselves and against the model's existing knowledge. Practitioners should always evaluate edits along all three axes — efficacy (does the edit take effect on the exact edited prompt), generalization (does it transfer to paraphrases and logically entailed queries), and specificity (do unrelated, superficially similar facts remain unaffected) — since optimizing for only one, particularly efficacy alone, can produce edits that look successful on the target prompt while quietly degrading broader model quality.
## 7. Benchmark Datasets & Evaluation
CounterFact and zsRE (zero-shot relation extraction) are the standard benchmarks for evaluating knowledge editing methods, each providing large sets of (subject, relation, original object, new counterfactual object) tuples along with paraphrase prompts (for measuring generalization) and "neighborhood" prompts referencing related-but-distinct facts (for measuring specificity). Reported metrics typically decompose into efficacy success rate (did the edit change the output on the exact edit prompt), paraphrase success rate (generalization), and neighborhood success rate or drawdown (specificity — how much unrelated, nearby knowledge was disturbed). Sequential/lifelong editing benchmarks additionally track how these three metrics degrade as a growing number of edits are applied cumulatively to the same model, which is often the more practically relevant regime since real deployments rarely need to make only a single edit in isolation.
## 8. Key Challenges & Limitations
Editing methods face a fundamental tension between generalization and specificity: pushing an edit to generalize more broadly to paraphrases and logical implications of the edited fact (a stronger, larger-magnitude update, or one directed along less-isolated key directions) inherently risks bleeding into and corrupting nearby, unrelated associations, and no method fully escapes this tradeoff. Sequential editing degrades over time — even individually well-behaved edits compound, and after enough edits many methods show measurable drift in general model capability well beyond the specifically targeted facts, an effect sometimes called "model collapse" under repeated editing. The locate-and-edit paradigm's assumption of clean localization is also imperfect: recent work has shown that the layer identified as causally important by activation patching is not always the layer where editing is most effective, undermining a fully principled way to choose edit targets from tracing results alone. Finally, most editing methods and benchmarks focus on simple, single-hop factual triples; editing more complex or multi-hop knowledge (facts that logically entail or depend on other facts) remains substantially harder and less well solved, since a correct edit may need to propagate consistently through everything that logically follows from the changed fact.
## 9. Hyperparameter Tuning
The regularization strength applied to the covariance matrix inversion ($C + \lambda I$ before inverting) directly trades off how aggressively the update exploits low-variance key directions (larger effective update magnitude, higher risk of instability) against how conservative and stable the edit is (smaller update magnitude, less specificity benefit from covariance weighting) — too little regularization can produce numerically unstable, overly large updates when the key covariance matrix is highly anisotropic. The target layer (or set of layers, for MEMIT-style distributed edits) is a critical choice best guided by causal tracing results on a validation set of similar facts, since editing too early or too late in the network relative to where the relevant association is actually stored degrades all three of efficacy, generalization, and specificity simultaneously. For batch editing, the number of facts edited per batch and the ridge regularization on the joint least-squares solve trade off computational efficiency (larger batches processed at once) against the accuracy of fitting every fact in the batch, particularly when many facts in a batch happen to share highly correlated key representations.
## 10. Real-World Applications & Case Studies
Deployed knowledge-editing systems are used to correct specific, high-visibility factual errors reported by users or discovered through auditing, without the cost, latency, and risk of a full retraining cycle — a practical necessity for any organization serving an LLM-based product where factual currency matters and full retraining happens only infrequently. Content moderation and safety teams use targeted editing to suppress specific harmful or incorrect associations a model has learned, offering a faster remediation path than waiting for the next full training run while a more thorough fix is prepared. Research use cases include using editing as a scientific probe into model internals — successfully editing a fact via a hypothesized locate-and-edit mechanism serves as causal evidence supporting (or, when it fails, evidence against) a specific interpretability hypothesis about how and where the model represents that category of knowledge, making editing methods a tool for mechanistic interpretability research as well as a practical maintenance technique.
## 11. Integration with Other Methods
Model editing is closely related to, but distinct from, parameter-efficient fine-tuning: both modify a small fraction of a model's parameters, but PEFT methods like LoRA typically target broad task adaptation via gradient-based training on a dataset, while editing methods target one or a few specific, precisely identified factual associations via closed-form or targeted optimization, often without any gradient descent at all. Editing also connects directly to mechanistic interpretability, since locate-and-edit methods both rely on and provide evidence for interpretability claims about where specific types of knowledge are stored. It intersects with hallucination mitigation as a potential remediation tool — if a hallucination or factual error is traced to a specific, correctable stored association, targeted editing offers a faster fix than retraining, complementing rather than replacing detection-and-filtering approaches for hallucinations that don't stem from a single identifiable stored fact.
## 12. Future Research Directions
Open problems include developing editing methods that remain stable and non-degrading over thousands of sequential edits applied across a model's deployed lifetime, rather than the current situation where most methods show measurable degradation well before that scale. Reconciling the dissociation between causal-tracing localization and effective-edit localization remains an important open theoretical question, since a fully principled account of where to edit (rather than an empirically-tuned choice) would make the whole pipeline more reliable and better understood. Extending editing to reliably handle multi-hop and logically entailed knowledge, where a single correct edit needs to propagate consistently through everything that depends on it, is a substantially harder unsolved problem than the single-fact editing that current benchmarks predominantly test. There is also growing interest in editing methods for capabilities and behaviors beyond narrow factual triples — including broader concept-level or safety-relevant edits — which raises both new technical challenges (these interventions are often far less cleanly localized than simple facts) and new evaluation challenges (defining efficacy, generalization, and specificity for a behavior rather than a fact is considerably less straightforward).
## 13. Summary & Key Takeaways
Model editing exploits the empirical finding that factual knowledge in transformer language models is substantially localized to identifiable components, primarily mid-layer MLPs functioning as key-value associative memories, enabling small, closed-form, targeted weight updates that insert or correct specific facts without the cost of full retraining. Covariance-weighted rank-one updates (ROME) and their batch generalization (MEMIT) achieve this by directing edits along key-space directions that are rarely used by other stored facts, substantially improving specificity over naive minimum-norm updates. The central, unresolved tension across the field is the three-way tradeoff between efficacy, generalization, and specificity, compounded by degradation under sequential editing at scale — meaning editing is best understood today as a fast, surgical remediation tool for specific, well-localized factual errors rather than a wholesale substitute for retraining or fine-tuning.
Keywords: model editing, knowledge editing, ROME, MEMIT, rank-one model editing, causal tracing, locate-and-edit, associative memory, key-value memory, MEND, SERAC, CounterFact, zsRE, efficacy generalization specificity, sequential editing, lifelong editing
---
## Appendix: Practical Labs
### Lab 1: Causal Tracing Localizes a Fact to the Correct Layer
import numpy as np
def simulate_causal_tracing(n_layers, true_layer, rng, noise_std=0.15):
"""Simulate per-layer contributions to a 'correct answer' logit for a clean run
vs. a corrupted run, then measure how much restoring each individual layer's
clean activation recovers the clean output (the causal tracing 'restoration
effect'). Only the layer that actually stores the fact should show a large effect."""
contribution_clean = rng.normal(0, noise_std, size=n_layers)
contribution_corrupt = contribution_clean.copy()
contribution_clean[true_layer] += 3.0
contribution_corrupt[true_layer] -= 3.0
for l in range(n_layers):
if l != true_layer:
contribution_corrupt[l] += rng.normal(0, noise_std * 0.5)
logit_clean = contribution_clean.sum()
logit_corrupt = contribution_corrupt.sum()
restoration_scores = np.zeros(n_layers)
for l in range(n_layers):
patched = contribution_corrupt.copy()
patched[l] = contribution_clean[l]
logit_patched = patched.sum()
restoration_scores[l] = (logit_patched - logit_corrupt) / (logit_clean - logit_corrupt)
return restoration_scores
def test_causal_tracing_localizes_fact_to_correct_layer():
rng = np.random.default_rng(1)
n_layers = 12
true_layer = 6
n_trials = 50
correct = 0
all_scores = []
for _ in range(n_trials):
scores = simulate_causal_tracing(n_layers, true_layer, rng)
all_scores.append(scores)
if np.argmax(scores) == true_layer:
correct += 1
mean_scores = np.mean(all_scores, axis=0)
print("mean restoration score per layer:")
print(" ".join(f"{s:.3f}" for s in mean_scores))
print(f"argmax-matches-true-layer rate: {correct}/{n_trials}")
assert correct / n_trials > 0.9
assert np.argmax(mean_scores) == true_layer
assert mean_scores[true_layer] > 2 * np.mean(np.delete(mean_scores, true_layer))
print("Causal tracing layer-localization test passed.")
if __name__ == "__main__":
test_causal_tracing_localizes_fact_to_correct_layer()### Lab 2: Covariance-Weighted Rank-One Edits (ROME) Improve Specificity
import numpy as np
def make_correlated_keys(n, d, rank, rng):
"""Existing fact keys drawn from an anisotropic (low-effective-rank) distribution,
mirroring how real hidden-state key vectors are far from isotropic."""
basis = rng.normal(size=(d, rank))
coeffs = rng.normal(size=(n, rank))
keys = coeffs @ basis.T
keys += rng.normal(scale=0.1, size=(n, d))
return keys
def rank_one_edit(W0, k_new, v_new, direction_matrix=None):
"""Rank-one update inserting (k_new -> v_new). If direction_matrix (e.g. C^-1) is
given, the update direction is C^-1 k_new (ROME-style); otherwise plain k_new (naive)."""
residual = v_new - W0 @ k_new
u = k_new if direction_matrix is None else direction_matrix @ k_new
denom = k_new @ u
return W0 + np.outer(residual, u) / denom
def test_covariance_weighted_edit_is_more_specific_than_naive_edit():
rng = np.random.default_rng(8)
d = 50
n_other = 400
other_keys = make_correlated_keys(n_other, d, rank=6, rng=rng)
W_true_proj = rng.normal(size=(d, d)) * 0.3
other_values = other_keys @ W_true_proj.T + rng.normal(scale=0.05, size=(n_other, d))
W0, *_ = np.linalg.lstsq(other_keys, other_values, rcond=None)
W0 = W0.T
baseline_error = np.mean((other_keys @ W0.T - other_values) ** 2)
k_new = make_correlated_keys(1, d, rank=6, rng=rng)[0]
v_new = rng.normal(size=d) * 0.5
C = (other_keys.T @ other_keys) / n_other
C_inv = np.linalg.inv(C + 1e-3 * np.eye(d))
W_naive = rank_one_edit(W0, k_new, v_new, direction_matrix=None)
W_cov = rank_one_edit(W0, k_new, v_new, direction_matrix=C_inv)
err_new_naive = np.linalg.norm(W_naive @ k_new - v_new)
err_new_cov = np.linalg.norm(W_cov @ k_new - v_new)
err_old_naive = np.mean((other_keys @ W_naive.T - other_values) ** 2)
err_old_cov = np.mean((other_keys @ W_cov.T - other_values) ** 2)
print(f"baseline old-key MSE={baseline_error:.5f}")
print(f"new-fact fit error: naive={err_new_naive:.6f} cov-weighted={err_new_cov:.6f}")
print(f"old-key MSE after edit: naive={err_old_naive:.5f} cov-weighted={err_old_cov:.5f}")
assert err_new_naive < 1e-6
assert err_new_cov < 1e-6
assert err_old_cov < err_old_naive
assert (err_old_naive - baseline_error) > 3 * (err_old_cov - baseline_error)
print("Covariance-weighted (ROME-style) edit specificity test passed.")
if __name__ == "__main__":
test_covariance_weighted_edit_is_more_specific_than_naive_edit()### Lab 3: Batch (MEMIT-Style) Editing Scales Better Than Sequential Editing
import numpy as np
def make_correlated_keys(n, d, rank, rng):
basis = rng.normal(size=(d, rank))
coeffs = rng.normal(size=(n, rank))
keys = coeffs @ basis.T
keys += rng.normal(scale=0.1, size=(n, d))
return keys
def sequential_rank_one_edits(W0, new_keys, new_values):
"""Naive approach: insert facts one at a time, each edit computed against the
CURRENT (already-edited) weight matrix, oblivious to the other pending edits."""
W = W0.copy()
for k_new, v_new in zip(new_keys, new_values):
residual = v_new - W @ k_new
denom = k_new @ k_new
W = W + np.outer(residual, k_new) / denom
return W
def batch_joint_edit(W0, new_keys, new_values, ridge=1e-3):
"""MEMIT-style: solve for ALL new key/value pairs simultaneously via a single
joint least-squares solve, rather than inserting them one at a time."""
K = new_keys
targets = new_values.T - W0 @ K.T
gram = K @ K.T
gram_reg = gram + ridge * np.trace(gram) / gram.shape[0] * np.eye(gram.shape[0])
coeff = np.linalg.solve(gram_reg.T, targets.T)
delta = K.T @ coeff
return W0 + delta.T
def retained_fact_error(W, keys, values):
return np.mean(np.sum((keys @ W.T - values) ** 2, axis=1))
def test_batch_editing_scales_better_than_sequential_editing():
rng = np.random.default_rng(15)
d = 50
n_other = 400
other_keys = make_correlated_keys(n_other, d, rank=6, rng=rng)
W_true_proj = rng.normal(size=(d, d)) * 0.3
other_values = other_keys @ W_true_proj.T + rng.normal(scale=0.05, size=(n_other, d))
W0, *_ = np.linalg.lstsq(other_keys, other_values, rcond=None)
W0 = W0.T
baseline_old_err = retained_fact_error(W0, other_keys, other_values)
edit_counts = [3, 8, 15]
seq_new_errors, batch_new_errors = [], []
for n_edits in edit_counts:
new_keys = make_correlated_keys(n_edits, d, rank=25, rng=rng)
new_values = rng.normal(size=(n_edits, d)) * 0.5
W_seq = sequential_rank_one_edits(W0, new_keys, new_values)
W_batch = batch_joint_edit(W0, new_keys, new_values)
seq_err = retained_fact_error(W_seq, new_keys, new_values)
batch_err = retained_fact_error(W_batch, new_keys, new_values)
seq_new_errors.append(seq_err)
batch_new_errors.append(batch_err)
print(f"n_edits={n_edits:2d} sequential new-fact MSE={seq_err:9.5f} "
f"batch new-fact MSE={batch_err:.6f}")
# batch editing stays accurate regardless of how many facts are edited at once
for e in batch_new_errors:
assert e < 0.05
# naive sequential editing increasingly fails to retain earlier-inserted facts
# as more edits are piled on, since each edit disturbs the keys of prior edits
assert seq_new_errors[-1] > seq_new_errors[0]
assert seq_new_errors[-1] > batch_new_errors[-1] * 50
print(f"baseline (unedited) old-key MSE={baseline_old_err:.5f}")
print("Batch (MEMIT-style) vs sequential editing interference test passed.")
if __name__ == "__main__":
test_batch_editing_scales_better_than_sequential_editing()### Lab 4: Edit Locality — Generalization to Paraphrases vs. Specificity for Unrelated Facts
import numpy as np
def make_correlated_keys(n, d, rank, rng):
basis = rng.normal(size=(d, rank))
coeffs = rng.normal(size=(n, rank))
keys = coeffs @ basis.T
keys += rng.normal(scale=0.1, size=(n, d))
return keys
def rank_one_edit(W0, k_new, v_new, C_inv):
residual = v_new - W0 @ k_new
u = C_inv @ k_new
denom = k_new @ u
return W0 + np.outer(residual, u) / denom
def test_edit_generalizes_to_paraphrases_but_preserves_unrelated_facts():
rng = np.random.default_rng(23)
d = 50
n_other = 400
other_keys = make_correlated_keys(n_other, d, rank=6, rng=rng)
W_true_proj = rng.normal(size=(d, d)) * 0.3
other_values = other_keys @ W_true_proj.T + rng.normal(scale=0.05, size=(n_other, d))
W0, *_ = np.linalg.lstsq(other_keys, other_values, rcond=None)
W0 = W0.T
C = (other_keys.T @ other_keys) / n_other
C_inv = np.linalg.inv(C + 2.0 * np.eye(d))
k_new = make_correlated_keys(1, d, rank=6, rng=rng)[0]
k_new = k_new / np.linalg.norm(k_new)
v_new = rng.normal(size=d) * 0.5
W_edited = rank_one_edit(W0, k_new, v_new, C_inv)
# paraphrase keys: small perturbations of the edited key (same underlying fact,
# different surface phrasing)
n_paraphrase = 200
paraphrase_keys = k_new + rng.normal(scale=0.05, size=(n_paraphrase, d))
paraphrase_keys /= np.linalg.norm(paraphrase_keys, axis=1, keepdims=True)
# unrelated keys: fresh random directions unrelated to the edited fact
n_unrelated = 200
unrelated_keys = make_correlated_keys(n_unrelated, d, rank=6, rng=rng)
def distance_to_target(keys, W, target):
return np.linalg.norm(keys @ W.T - target, axis=1)
para_before = distance_to_target(paraphrase_keys, W0, v_new)
para_after = distance_to_target(paraphrase_keys, W_edited, v_new)
unrel_before = distance_to_target(unrelated_keys, W0, v_new)
unrel_after = distance_to_target(unrelated_keys, W_edited, v_new)
para_improvement = np.mean(para_before - para_after) / np.mean(para_before)
unrel_improvement = np.mean(unrel_before - unrel_after) / np.mean(unrel_before)
print(f"paraphrase keys: mean distance-to-target before={np.mean(para_before):.4f} "
f"after={np.mean(para_after):.4f} (relative improvement={para_improvement:.3f})")
print(f"unrelated keys: mean distance-to-target before={np.mean(unrel_before):.4f} "
f"after={np.mean(unrel_after):.4f} (relative improvement={unrel_improvement:.3f})")
# the edit generalizes: paraphrases move substantially closer to the new target
assert para_improvement > 0.3
# the edit is specific: unrelated keys are barely affected
assert abs(unrel_improvement) < 0.05
assert para_improvement > 5 * abs(unrel_improvement)
print("Edit locality vs. generalization test passed.")
if __name__ == "__main__":
test_edit_generalizes_to_paraphrases_but_preserves_unrelated_facts()