Feed-Forward as Key-Value Memory

The feed-forward block is usually written as two matmuls with something in between. Expand the second one by column and a different object appears.

where $k_i$ is the $i$-th row of $W_1$ and $v_i$ the $i$-th column of $W_2$. The block holds $d_{ff}$ key-value pairs. A token’s activation is a soft address into them, and the output is whatever the matching entries say to write to the residual stream.

This is not an analogy imposed from outside — it is the same arithmetic, regrouped. What makes it a claim rather than a rewrite is the empirical part: the keys correspond to human-recognisable input patterns, and the values correspond to recognisable output distributions.

What the keys match

Geva et al. took each $k_i$, found the training examples that activate it most strongly, and had annotators describe them. Almost every neuron examined had a describable trigger, and the triggers were stratified by depth: shallow layers keyed on surface patterns — a specific suffix, a bigram, a formatting convention — and deep layers on semantic ones, such as “a sentence about a military rank” or “the object of a purchase”.

What the values write

The second paper is the sharper result. Project $v_i$ through the unembedding matrix and read the top tokens — the same operation as the logit lens, applied to a weight rather than an activation. The result is interpretable: a value vector’s top tokens form a coherent category, and it is usually the category that plausibly follows the key’s trigger.

The block does not compute a distribution. It casts weighted votes over the vocabulary, and the residual stream accumulates them across layers — which is why the logit lens shows the prediction sharpening layer by layer rather than appearing at the end.

Editing

If a fact is stored as a small set of entries, changing it should not require retraining. ROME formalises this: treat the block as a linear associative memory $W_2$ satisfying $W_2 K \approx V$, and solve for the minimal update that maps a chosen key to a chosen value while disturbing the others as little as possible.

It works, and the caveats are instructive. Causal tracing localises a fact to a specific layer range, and edits applied elsewhere do not take. Edits generalise to paraphrases but frequently not to the reversed relation. And at scale the edits interfere — MEMIT extends the method to thousands, and past that the model degrades.

import torch
from torch import Tensor


@torch.no_grad()
def describe_entry(w1: Tensor, w2: Tensor, unembed: Tensor, i: int, k: int = 10):
    """What entry i matches on, and what it writes to the vocabulary."""
    key, value = w1[i], w2[:, i]                     # [D], [D]

    # The value's vote, read directly in token space.
    promoted = (unembed @ value).topk(k).indices     # [k] token ids
    return key, promoted


@torch.no_grad()
def entry_contributions(w1: Tensor, w2: Tensor, x: Tensor, k: int = 10):
    """Which entries actually fired for this token, ranked by contribution."""
    h = torch.relu(w1 @ x)                           # [d_ff]
    # Not just activation: an entry with a large activation and a small value
    # vector contributes little.
    strength = h * w2.norm(dim=0)
    return strength.topk(k)

That second function is the one worth keeping. Ranking by activation alone is the common mistake and it is misleading — what reaches the residual stream is $h_i \|v_i\|$, and the two orderings differ substantially.