🧠 MLPs are Hebbian Memories: A Simple Recipe for Fact-Storing Transformers

Roberto Garcia*, Jerry Liu*, Ronny Junkins*, and Chris RΓ©

⚑ TL;DR

It sounds like sci-fi, but it's real: we can instantly build knowledge into a Transformer block, no training required!

Previously, we showed that MLPs can be constructed to store facts. Since then, we found a much simpler and more powerful explanation: a Transformer MLP is naturally a Hebbian memory. This view lets us develop the first closed-form fact-storing MLP construction, requiring no gradient descent 🀯, that packs facts at the information-theoretically optimal rate and is usable within Transformer blocks for factual recall!

Full paper team: Roberto Garcia*, Jerry Liu*, Ronny Junkins*, Sabri Eyuboglu, Atri Rudra, Chris RΓ©

πŸ“„ Paper
πŸ’» GitHub

The Hazy Research Bot constructing a fact-storing MLP

The Hazy Research Bot Constructing a Fact-Storing MLP.

Beyond efficient fact-storing MLPs (click to expand!)

Excitingly, our construction's impact goes beyond efficiently building MLPs! This post focuses on the MLP construction and its storage capacity, but in our paper we discuss how it also:

  • Gives the first theoretical account of why LLMs can store facts at an information-theoretically optimal rate, as observed in prior work.
  • Offers a constructive, rather than mechanistic, interpretability of fact storage inside Transformers. Since we can construct fact storing MLPs in closed-form, we can exactly understand how facts live in them.
  • Enables seamless fact editing in Transformer blocks without retraining.

πŸ”‘ Fact Storage

Prior work (e.g., Allen-Zhu and Li 2024, Morris et al. 2025) has empirically found that LLMs store facts near the information-theoretically optimal rate of Θ(Flog⁑F)\Theta(F\log F) parameters for FF facts. Our goal is to match this rate by construction in our MLPs, and to show that Transformers can actually use them.

What do we mean by facts? We study facts through the lens of a fact set, which is simply a mapping from keys to values. Concretely, a fact set is defined by the key embeddings K={k1,…,kF}K=\{k_1,\ldots,k_F\}, value embeddings V={v1,…,vF}V=\{v_1,\ldots,v_F\}, and a mapping between them f:[F]β†’[F]f:[F]\to[F]. For example, the key embedding for "France" might map to the value embedding for "Paris" in the capitals fact set.

What do we mean by a fact-storing MLP? We say an MLP\mathrm{MLP} stores the fact set if, when queried on key kik_i, its output has the highest dot product with the correct value vf(i)v_{f(i)}:

⟨MLP(ki),vf(i)⟩>⟨MLP(ki),vj⟩forΒ allΒ jβ‰ f(i).\langle \mathrm{MLP}(k_i), v_{f(i)} \rangle > \langle \mathrm{MLP}(k_i), v_j \rangle \quad \text{for all } j \neq f(i).

This matches how language models decode tokens: the model produces a hidden vector, then scores possible next tokens by dot product with output embeddings.

πŸ”§ A simple MLP construction

How do we build an MLP that stores a fact set? Surprisingly, a very simple closed-form recipe works.

The construction. Given keys kik_i and values vf(i)v_{f(i)}, sample random matrices A,G∈RmΓ—dA,G \in \mathbb{R}^{m\times d} from a Gaussian distribution and construct the gated MLP:

MLP(x):=B ϕ(x),Ο•(x):=(Ax)βŠ™(Gx),B:=1Fβˆ‘ivf(i) ϕ(ki)⊀\mathrm{MLP}(x):=B\,\phi(x), \qquad \phi(x) := (Ax)\odot(Gx), \qquad B := \frac{1}{F}\sum_i v_{f(i)}\,\phi(k_i)^\top

That's the whole construction, no gradient descent!

Our Hebbian memory view of MLPs (click to expand!)

A classical Hebbian memory stores key-value pairs as a sum of outer products, W=βˆ‘iviki⊀W = \sum_i v_i k_i^\top, and retrieves values with Wq=βˆ‘ivi⟨ki,q⟩Wq = \sum_i v_i \langle k_i, q\rangle. This is a weighted sum of stored values, where each value's weight is its key's similarity to the query. Our construction is exactly a Hebbian memory, but with the similarity measured in the MLP's feature space Ο•\phi instead:

MLP(q)=B ϕ(q)⏟MLP=1Fβˆ‘ivf(i)β€‰βŸ¨Ο•(ki),Ο•(q)⟩=1Fβˆ‘ivf(i) K(ki,q)⏟HebbianΒ Memory,\underbrace{\mathrm{MLP}(q) = B\,\phi(q)}_{\text{MLP}} = \frac{1}{F}\sum_i v_{f(i)}\,\langle \phi(k_i), \phi(q)\rangle = \underbrace{\frac{1}{F}\sum_i v_{f(i)}\,K(k_i, q)}_{\text{Hebbian Memory}},

Intuitively, the kernel similarity K(ki,q)K(k_i, q) acts as a soft lookup. When we query with a stored key q=kjq = k_j, the matching term K(kj,kj)K(k_j, k_j) is large while the mismatched terms K(ki,kj)K(k_i, k_j) for i≠ji \neq j are small, so the weighted sum collapses onto the correct value vf(j)v_{f(j)}.

Hebbian kernel memory retrieval flow from stored keys through kernel comparison to weighted sum of stored values
MLPs are Hebbian memories. An MLP compares a query to stored keys in feature space, then returns a similarity-weighted sum of stored values.

πŸ“ˆ MLP storage capacity. How many parameters does this construction need to store a fact set? Excitingly, we find that our simple MLP construction with WW parameters stores a fact set of size FF exactly at the information-theoretically optimal rate when the key and value embeddings are isotropic (e.g. uniform-spherical):

W=Θ(md)=Θ(Flog⁑F)W=\Theta(md)=\Theta(F\log F)

We show an analogous result for non-isotropic key and value embeddings (like those from LLMs) in our paper. There, the storage capacity scales at the same rate, up to penalization factors depending on the embedding geometry.

Standalone fact-storage capacity scaling under isotropic embeddings
MLP capacity scaling. Our constructed Hebbian MLPs achieve optimal W=Θ(Flog⁑F)W=\Theta(F\log F) capacity under isotropic embeddings, improving over prior constructions (NTK) by 1010-104Γ—104\times while staying within 66-10Γ—10\times of gradient descent (GD)-trained MLPs. We also include two enhanced flavors of our construction: a whitened variant that uses a covariance-whitened kernel, and a data-dependent variant that solves a least-squares problem to store the fact set. Check out our paper to learn more!

🎯 Transformer block storage capacity. Storing facts in a standalone MLP is one thing; using them inside a Transformer is another. When an MLP sits inside a Transformer block, attention doesn't hand it the exact key kik_i, but a noisy version instead. Excitingly, our MLP construction can handle noisy fact queries gracefully.

Transformer block schematic: attention produces a noisy query that a constructed MLP decodes to the correct value
Fact recall inside a Transformer block. Attention hands the MLP a noisy query q=k1+Ξ΅q=k_1+\varepsilon instead of the exact key; our constructed MLP, and thus the whole Transformer block, still decodes the correct value v1=Emb(Paris)v_1=\mathrm{Emb}(\text{Paris}).

Using our MLP construction, we derive that a Transformer block still retrieves the correct fact as long as the attention noise (i.e. the L2 error between a query vector and the exact key embedding of the fact being queried) stays below d/(Flog⁑F)\sqrt{d/(F\log F)}. When it does, the block inherits the same information-theoretically optimal scaling as MLPs,

W=Θ(md)=Θ(Flog⁑F).W = \Theta(md) = \Theta(F\log F).
Transformer fact-storage capacity scaling
Transformer capacity scaling. Under bounded attention noise, Transformer blocks with our MLP construction achieve optimal W=Θ(Flog⁑F)W=\Theta(F\log F) capacity, with the data-dependent variant up to 63Γ—63\times better than prior constructions. As in the MLP capacity figure, we also include our whitened and data-dependent constructions.

This is the first MLP construction demonstrating that Transformer blocks can store facts at the information-theoretically optimal rate. Compared to prior work, our MLP construction is not only simpler but scales better inside Transformers: our data-dependent construction stays within 3Γ—3\times of GD-trained MLPs and is up to 63Γ—63\times better than prior constructions as the number of facts grows at fixed model dimension.

πŸ”­ What's next

We're excited about these results, since a simple mathematical object, a Hebbian kernel memory, allows us to develop fact-storing MLPs and Transformer blocks with provable capacity guarantees! But there's still a ton more to do:

  • Understanding how pretrained LLMs store facts. Can we use our theory to understand the MLPs inside pretrained language models? Can we directly extract or edit the knowledge stored within these MLPs?

  • Studying memory in sequence mixers and multi-layer Transformers. Can we extend our theory to fact storage in the parameters of sequence mixers, such as attention's KV caches? How do attention and MLP layers cooperate to store knowledge in a multi-layer model?