Parameter-Efficient Fine-Tuning advanced 8 min read 6 flashcards

Locate-Then-Edit Knowledge Editing

ROME and MEMIT change a fact by solving a small constrained least-squares problem on one MLP weight matrix, which is the cheapest possible way to make a model say something new and the reason the change is so hard to contain.

A model says the Eiffel Tower is in Paris. You want it to say Rome, or, less mischievously, you want a model trained before an acquisition to name the right parent company. Fine-tuning on the new fact takes gradient steps that touch every parameter and tends to either not stick or to drag unrelated behaviour with it. Locate-then-edit asks a narrower question: which specific weights carry this association, and what is the smallest change to them that makes the new fact come out?

The premise comes from the observation that transformer feed-forward layers behave like key-value memories, with each key matching an input pattern and the corresponding value pushing probability mass onto particular output tokens (Geva et al., 2021, Transformer Feed-Forward Layers Are Key-Value Memories, EMNLP 2021). If a fact lives in a key-value pair, editing it is a linear algebra problem rather than a training problem.

Finding the site: causal tracing

ROME locates the relevant computation with a causal intervention. Run the prompt cleanly, run it again with the subject tokens corrupted by noise, then restore individual hidden states from the clean run into the corrupted run and measure how much of the correct prediction comes back. States whose restoration recovers the prediction are causally decisive. In GPT-style models this procedure lights up mid-layer MLP modules at the last subject token (Meng et al., 2022, Locating and Editing Factual Associations in GPT, NeurIPS 2022, arXiv:2202.05262).

That is the "locate" half, and it is also the half that turned out to be weaker than it looked. Causal tracing tells you where information is used; it does not tell you where an edit will succeed. Hase et al. found that edit success at a layer is statistically uncorrelated with the tracing signal at that layer, and that facts can be changed by editing weights in an entirely different location from where tracing points (Hase et al., 2023, Does Localization Inform Editing?, arXiv:2301.04213).

The edit itself: rank-one, closed form

Treat the second MLP matrix \(W\) as an associative memory that maps keys to values, \(W K \approx V\). You want a new pair \((k_*, v_*)\) inserted while the existing mapping is disturbed as little as possible. Minimise \(\|\hat{W}K - V\|_F\) subject to \(\hat{W}k_* = v_*\), and the Lagrangian gives a closed form:

\[\hat{W} = W + \frac{(v_* - W k_*)\,(C^{-1} k_*)^{\top}}{(C^{-1} k_*)^{\top} k_*}\]

where \(C = \mathbb{E}[k k^{\top}]\) is the uncentered second moment of keys estimated from a sample of Wikipedia text. The update is rank one: a single outer product. Computing it costs a forward pass, an optimisation over \(v_*\) of a few dozen steps, and one precomputed covariance.

MEMIT keeps the same objective but spreads the residual across a range of critical layers instead of dumping it into one, which is what lets it absorb thousands of facts in a single pass rather than dozens (Meng et al., 2023, Mass-Editing Memory in a Transformer, ICLR 2023, arXiv:2210.07229).

Why it is a projection, not a rewrite

The rank-one form explains both the appeal and the pathology. Written as \(\Delta = u\,\tilde{k}^{\top}\), the edit adds \(u\,(\tilde{k}^{\top}k)\) to the output for any key \(k\). Keys orthogonal to \(\tilde{k}\) are untouched, which is the containment story. Keys that merely overlap with it get a scaled dose of the same output shift, which is the contamination story, and nothing in the objective distinguishes "related fact that should also change" from "unrelated fact that should not".

The preservation term is only as good as \(C\). The covariance is estimated from a generic corpus, so it encodes what typical keys look like in aggregate, not which specific facts you care about keeping. An edit that is well-behaved in expectation can still be badly behaved on the handful of neighbours that matter.

When it breaks

Success is measured where the edit was made. Prompt the edited model with the exact subject-relation phrasing and it answers correctly. Ask a paraphrase, a two-hop question, or a question about a consequence of the fact, and the edit often fails to show up at all. See ripple effects and edit portability.

The edit is tuned on a single prompt. \(v_*\) is optimised to maximise the target token's probability given one prompt template, with a KL term to limit drift on the subject's other associations. That is a very local objective for a change the reader expects to be global.

It degrades with repetition. Each edit adds another rank-one term to the same matrix, and the terms accumulate. Downstream ability collapses after enough sequential edits, in a pattern specific enough to have its own literature (Gupta et al., 2024, ACL Findings). See sequential editing and model collapse.

Locality assumes localisation. If a fact is not stored in one place, editing one place cannot fully change it. The tracing result that motivated the method does not survive as a guide to where editing works.

Check yourself

6 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track