Blog post derives memory capacity of linear layers, MLPs and attention

The post starts from a familiar experience: a language model can recall obscure facts, and the author asked a model called Astra to recite the rules of the board game Dune: Imperium on the drive to a friend's game night. The question it sets out to answer is how models remember so much, and how much they can store.

Its answer is a way of looking at linear maps, MLPs and self-attention as different kinds of memory stores, from which the author derives heuristics (not exact results) for how much each can hold. A memory is a key-value pair: the key is used for lookup by similarity, the value is what we want back. A store that recalls by similarity is called an associative memory. The quality of a store is judged by one quantity, the interference-to-signal ratio, which measures how much the other memories get in the way when one is recalled. Capacity is the number of memories a store can hold before interference grows as large as the signal. The headline claim in the introduction is that moving from a linear layer to an MLP to self-attention, with 64-dimensional keys, is a jump from about 64 memories, to about 2,700, to over a million. Self-attention is described as having very good recall and a potentially unbounded capacity.

The simplest store is the outer-product store: the weights are the sum of v_i k_i transposed over all memories, so storing a new memory just means adding an outer product. Querying multiplies the input by the matrix, which compares the query with every key and returns the sum of the values weighted by those similarities. A short aside shows that a gradient descent step on a linear layer is itself a memory write, with key x and value equal to minus the learning rate times the output gradient. The author points to Irie, Csordás & Schmidhuber for the view of a gradient-trained linear layer as attention over its training inputs.

With unit-length keys, recall returns the target value (the signal) plus interference from every other memory whose key resembles the query. If keys were orthonormal, interference would vanish, but a space of dimension d_k has only d_k orthogonal directions, so perfect recall is impossible once M exceeds d_k. For random unit keys, each similarity is about plus or minus 1/sqrt(d_k). With random values of similar length, the M-1 interference terms add like a random walk, giving an interference-to-signal ratio of about sqrt(M/d_k). Setting that to 1 gives a capacity of M ≈ d_k. Even below that, the author notes, the outer-product store produces significant interference.

There are three ways to reduce it: make d_k larger (limited by hardware), store memories more carefully, or recall more selectively by suppressing weakly similar keys. The first careful option gives the optimal linear store: choose the weights by least-squares regression so that recall fits all memories at once. For M > d_k the solution is W = V transposed K (K transposed K) inverse, which is the outer-product store with the query first multiplied by the inverse of the key covariance. That whitens the keys so shared directions are not overweighted. When M ≤ d_k and the keys are linearly independent, it fits every memory exactly, so there is no interference even for non-orthogonal keys. Beyond d_k the ratio is about sqrt(M/d_k - 1), which reaches 1 at M = 2d_k: twice the capacity of the outer-product store. An aside derives this using the projection matrix onto the column space of the keys; the signal weights sum to d_k, and the extra -1 is the d_k memories' worth of interference the optimal store cancels.

A numerical check samples random unit-length keys in d_k = 64 dimensions and queries with a stored key. The outer-product store hits capacity at M ≈ d_k = 64. The optimal store recalls exactly up to M = d_k, its interference then rises from zero like sqrt(M/d_k - 1), and it reaches capacity near M = 2d_k = 128, as predicted. Because it cancels a fixed number of memories' worth of interference, the gap between the two stores shrinks as M grows.

The optimal store is still linear in the similarities, scaling small and large ones by the same factor. The MLP section changes that: pass each similarity through a nonlinearity sigma before the weighted sum, f(x) = V transposed sigma(Kx). The author says this rediscovers a two-layer MLP from first principles, with the rows of the first weight matrix as keys and the columns of the second as values. The post cites Sukhbaatar et al., and Geva et al., for noticing the resemblance between feedforward layers and memory stores. The text available here breaks off at that point, so the MLP and self-attention capacity derivations are not covered below.

Key facts

The post frames linear maps, MLPs and self-attention as associative memories and judges each by its interference-to-signal ratio; capacity is the number of memories held before interference equals the signal.
The outer-product linear store has capacity M ≈ d_k, from an interference-to-signal ratio of about sqrt(M/d_k) for random unit keys and random values of similar length.
A least-squares (optimal) linear store recalls exactly up to M = d_k and has ratio about sqrt(M/d_k - 1) beyond, giving capacity M = 2d_k, twice the outer-product store.
A numerical check with d_k = 64 random unit keys matches: outer-product capacity near 64, optimal store near 128.
With 64-dimensional keys the introduction claims a progression of about 64, about 2,700, then over a million memories for linear, MLP and self-attention; applying a nonlinearity to key similarities yields a two-layer MLP.

Why it matters

The post offers a compact intuition for why transformer components can hold so much knowledge: each layer type is a memory store, and the differences between them come down to how much recall interference they allow. It links a plain linear layer, the MLP and attention through one idea (recall by similarity), and shows that a gradient descent step on a linear layer can be read as writing a memory. Its capacity numbers (about 64, about 2,700, over a million at 64-dimensional keys) give a feel for the scale of the gap between layers.

Who it affects

Readers who want a first-principles picture of what MLP and attention layers do, such as students and engineers new to transformer internals. The material is a derivation with some linear algebra (least squares, projection matrices, pseudoinverses), so it suits people comfortable with that notation.

How to use it

As a mental model rather than a tool. When thinking about a layer, ask what its keys and values are and how recall suppresses interference from dissimilar keys. The post's own framing: rows of the first MLP weight matrix act as keys and columns of the second as values, and the interference-to-signal ratio is the quantity to track. The post includes an interactive visual where the number of stored memories M can be dragged past d_k = 64.

How solid is it

The linear-store results are derived in the text and checked numerically at d_k = 64, where measured capacities match the predictions of about 64 and about 128. The post itself calls its results heuristics, resting on assumptions of random unit-length keys and random values of similar length. The ~2,700 (MLP) and over-a-million (self-attention) figures are stated in the introduction, but their derivation is not in the text available here.

Risks and caveats

The capacities are idealised estimates for random keys and values, not measurements of real language models; no experiments on real large language models are reported in the text available here, and it does not claim that real models store facts this way. The MLP and self-attention figures cannot be checked from the available text. The sentence crediting earlier work on feedforward layers as memories is cut off after the mention of Geva et al.

“We just rediscovered a two-layer MLP from first principles!”

— From the post on alexwang.ai