nucleation

THEORY-1030: In transformer language models, each hidden unit of a feed-forward layer acts as a key–value memory, its input weights detecting patterns in the prefix and its output weights promoting the tokens that follow them, and a layer's output composes many such memories while the residual stream refines the prediction from layer to layer, as Geva, Schuster, Berant and Levy propose

2 min read

Properties4
StatusProposed
Tagsrepresentation-learning
SourceLIT-1576 — Transformer Feed-Forward Layers Are Key-Value Memories
ExtendsTHEORY-1033 — In a 16-layer transformer language model trained on WikiText-103, the training prefixes that most activate a sampled feed-forward hidden unit share patterns that human annotators recognize, mostly shallow in layers 1–9 and semantic in layers 10–16; projected through the output embedding, a unit's output weights put the next token of its top prefix first in up to about 3.5% of units in the upper layers and almost none below; and a layer's top prediction differs from every unit's in at least about 68% of examples

Graph View

Source

Geva et al. (2021), LIT-1576, read in full in NOTE-1200.

What is asserted

A functional reading of the algebra. FF(x) = f(x·Kᵀ)·V is an unnormalized key–value memory (demonstrated, and Sukhbaatar et al.’s). The claim is that trained models use it as one: the units’ keys are pattern detectors and their values next-token distributions, so that the feed-forward layers are a store of pattern-to-continuation associations, composed within a layer and refined across layers.

What the record holds toward it

  • For: THEORY-1033, in one WikiText-103 model: recognizable patterns for every sampled key, a shallow-to-semantic gradient, and value–key agreement far above chance in the upper layers.
  • Weak where it matters: that agreement is about 3.5% of units at best, and there is no causal test.
  • Against the unit as the unit of memory: the dictionary-learning readings (LIT-1545, LIT-1533) find MLP neurons largely polysemantic and read cleaner features off a learned overcomplete basis. That bears on “each hidden unit”, not on the memory form. The record’s basis question (THEORY-802) applies to the choice of units as memories.
  • Borrowed at another level: ROME (LIT-1536, THEORY-1011) reads the whole output matrix as a linear associative memory keyed by activation vectors. That reading is not this one, and evidence for one is not evidence for the other.

Why it stays Proposed

The evidence is from one small model and is correlational throughout. The reading’s “each key correlates with textual patterns” is shown; “operate as key–value memories”, in the sense that the model’s predictions are produced by those associations, is not.

What would confirm or reject it

The promote_when above.

Backlinks