FFN as Key-Value Memory | Transformers Under the Hood

The Feed-Forward Network as Key-Value Memory

Part 5 of the Transformers Under the Hood series.

For years, the standard mental model of a Transformer’s Feed-Forward Network (FFN) was that of a feature transformer. In this view, the Multi-Head Attention (MHA) layer aggregates context from the sequence, and the FFN applies a non-linear transformation to that aggregated vector. It was seen as a processing unit: input goes in, complex features emerge out. However, recent mechanistic interpretability research suggests a different, more structural role. The FFN functions as a Key-Value (KV) memory store within the transformer layer.

This shift in perspective changes how we understand information flow. Instead of just transforming data, the FFN stores and retrieves specific pieces of information based on signals sent by the attention heads. This mechanism is particularly evident in architectures using swiglu activations, where the gating mechanism acts as a precise addressing system for memory retrieval.

The Retrieval and Write Mechanism

To understand the FFN as memory, we must look at the data flow between the attention layer and the FFN. The MHA layer produces an output vector $z$ for each position. This vector is a weighted sum of value vectors from previous tokens. Crucially, this vector $z$ also serves as the input to the FFN.

In the KV memory interpretation, the attention output $z$ acts as a set of "keys." These keys query the internal weights of the FFN. The FFN consists of two weight matrices: $W_{in}$ (expansion) and $W_{out}$ (projection). The first matrix projects the input into a higher-dimensional space, and the second projects it back down.

When $z$ is multiplied by $W_{in}$, specific neurons activate. If we view the columns of $W_{in}$ as memory addresses, the attention output determines which addresses are accessed. The activation pattern then determines what data is read from those addresses. The final output of the FFN is a linear combination of these activated neurons, effectively retrieving stored values associated with the specific keys provided by the attention mechanism. This retrieval pulls from pre-trained knowledge encoded in the static weight matrix, distinct from the dynamic context in the residual stream.

This creates a feedback loop. Attention retrieves context from previous layers, formats it into keys, and the Multi-Layer Perceptron (MLP) uses those keys to access stored knowledge (values) from its weights. This knowledge is then added back to the residual stream, influencing subsequent attention heads.

SwiGLU as an Addressing Mechanism

The choice of activation function in the FFN significantly impacts this memory mechanism. Standard transformers often use ReLU or GELU. Modern high-performance models, such as LLaMA and Mistral, use SwiGLU (Swish-Gated Linear Unit), originally introduced by Shazeer (2020). SwiGLU changes the FFN structure from a single path to a gated path:

$$\text{SwiGLU}(x) = \text{Swish}(x W_B) \otimes (x W_G)$$

Here, $x$ is the input from the residual stream. $W_B$ and $W_G$ are separate weight matrices. The Swish function is applied to one branch, and the other branch passes through a linear projection. The results are multiplied element-wise.

This gating mechanism is critical for KV memory. In a standard ReLU FFN, many neurons activate simultaneously, leading to a "noisy" memory access where many values are retrieved at once. SwiGLU introduces sparsity and selectivity. The gate branch ($x W_G$) acts as a filter. If the gate is close to zero, the corresponding neuron in the expansion branch is suppressed.

This allows the model to implement sparse key-value pairs. A specific attention pattern might activate only a few relevant neurons in the FFN, retrieving a specific fact or rule without interference from unrelated memories. This precision reduces noise in the residual stream and allows the model to store more distinct facts within the same parameter budget. Research indicates that SwiGLU enables more efficient routing of information, treating the FFN as a sparse lookup table rather than a dense transformer.

Interpretability and Probing

Viewing the FFN as KV memory makes the model more interpretable. If we can identify which neurons correspond to specific keys, we can trace how the model "remembers" information.

Probing techniques involve analyzing the activation patterns of FFN neurons across different inputs. Researchers have found that individual neurons often respond to specific semantic concepts or syntactic structures. For example, a neuron might activate strongly when the attention mechanism provides a key related to a specific entity type (e.g., "capital city").

By mapping these activations, we can construct a "dictionary" of the model's memory. We can see that the attention layer identifies a context (the key), and the FFN layer retrieves a fact (the value). This separation of concerns—retrieval via attention, storage via FFN—explains why transformers can handle long-range dependencies. The attention mechanism handles the "where" (which previous token is relevant), and the FFN handles the "what" (what fact is associated with that context).

This interpretability also helps in debugging. If a model fails on a task, we can check if the attention heads are generating the correct keys. If the keys are correct but the output is wrong, the issue lies in the FFN's ability to retrieve the correct values. Conversely, if the FFN activates the right neurons but the attention didn't provide the key, the failure is in context aggregation.

Conclusion

The Feed-Forward Network is not just a passive transformer of features. It is an active memory component within the transformer architecture. The attention heads provide keys, and the FFN provides values. This KV mechanism, enhanced by SwiGLU gating, allows for sparse, precise, and efficient storage of knowledge. Understanding this mechanism is crucial for ML engineers working on model optimization, interpretability, and fine-tuning. It shifts the focus from treating the FFN as a black box to treating it as a structured database that can be queried, analyzed, and improved.

This perspective clarifies the distinct roles of attention and feed-forward layers. Attention manages the flow of information across the sequence, while the FFN manages the lookup of specific facts within the network's internal state. Together, they form a coherent system for processing and storing knowledge, enabling the remarkable capabilities of modern large language models.