Feed-Forward Networks as Key-Value Memories: How Transformers Store, Retrieve, and Edit Parametric Knowledge

In transformer language models, multi-head self-attention and feed-forward networks (FFNs) perform distinct computational roles. While self-attention mechanisms route contextual information dynamically across sequence positions, feed-forward layers account for roughly two-thirds of total model parameters in standard architectures. Despite early interpretations characterizing FFNs simply as generic non-linear projections, mechanistic analysis shows that feed-forward layers operate as unnormalized, key-value associative memories.

Understanding the mechanics of FFN memory retrieval explains how neural networks store factual associations, how representation updates accumulate across network depth, and why model editing methods can rewrite stored facts through closed-form weight updates.

Deconstructing the FFN as an Associative Memory

A standard two-layer feed-forward network in a transformer block processes a token representation $x \in \mathbb{R}^d$ from the residual stream using an intermediate hidden dimension $d_m$ (typically $4d$ in standard architectures or $\frac{8}{3}d$ in gated architectures). The canonical layer transformation is defined as:

$$\text{FFN}(x) = \sigma(x W_{\text{in}} + b_{\text{in}}) W_{\text{out}} + b_{\text{out}}$$

where $W_{\text{in}} \in \mathbb{R}^{d \times d_m}$, $W_{\text{out}} \in \mathbb{R}^{d_m \times d}$, and $\sigma$ is a non-linear activation function such as ReLU, GeLU, or Swish.

As demonstrated by Geva et al. (2021) and Sukhbaatar et al. (2019), this matrix multiplication can be decomposed into an explicit summation over individual parameter vectors. Let $k_i \in \mathbb{R}^d$ denote the $i$-th column of $W_{\text{in}}$, and let $v_i \in \mathbb{R}^d$ denote the $i$-th row of $W_{\text{out}}$. The operation becomes:

$$\text{FFN}(x) = \sum_{i=1}^{d_m} m_i(x) v_i + b_{\text{out}}$$

where the scalar coefficient $m_i(x)$ is computed as:

$$m_i(x) = \sigma(x^\top k_i + b_{i,\text{in}})$$

In this formulation, the FFN operates as an unnormalized key-value associative memory:

1. Keys ($k_i$): Vector $k_i$ acts as a learned pattern detector in the representation space. The inner product $x^\top k_i$ measures the degree to which the current token state matches the trigger condition.
2. Memory Activations ($m_i$): The non-linear function $\sigma$ thresholds the match score into an activation coefficient $m_i(x)$, determining how strongly memory slot $i$ fires.
3. Values ($v_i$): Vector $v_i$ acts as the stored memory payload. When triggered, it injects a directional update directly into the residual stream.
4. Aggregation: The layer output is a linear combination of retrieved value vectors weighted by their respective activation coefficients.

Unlike multi-head self-attention, which constructs dynamic keys and values from the input sequence at inference time and normalizes attention weights via softmax across sequence positions, feed-forward layers store static keys and values directly in model weights learned during pre-training.

What Keys and Values Encode

Mechanistic analysis of trained language models reveals distinct semantic structures in the input keys and output values.

Input Keys as Pattern Detectors

Empirical evaluations across models such as GPT-2, RoBERTa, and Llama show that keys $k_i$ correlate with specific textual, syntactic, and conceptual triggers. Research by Dai et al. (2022) on knowledge neurons and Geva et al. (2021) demonstrates that individual keys respond selectively to:

• Lexical and syntactic n-grams: Specific prefixes, punctuation structures, or recurring word sequences (e.g., phrases ending with prepositions or numerical constructs).
• Entity and domain categories: Representations indicating specific entity types, such as programming syntax tokens, geographical entities, or medical terminology.
• Relational contexts: Prompt configurations that bind a subject to a relation (e.g., "The capital of [Entity] is").

Output Values as Concept and Vocabulary Promoters

While keys detect patterns, value vectors $v_i$ define the representation update written to the residual stream. By projecting value vectors into the vocabulary space using the language model's unembedding matrix $W_U \in \mathbb{R}^{d \times |V|}$, the direct effect of an individual memory slot on output token probabilities can be computed:

$$E(v_i) = v_i W_U$$

Research by Geva et al. (2022) and Dar et al. (2023) demonstrated that the top logits in $E(v_i)$ correspond to semantically coherent token clusters. When an activation $m_i(x) > 0$ occurs, the value vector $v_i$ directly boosts the probability of specific tokens in vocabulary space, acting as an additive contribution to the model's final token prediction.

Hierarchical Memory Distribution Across Layers

Parametric memory is not distributed uniformly across transformer layers. Instead, language models organize key-value retrieval hierarchically by network depth:

• Lower Layers (1 to 1/3 depth): Keys trigger on subword n-grams, local syntactic markers, and shallow structural patterns. Value vectors apply baseline corrections to token representations, grounding lexical properties and positional context.
• Middle Layers (1/3 to 2/3 depth): Keys match specific subject-relation bindings, entity types, and domain topics. Value vectors inject factual attributes and relational knowledge into the residual stream (for example, associating an entity with its attributes or geographic facts).
• Upper Layers (2/3 to final depth): Keys detect contextual task framing, high-level intent, and constraint states. Value vectors directly push probability mass onto candidate output tokens in the vocabulary space.

Causal Localization and Model Editing

The key-value memory formulation provides the theoretical foundation for model editing and localized weight intervention.

Using causal tracing across autoregressive models, Meng et al. (2022) demonstrated that factual recall tasks (e.g., answering "The Eiffel Tower is located in...") exhibit localized causal states concentrated in the MLP layers at the final token of the subject entity.

Because the MLP functions as a linear associative memory mapping key states to value updates, factual editing can be formulated as a rank-one matrix update to $W_{\text{out}}$. In algorithms such as ROME (Rank-One Model Editing) and MEMIT (Mass-Editing Memory in a Transformer) (Meng et al., 2022):

1. The target key representation $k_*$ corresponding to the subject entity is extracted at the critical middle layer.
2. The desired target value $v_*$ that maximizes the probability of the new target token is computed.
3. The projection matrix $W_{\text{out}}$ is updated via:

$$W_{\text{out}}^{\text{new}} = W_{\text{out}} + \Delta W$$ where $\Delta W$ satisfies $\Delta W k_* = v_* - W_{\text{out}} k_*$, while minimizing interference with previously learned associations using the uncentered covariance matrix of input activations.

This confirms that factual associations reside within the linear weight mappings of feed-forward networks rather than being diffused arbitrarily throughout attention projections.

Modern Extensions: Gating, MoE, and Activation Sparsity

Modern LLM architectures build on the key-value memory paradigm with several structural refinements:

SwiGLU Gating

Architectures such as Llama, Gemma, and Mistral replace standard FFNs with Gated Linear Units (Shazeer, 2020):

$$\text{FFN}_{\text{SwiGLU}}(x) = \left( \text{Swish}(x W_{\text{gate}}) \odot (x W_{\text{in}}) \right) W_{\text{out}}$$

SwiGLU introduces an explicit gating projection $W_{\text{gate}}$ that acts as a multiplicative filter over the input key matrix $W_{\text{in}}$. This dual-projection design tightens the selectivity of memory retrieval, ensuring that memory values are retrieved only when both the gate and the key conditions align.

Mixture-of-Experts (MoE) Memory Scaling

Scaling parametric capacity without increasing per-token compute is achieved by multiplying the number of FFN memory banks. In MoE architectures such as DeepSeek and Mixtral, the single FFN module is replaced by $E$ independent expert FFNs (Fedus et al., 2022):

$$\text{MoE}(x) = \sum_{e \in \text{TopK}(g(x))} g_e(x) \text{FFN}_e(x)$$

Each expert $\text{FFN}_e$ represents a specialized memory bank. The routing gate $g(x)$ directs tokens to the top-$k$ relevant memory partitions, scaling total available key-value parameters to hundreds of billions while activating only a sparse subset during inference.

Natural Activation Sparsity

Empirical profiling of transformer FFNs indicates substantial activation sparsity during standard generation. Across ReLU and thresholded GeLU/SwiGLU models, typically between 80% and 90% of intermediate activations $m_i(x)$ evaluate to zero or near-zero for any individual token (Song et al., 2024).

This property enables dynamic neuron pruning, activation caching, and speculative memory loading in specialized inference engines, demonstrating that transformer generation relies on sparse lookups from a massive dictionary of parameter-stored knowledge.

Sources

• Geva, M., Schuster, R., Berant, J., & Levy, O. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing (EMNLP).
• Geva, M., Caciularu, A., Wang, K. R., & Berant, J. (2022). Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNLP).
• Sukhbaatar, S., Grave, E., Bojanowski, P., & Joulin, A. (2019). Augmenting Self-attention with Persistent Memory. arXiv preprint arXiv:1907.01470.
• Dai, D., Dong, L., Hao, Y., Sui, Z., Chang, B., & Wei, F. (2022). Knowledge Neurons in Pretrained Transformers. Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (ACL).
• Dar, G., Geva, M., Gupta, A., & Berant, J. (2023). Analyzing Transformers in Embedding Space. Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (ACL).
• Meng, K., Bau, D., Andonian, A., & Belinkov, Y. (2022). Locating and Editing Factual Associations in GPT. Advances in Neural Information Processing Systems (NeurIPS).
• Meng, K., Sharma, A. S., Andonian, A., Yonchev, Y., Belinkov, Y., & Bau, D. (2022). Mass-Editing Memory in a Transformer. arXiv preprint arXiv:2210.07229.
• Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv preprint arXiv:2002.05202.
• Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research (JMLR).
• Song, M., et al. (2024). ProSparse: Introducing and Exploiting Activation Sparsity in Large Language Models. arXiv preprint arXiv:2401.04088.