Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Feed-Forward Networks › Interpreting the FFN as Key-Value Memory

Michael BrenndoerferUpdated June 10, 20257 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 107Feed-Forward Networks › Interpreting the FFN as Key-Value Memory

Reading section 5 of 13 · Interpreting the FFN as Key-Value Memory · 7 minSection 5 of 13 · 7 min

Section 5 of 13 · 7 min read

Feed-Forward Networks › Interpreting the FFN as Key-Value MemoryOne section · Sections in this chapter5 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

One of the most illuminating ways to understand the FFN is to see it as a neural network layer and as an associative memory: a structure that stores key-value pairs in its weights and retrieves the appropriate values when presented with matching inputs. This interpretation, developed by researchers at Tel Aviv University and other institutions, helps explain how transformers store and retrieve factual knowledge, and why simply scaling up the FFN increases a model's capacity to "know" things.

The idea is elegant. Consider the first layer's weight matrix $W_1$ . Each column of $W_1$ can be thought of as a "key" that matches certain input patterns. When we compute $xW_1 + b_1$ , we are essentially computing the dot product between the input $x$ and each of the $d_{ff}$ keys. Each element of the resulting vector measures how well the input matches the corresponding key.

To make this precise: the pre-activation value for hidden dimension $i$ measures the alignment between the current input and the $i$ -th key pattern. Given an input $x$ and the $i$ -th column of $W_1$ (call it $\mathbf{k}_i$ ), this alignment is:

h_i = x \cdot \mathbf{k}_i + b_{1,i}

where:

$h_i$ : the pre-activation value for hidden dimension $i$ , measuring how well the input matches key $i$
$x$ : the input vector ( $d_{\text{model}}$ -dimensional), representing the current token's contextual state
$\mathbf{k}_i$ : the $i$ -th column of $W_1$ , encoding the pattern that key $i$ is tuned to detect
$b_{1,i}$ : the $i$ -th element of the bias vector $b_1$ , which acts as a threshold, biasing the key toward activation or suppression

When the input $x$ aligns well with key $\mathbf{k}_i$ (high dot product), the corresponding hidden dimension activates strongly. After ReLU, only positive activations survive, so keys that match the current input "fire" while keys that don't match are silenced. This is exactly how an associative memory works: present a query, retrieve only the memories that match.

The second layer's weight matrix $W_2$ contains the "values" associated with each key. Each row of $W_2$ (call it $\mathbf{v}_i$ for the $i$ -th row) is the value vector that gets added to the output when hidden dimension $i$ is active. The final FFN output is a weighted sum of these values, where the weights are the hidden activations:

\text{FFN}(x) = \sum_{i=1}^{d_{ff}} h_i^+ \cdot \mathbf{v}_i

where:

$h_i^+ = \text{ReLU}(h_i) = \max(0, h_i)$ : the post-activation strength of the $i$ -th key match
$\mathbf{v}_i$ : the $i$ -th row of $W_2$ , encoding the value retrieved when key $i$ fires
The sum runs over all $d_{ff}$ key-value pairs, but in practice only the active (nonzero) $h_i^+$ values contribute

FFN as Key-Value Memory

The feed-forward layer can be interpreted as an associative memory where $W_1$ columns are keys, hidden activations are match scores, and $W_2$ rows are values. Input patterns that match certain keys retrieve their associated values. With $d_{ff}$ key-value pairs per layer and dozens of layers per model, the total storage capacity of the FFN weights is enormous.

This memory interpretation explains several important empirical observations about transformers. Researchers have found that specific neurons in FFN layers activate for particular concepts: individual neurons that fire consistently for "the Eiffel Tower," for "programming languages," or for "past tense verbs." These neurons are not hand-designed. They emerge from training, as the model discovers that certain hidden dimensions reliably encode certain patterns. The FFN learns to store factual associations in its weights, and the attention mechanism retrieves them by routing relevant input patterns through the FFN.

This interpretation also explains a peculiar phenomenon called "knowledge editing" in language models. When researchers want to update a model's stored facts (for instance, changing "The president of France is X" to a newer value), they can sometimes accomplish this by modifying specific rows of FFN weight matrices. The change propagates naturally because those rows encode the associated value that gets retrieved when the "president of France" key fires.

The memory interpretation also clarifies why larger FFN dimensions store more knowledge. More hidden dimensions means more key-value pairs, which means more distinct patterns the model can recognize and more distinct responses it can retrieve. This is why scaling the FFN dimension is one of the primary ways to increase a model's factual capacity.

Let's visualize this interpretation with a concrete example:

In[12]:

Code

# Interpret FFN as key-value memory
d_model_small = 8
d_ff_small = 16

W1_small = np.random.randn(d_model_small, d_ff_small) * 0.5
b1_small = np.zeros(d_ff_small)
W2_small = np.random.randn(d_ff_small, d_model_small) * 0.5
b2_small = np.zeros(d_model_small)

# Create an input that strongly activates certain hidden dimensions
x_test = np.random.randn(d_model_small)

# Compute hidden activations (before and after ReLU)
pre_activation = x_test @ W1_small + b1_small
hidden_activations = relu(pre_activation)

# See which "keys" were matched (high activation)
active_dims = hidden_activations > 0.5

Out[13]:

Console

FFN as Key-Value Memory:
  Input dimension: 8
  Number of 'keys' (hidden dim): 16

Pre-activation (dot product with keys):
  [ 2.89 -0.3  -1.39  1.7   0.28  0.9   1.25 -0.09 -0.54  1.3  -2.01 -0.44
 -0.15  3.36  0.76 -1.39]

After ReLU (matched keys):
  [2.89 0.   0.   1.7  0.28 0.9  1.25 0.   0.   1.3  0.   0.   0.   3.36
 0.76 0.  ]

Strongly activated dimensions: [ 0  3  5  6  9 13 14]

Let's visualize this sparsity pattern across multiple inputs to see how different inputs activate different subsets of hidden dimensions:

Out[14]:

Visualization

Heatmap showing hidden activations across inputs and dimensions, with many zero values demonstrating sparsity. — Hidden activation patterns for 10 different random inputs passing through a small FFN with 16 hidden dimensions. Each row shows the activation values after ReLU for one input. White cells indicate zero activation (the feature did not match), while blue cells indicate positive activations (the key fired). Different inputs activate different subsets of dimensions. This shows that the FFN implements input-dependent sparse retrieval, similar to a content-addressable memory.

The heatmap reveals the key-value memory character directly. Each column represents one hidden dimension (one key-value pair). The white cells are zero, meaning that key did not fire for that input. The blue cells indicate firing keys, with darker blue indicating stronger activation. Notice that no two input samples have exactly the same activation pattern: different inputs retrieve different combinations of stored associations. This input-dependent sparse access is what makes the FFN such a powerful memory structure.

This sparsity also has practical implications for efficiency. In a model with large $d_{ff}$ , only a fraction of hidden dimensions are nonzero for any given input. Methods that exploit this sparsity can skip computation for inactive dimensions, reducing the compute needed well below the theoretical maximum. Recent work on ReLU-based sparse FFNs has shown that sparsity rates of 90% or higher can be achieved with minimal quality loss, meaning only 10% of hidden dimensions are needed for most tokens.