AI/ML Explained: Knowledge Storage & Model Editing

Picture a 70B-parameter model as a giant distributed KV store. The fact "the Eiffel Tower is in Paris" isn't smeared across the whole cluster — it lives on a few specific nodes. The catch: no schema, no index, no docs. How do you find that node? The answer is chaos-engineering style (fault injection): deliberately corrupt a layer's activations, watch whether the model now gets it wrong, and reverse-engineer where the data is stored.

The old intuition says a neural net is a "black box where knowledge is globally entangled." But in 2021 Geva et al. gave a key insight: a Transformer's FFN (feed-forward layer) is essentially a key-value memory. The FFN has two matrices — each row of the first (W_in) is a key (matching an input pattern, e.g. "the subject is a landmark"), and the matching column of the second (W_out) is a value (pouring tokens like "Paris" into the output vocabulary). This turns "where is knowledge stored" from a philosophy question into a locatable engineering problem.

How to localize? Two main paths:

• Knowledge Neurons (Dai et al. 2022) — use gradient attribution (integrated gradients) to compute which neuron's activation contributes most to "producing the correct answer." Zero it out, and the model's confidence in that fact drops.
• Causal Tracing (the ROME paper) — a cleverer "corrupt-then-restore" experiment: ① add noise to the whole sentence's representation so the model gets it wrong; ② layer by layer, token by token, restore one position's clean activation; ③ see which restoration brings back the right answer — that position is the causal mediator storing the fact.

The ROME paper's core empirical finding: the mid-layer MLP, active while processing the last token of the subject name, is decisive for fact recall. Not a guess — a statistical result from causal experiments over thousands of facts. Once you know "where it's stored," the natural next step is "can you edit it precisely."

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Use a hook to grab FFN activations — observe the KV-memory structure
model = AutoModelForCausalLM.from_pretrained("gpt2-xl")
tok = AutoTokenizer.from_pretrained("gpt2-xl")

acts = {}
def grab(name):
    def hook(mod, inp, out): acts[name] = out.detach()
    return hook

# Hook after the mid-layer MLP activation (where keys fire)
layer = 17  # the "decisive mid layer" range ROME found
model.transformer.h[layer].mlp.act.register_forward_hook(grab("mlp_key"))

ids = tok("The Eiffel Tower is located in the city of", return_tensors="pt")
model(**ids)
# Which neurons fire hardest at the last subject token = candidate "knowledge neurons"
key = acts["mlp_key"][0, -1]
print("top-5 most-active neurons:", key.topk(5).indices.tolist())

The model learned "the current US president is X" wrong. Three options: re-pretrain (rebuild the whole database — astronomically expensive), fine-tune (a full-table update that easily hits other rows and causes forgetting), or ROME — like a surgical UPDATE ... WHERE id=?: change one MLP weight matrix by adding a single "rank-one" correction, so "this subject's" key maps to the "new object's" value, ideally touching no other record.

Facts go stale (elections, renames, data errors). Fine-tuning one fact is costly: build data, risk overfitting, and cause catastrophic forgetting of other knowledge. ROME pushes the previous card's insight to its conclusion — if the FFN is a linear associative memory, then editing it is a linear-algebra problem.

Treat the FFN's output matrix W as an associative memory mapping key vectors k to value vectors v (i.e. W k ≈ v). We want a specific subject key k* ("Eiffel Tower") to output a new value v* (pointing to "Rome"), while disturbing as little as possible what W does for all other keys. This becomes a constrained least-squares:

Intuition: a rank-one matrix = a column vector times a row vector — the lowest-information kind of edit — it changes W's behavior in one direction only. The C⁻¹ term is the key guardrail: it uses "the statistics of all historical keys" to project the update into a direction that doesn't disturb common keys — like adding a WHERE clause to the UPDATE so other rows aren't hit. After editing, the model not only says "Rome" on the original sentence but also generalizes to paraphrases like "which country is the tower in" — because it edited the internal representation, not memorized a string.

# EasyEdit: a unified implementation of ROME/MEMIT and other editors
# pip install easyeditor
from easyeditor import BaseEditor, ROMEHyperParams

hparams = ROMEHyperParams.from_hparams("./hparams/ROME/gpt2-xl")
editor = BaseEditor.from_hparams(hparams)

# Edit one fact at a time: change the Eiffel Tower's location to Rome
metrics, edited_model, _ = editor.edit(
    prompts=["The Eiffel Tower is located in the city of"],
    ground_truth=["Paris"],       # old fact
    target_new=["Rome"],           # new fact (to write into weights)
    subject=["Eiffel Tower"],     # key: locate the subject token
)
# metrics includes efficacy (edit success) / generalization (paraphrase)
#          / specificity (did unrelated knowledge get hit) — the 3 core metrics
print(metrics)

If ROME is a single UPDATE, MEMIT is batch UPDATE + load balancing. Looping edits a thousand times locks the table and accumulates error; MEMIT solves for the thousand corrections at once and spreads them across several mid-layer MLPs — like distributing write pressure across shards instead of dumping it all on one node.

Real needs are often batched: a knowledge base updated thousands of facts, sync them into the model at once. Editing one by one with ROME makes them fight. MEMIT (Meng et al. 2022) made two key improvements:

• ① Joint solve, not sequential stacking — put all edited facts into one least-squares objective solved together, so the updates "know about each other" and a later one won't wipe out an earlier one.
• ② Spread across multiple layers — instead of editing one layer, distribute the total correction across a contiguous range of mid-layer MLPs (e.g. layers 13–17). Each layer carries only part of the change, so per-layer disturbance is small and total capacity is large. This is the key to editing thousands at once, orders of magnitude beyond ROME.

MEMIT's empirical result: on GPT-J (6B) and GPT-NeoX (20B), it can inject thousands of associations at once, far beyond prior methods. It remains the mainstream baseline for "batch knowledge injection" research.

from easyeditor import BaseEditor, MEMITHyperParams

hparams = MEMITHyperParams.from_hparams("./hparams/MEMIT/gpt-j-6b")
editor = BaseEditor.from_hparams(hparams)

# Key difference: prompts is a batch, injected in one shot
metrics, edited_model, _ = editor.edit(
    prompts=[
        "The president of the USA is",
        "The capital of Australia is",
        "The CEO of Twitter is",
        # ... can be thousands
    ],
    target_new=["Jane Doe", "Sydney", "John Roe"],
    subject=["USA", "Australia", "Twitter"],
)
# MEMIT jointly solves this batch and spreads it across several mid MLPs
print(metrics)

You changed one master record in a database but didn't run the cascade update — every materialized view depending on it still holds the old value. Model editing has exactly this disease: you change "the club Messi plays for" to a new team, but the model still answers old facts for "who are Messi's teammates" and "which league is he in" — the logically entailed facts. You changed the primary key, the derived views didn't refresh, the data is now inconsistent.

ROME/MEMIT look beautiful, but Cohen et al. 2023 posed a stinging question: a fact is never isolated. Changing "A's capital is B" should ripple through a chain of inferences: "A's seat of government," "B is the capital of which country," "which way to travel from A to reach the capital"... They built the RippleEdits benchmark (~5000 edits, specifically testing these chain reactions) and found:

• Mainstream editing methods fail broadly on the "ripple" — they change the target fact correctly, but the logically entailed related facts don't follow, leaving the model self-contradictory.
• Counterintuitive result: a simple in-context edit (just put the new fact in the prompt, touch no weights at all) scored best on this benchmark — because a fact placed in context can be used consistently through reasoning, which an isolated fact welded into weights cannot.

Another layer of "knowledge conflict" is parametric memory vs context: when a fact retrieved by RAG clashes with the old memory in the model's parameters, which does the model trust? Research shows models often stubbornly trust their own parametric memory even when context clearly gives an update — the hardest part of hallucination mitigation, and proof that "writing knowledge into parameters" and "getting the model to use knowledge correctly" are two different things.

# After editing you MUST test the "ripple," not just the target fact
# edit: the club Messi plays for → Inter Miami

target = "Where does Lionel Messi currently play? "
print(query(edited_model, target))   # ✓ likely correct — the target fact

# The real test: logically entailed "ripple facts"
ripples = [
    "Which league does Lionel Messi play in? ",      # league
    "Who are Lionel Messi's teammates? ",            # teammates
    "In which country does Lionel Messi work? ",     # country
]
for q in ripples:
    # editing methods often crack here: target right, ripple all wrong
    print(q, "→", query(edited_model, q))

• Geva et al. 2021 — Transformer Feed-Forward Layers Are Key-Value Memories — the foundation of this whole "FFN = KV memory" editing line
• Dai et al. 2022 — Knowledge Neurons in Pretrained Transformers — locating fact-expressing neurons via gradient attribution
• Meng et al. 2022 — Locating and Editing Factual Associations in GPT (ROME) — the seminal causal-tracing + rank-one-editing work
• Meng et al. 2022 — Mass-Editing Memory in a Transformer (MEMIT) — scaling to thousands of edits
• Cohen et al. 2023 — Evaluating the Ripple Effects of Knowledge Editing — the RippleEdits benchmark, revealing cascade failures
• Hase et al. 2023 — Does Localization Inform Editing? — counterintuitive: the localized layer ≠ the layer to edit