MLPs are Hebbians: Constructing Efficient Fact-Storing MLPs for Transformers

TL;DR

Hazy Research proves that a Transformer’s gated MLP is structurally a Hebbian associative memory and gives a closed-form construction — no gradient descent — that stores a fact set at the information-theoretically optimal parameter rate. Given M key/value pairs and random Gaussian matrices, the recipe writes value embeddings into a gated MLP so that querying with a key returns the correct value by maximum dot product. Under isotropic embeddings the construction needs 10–104× fewer parameters than prior fact-storing MLP constructions at matched fact count, and 15–63× fewer when embedded inside a Transformer block that must tolerate attention-side noise on the query. The work explains why LLMs empirically store facts at ~2 bits/parameter and enables modular fact editing by hot-swapping the MLP.

Key claims

A gated MLP written as MLP(x) = W_down · (σ(W_gate · x) ⊙ (W_up · x)) with W_up = V·K^T / √d and W_gate = K^T / √d for random Gaussian projections is exactly a Hebbian outer-product memory where similarity is measured in the MLP’s feature space — no training required [§Construction, §“Our Hebbian memory view of MLPs”].
Under isotropic (uniform-spherical) key and value embeddings the construction stores M facts with a parameter count that matches the information-theoretic lower bound of ~2 bits/parameter [§“MLP storage capacity”].
For arbitrary (non-isotropic) key/value geometries the same rate holds up to penalization factors that depend on the embedding geometry — a first Transformer-compatible fact-storing MLP construction that satisfies all three empirical LLM properties: optimal-rate storage, arbitrary embedding geometries, and usability inside a Transformer [Abstract].
Inside a Transformer block the construction still retrieves the correct fact as long as the attention-produced query has L2 error below √(1/log M) from the true key — so the block (not just a standalone MLP) inherits the optimal-rate scaling [§“Transformer block storage capacity”].
At matched fact count the construction uses 10–104× fewer parameters than prior fact-storing MLP constructions in the isotropic setting, and 15–63× fewer inside a Transformer block; it also stays within √2 of the parameter count of MLPs actually trained by gradient descent [Abstract].
Because the MLP is constructed rather than trained, one can swap a Transformer’s MLP for a new one to insert or edit facts without retraining — demonstrated as a proof-of-concept [Abstract, §“Beyond efficient fact-storing MLPs”].

Method

Fix a fact set as a mapping from M key embeddings k_i ∈ R^d to value embeddings v_i ∈ R^d. Sample a projection matrix W ∈ R^{h×d} with Gaussian entries — this is the only random primitive. Set W_gate = W, W_up = W, and W_down = Σ_i v_i · (σ(W k_i) ⊙ (W k_i))^T. The resulting gated MLP is a Hebbian outer-product store where each key is written as its feature-space image φ(k_i) := σ(W k_i) ⊙ (W k_i) and each value is retrieved by a soft kernel lookup MLP(x) = Σ_i v_i · ⟨φ(x), φ(k_i)⟩. When x = k_j the diagonal term dominates and the readout collapses onto v_j; when x is a noisy version of k_j (as delivered by attention inside a Transformer block), the same argument goes through as long as the noise stays below ~1/√log M. The storage-capacity analysis measures the decoding margin — the gap between the correct value’s dot product and the largest incorrect one — rather than just reconstruction error, which is what closes the gap to the information-theoretic bound and to the empirical LLM regime.

Results

Storage rate: matches the information-theoretically optimal ~2 bits/parameter under isotropic embeddings; matches the same rate up to embedding-geometry penalization factors under arbitrary embeddings.
Parameter efficiency vs prior constructions: 10–104× fewer parameters at matched fact count for standalone MLPs, 15–63× fewer inside Transformer blocks.
Parameter efficiency vs GD-trained MLPs: stays within √2 of the parameter count of MLPs trained end-to-end by gradient descent to store the same fact set.
Editing: proof-of-concept demonstration that swapping a Transformer’s MLP with a freshly constructed one edits factual behavior without any retraining or model surgery beyond the swap.
Robustness inside Transformers: correct-fact retrieval is guaranteed as long as attention’s L2 error on the query stays below √(1/log M), giving the first block-level optimal-rate result rather than a standalone-MLP result.

Why it’s interesting

This is the first closed-form, Transformer-compatible fact-storing MLP that hits the ~2 bits/parameter information-theoretic ceiling — the same ceiling that Incompressible Knowledge Probes: Estimating Black-Box LLM Parameter Counts via Factual Capacity uses empirically as a black-box parameter-count estimator on 89 open models. Together the two papers form a matched pair: IKP measures the fact-per-parameter capacity of trained LLMs and finds it saturated at the optimum; MLPs-are-Hebbians explains why by giving a construction that reaches the same rate without training. It also sharpens the geometric-vs-associative-memory question raised by Deep sequence models tend to memorize geometrically; it is unclear why: the Hebbian view is associative in the classical Hopfield sense, but the similarity kernel lives in the MLP’s non-linear feature space, which is the ingredient that lets M scale linearly with parameter count rather than quadratically as in vanilla Hopfield networks — a middle position between the flat associative store and the emergent-geometric-embedding hypothesis. Compared to STEM: Scaling Transformers with Embedding Modules‘s STEM (which makes FFN memory architectural by token-indexed embedding tables), this work keeps the standard MLP shape and instead shows that the ordinary gated MLP was already the right architecture — the improvement is in how one sets the weights. On the mech-interp side, it delivers what the field has been asking for: a constructive rather than merely mechanistic account of factual recall, with modular fact editing falling out for free.