Can We Trade MLPs for Engrams?

A transformer holds two kinds of knowledge. Attention reads what is in the context. The weights hold what is not: that a thermostat takes a temperature, that living room is a place and warm white is a colour. In a standard transformer those facts live in the feed-forward layers, which act as key-value memories, and every token multiplies through all of them. Needle replaced its feed-forward layers with the Hadamard MLP, so the facts need another home. The engram is that home: hashed n-gram tables that a token reads by index.

Two kinds of parameter

A parameter in a matrix is read by every token. A parameter in a table is read by the tokens that need it.

For a matrix $W \in \mathbb{R}^{m \times n}$ and a token $x$, the product $Wx$ costs $mn$ multiply-adds and moves $mn$ weights from memory, every token, whatever the token is. For a table $T \in \mathbb{R}^{S \times c}$ read at row $i$, the fetch $T[i]$ costs $c$ numbers moved and no multiply-adds, and it costs the same for any $S$:

$\mathrm{cost}(Wx) = mn \text{ per token}, \qquad \mathrm{cost}(T[i]) = c \text{ per token, for any } S.$

Capacity in a matrix costs every token. Capacity in a table costs only the tokens that read it. So the design is to keep the matrices small and let the tables be large. The engram is the table: a conditional memory of hashed n-gram rows, in Needle's case sized for a phone.

The lookup

A table needs a row index, and the index has to come from the tokens with no learned computation in the way. The engram hashes the last few tokens. For a position $n$ with tokens $t_n, t_{n-1}, \dots$ and an order $o$, the row is

$i_o(n) = \mathrm{mix}(t_n, t_{n-1}, \dots, t_{n-o+1}) \bmod S.$

$\mathrm{mix}$ is the FNV hash, a few integer operations per token. Start from a seed $a$, then for each token in the window:

$a \leftarrow (a \oplus t_j) \cdot 16777619 \pmod{2^{32}}, \qquad \text{then } a \leftarrow a \oplus (a \gg 15).$

Needle uses orders 2 and 3: one row indexed by the last two tokens, one by the last three. Each order has three heads with different seeds, so a site has six tables, each with $S = 18{,}432$ rows of 128 numbers. The six rows a token fetches are concatenated into one vector

$e = [\,T_1[i_1],\ T_2[i_2],\ \dots,\ T_6[i_6]\,] \in \mathbb{R}^{768},$

which is exactly the model width. A row that would reach past the start of the sequence is zeroed.

The vocabulary has 8,192 tokens, so there are 67 million possible bigrams and a table has 18,432 rows. Most rows are shared by many n-grams. Three things make that work:

• the rows are learned, so a row becomes whatever the n-grams that land on it need, and the common n-grams win;
• the three heads of an order have different seeds, so two n-grams that collide in one head almost never collide in the other two, and the concatenated $e$ still tells them apart;
• the gate below decides, per token, how much of the fetched row to believe.

This is the hashing trick applied to n-grams, close to hash embeddings with several hash functions per key, and the table read is the gather that product-key memories also rely on, without the learned key search.

The gate

The fetched vector $e$ is projected to a key and a value, two ordinary $768 \times 768$ matrices:

$k = W_k\, e, \qquad v = W_v\, e.$

The value is smoothed over earlier positions with four learned taps at stride 3, so the value at a position also carries the rows fetched three, six and nine tokens back:

$v_n \leftarrow \sum_{j=0}^{3} \tau_j \odot v_{n-3j}.$

Then the gate compares the key with the residual stream $x$ of the block the site belongs to, and adds the value in proportion:

$\alpha = \sigma\!\left(\frac{\hat{x} \cdot \hat{k}}{\sqrt{d}}\right), \qquad x \leftarrow x + \alpha\, v, \qquad \hat{x} = \frac{x}{\mathrm{rms}(x)},\ \hat{k} = \frac{k}{\mathrm{rms}(k)}.$

If the key points the way the stream already points, $\alpha$ is near one and the value is written in. If it does not, $\alpha$ is small and the row is mostly ignored. A hash collision is a key that does not agree with the stream, and a gate that stays closed. The model never has to know which rows are shared; it learns keys that only open for the right context.

Where the sites sit

Needle 3 has five sites, at layers 3, 7, 11, 15 and 19 of its 20 blocks. Every site has its own six tables and its own $W_k$, $W_v$ and taps. The hash indices depend only on the tokens, so they are the same at all five sites; what differs between sites is the rows, and what the stream looks like when the key is compared with it. A site fires at the start of its block, before that block's attention, so the attention and the mixer of the block work on a stream that already contains the fetched facts.

The sites belong to blocks, and the intelligence ladder keeps them with their blocks. needle build --layers 8 keeps the sites whose layers survive and drops the other tables from the file, so a shallower model is also a smaller file.

The cost

Per site, with $d = 768$:

Over five sites the engram is 76.7M parameters, 70.8M of them in tables, for 5.9M multiply-adds per token. The dense feed-forward layers it replaces, twenty of them at $d = 768$, would be 94M parameters for 94M multiply-adds per token. The tables hold 92% of the engram's parameters and cost nothing to multiply; a token reads six rows of the 110,592 in a site, 768 numbers out of 14 million.

Results

Four questions, each answered with one run per setting on the SYNTH mixture, validation cross-entropy in nats per token, general text and tool calls scored separately.

Does it help. The same model trained twice for 4,000 steps, once with two memory sites of 8,192 rows and once without, sizes at 16 bits:

1.9% of the loss for 27% more file, almost all of it table rows that no token multiplies.

How many rows. An 18-layer, $d = 1024$ model with sites at layers 2 and 10, trained for 3,000 steps, tables of three sizes:

Where the first site sits. The same model with the second site fixed at layer 10 and the first moved:

Layer 2 was run twice; the two runs differ by 0.004, 0.2% of the loss, which is the noise floor for this series. Layer 1 is the only clearly bad position, 0.7% above layer 2; from layer 2 on, the differences are within that noise.

Sites against rows. At 20 layers, two layouts of the tables:

More sites with fewer rows each beat fewer sites with more rows, on tool calls by a wide margin, for a file 22% larger. Needle 3 ships five sites of 18,432 rows, where the tables at 2 bits are already two thirds of a 29 MB file; doubling the sites would add about 19 MB.

What the bytes buy, against the other knobs in the same 18-layer series. From 8,192 to 28,672 rows adds 45% to the file for 2.2% less loss. Widening the value heads from 64 to 128 adds 46% for 6.5% less loss, nearly three times as much for the same bytes. But the value bytes are multiplied by every token, and the table bytes are not. Table rows are the cheapest bytes to run, not the cheapest bytes to store, and Needle spends its byte budget on both.

On the device

In the shipped needle3.cact the five table tensors are 2-bit: 110,592 rows of 128 each, 17.7 MB of the 29 MB file. A token reads 30 rows, 3,840 numbers, 960 bytes at 2 bits, 32 bytes per row. Each row is decoded on fetch through a 4-entry codebook and the six rows are concatenated straight into $e$. The hash is a few integer operations per table. $W_k$ and $W_v$ are 2-bit as well, and the engine runs both projections from one quantised copy of $e$. The projected value is kept in an int8 cache per site, so the taps read the values three, six and nine tokens back instead of recomputing them.

The dense feed-forward layer, had it stayed, would have been 94M weights moved per token. The engram moves 5.9M through its projections and 3,840 through its tables. The rest of the tables stay where they are until a token asks for them.

What it gives up

• A row is indexed by the surface form of an n-gram. set the lights and turn on the lights are different rows; a paraphrase is a different key, and the model has to learn the fact twice or let attention carry it.
• Rows collide. The gate and the multiple heads make collisions cheap, not free; a rare n-gram that shares a row with a common one mostly reads the common one's content.
• The tables are two thirds of the file. On a device where flash is the limit and the model is never sliced, that is the cost of capacity that costs no arithmetic.
• The projections are ordinary matrices. 8% of the engram's parameters still cost 5.9M multiply-adds per token.
• The tables hold what the training data put in them. The attention-only study found that what a small feed-forward layer contributes is mostly parameter count, and the engram is the same bet made in a cheaper place; whether that holds for knowledge-heavy text at larger scale is the open question it left.

References

1. Cheng et al., Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models, 2026.
2. Geva et al., Transformer Feed-Forward Layers Are Key-Value Memories, 2020.
3. Weinberger et al., Feature Hashing for Large Scale Multitask Learning, 2009.
4. Svenstrup et al., Hash Embeddings for Efficient Word Representations, 2017.
5. Lample et al., Large Memory Layers with Product Keys, 2019.
6. Eastlake, Hansen, Fowler, Noll and Vo, The FNV Non-Cryptographic Hash Algorithm, IETF Internet-Draft.
7. Ndubuaku et al., A Controlled Study of Attention-Only Transformers, 2026.