Source-linked AI summary
Transformer Feed-Forward Layers Are Key-Value Memories
Mor Geva, Roei Schuster, Jonathan Berant, Omer Levy
TL;DR
The function of transformer feed-forward layers has been under-explored despite their large parameter share. The paper analyzes them as key-value memories and finds interpretable pattern detectors whose outputs are composed within layers and refined through residual connections.
Problem
Feed-forward layers contain most transformer parameters, yet little is known about their function in the network.
Method
The paper analyzes feed-forward layers as key-value memories, treating parameter matrices as keys and values and examining their stored patterns and output distributions.
Results
Keys correlate with human-interpretable input patterns, values correlate with corresponding next-token distributions especially in upper layers, and layer outputs aggregate then refine these distributions.
Takeaways & Limitations
Feed-forward layers act as pattern detectors across the model, while the final output distribution is constructed gradually from lower to upper layers.
Takeaways & Limitations
The increasing correlation between value distributions and key patterns across layers may reflect transformations involving self-attention, not feed-forward layers alone.
Abstract
from arXiv · showhide
Feed-forward layers constitute two-thirds of a transformer model's parameters, yet their role in the network remains under-explored. We show that feed-forward layers in transformer-based language models operate as key-value memories, where each key correlates with textual patterns in the training examples, and each value induces a distribution over the output vocabulary. Our experiments show that the learned patterns are human-interpretable, and that lower layers tend to capture shallow patterns, while upper layers learn more semantic ones. The values complement the keys' input patterns by inducing output distributions that concentrate probability mass on tokens likely to appear immediately after each pattern, particularly in the upper layers. Finally, we demonstrate that the output of a feed-forward layer is a composition of its memories, which is subsequently refined throughout the model's layers via residual connections to produce the final output distribution.
1 Introduction
Feed-forward layers, which contain most transformer parameters yet remain under-explored, emulate key-value memories that detect interpretable input patterns and produce output-token distributions. Across layers, memories are composed and residual connections refine the prediction into the final distribution.
- Feed-forward layers emulate key-value memories: keys interact with inputs to produce coefficients that weight values into the layer output.The first parameter matrix corresponds to keys, and the second to values.
- Each key correlates with human-interpretable input patterns, while its value induces an output-vocabulary distribution associated with likely next tokens.For example, one key responds to prefixes describing a period of time ending with “a”, while its value concentrates probability on “while”.
- Each layer combines hundreds of active memories, producing a distribution qualitatively different from any individual value.
- Residual connections refine predictions across layers while retaining most of the residual information.
- The model’s final output distribution is gradually constructed in a bottom-up fashion, with feed-forward layers acting as pattern detectors across layers.
2 Feed-Forward Layers as Unnormalized Key-Value Memories
A transformer feed-forward layer processes each input position independently and is nearly identical to an unnormalized key-value neural memory. Its hidden dimension indexes memory cells, whose input-dependent coefficients weight the stored values.
- A feed-forward layer is a position-wise function that processes each input vector independently.
- Neural memories consist of key-value pairs, with keys and values stored as rows of parameter matrices K and V.
- Feed-forward layers are almost identical to key-value neural memories, differing because canonical transformers omit the memory’s softmax normalization.The hidden dimension d_m corresponds essentially to the number of memories.
- The hidden activation m = f(x · K^⊤) contains an unnormalized non-negative memory coefficient for each memory cell.
- The proposed interpretation asks what memories canonical transformer feed-forward layers store, beyond treating their parameters as persistent memory cells.
3 Keys Capture Input Patterns
Feed-forward keys act as pattern detectors: their strongest training-example triggers share human-recognizable shallow or semantic patterns. Lower layers favor shallow patterns, whereas upper layers increasingly capture semantic ones.
- Each key correlates with a specific set of human-interpretable input patterns, including recurring n-grams and semantic topics.The analysis retrieves examples with the highest memory coefficients and examines the patterns shared among them.
- The study retrieves top-triggering prefixes by ranking training-set prefixes according to their memory coefficients for each key.For each key, researchers computed coefficients over prefixes and annotated the top-25 retrieved examples for recurring patterns.
- Experts identified at least one pattern for every key, averaging 3.6 patterns per key.A majority of retrieved prefixes were linked to recognized patterns through expert annotation.
- 65%-80% of retrieved prefixes were associated with at least one identified pattern, indicating that key-triggering examples commonly share recognizable structure.The reported coverage concerns the top examples associated with each key.
- Lower layers (layers 1-9) are dominated by shallow patterns, often sharing surface forms such as the last word.Upper layers (layers 10-16) instead contain more semantic patterns from similar contexts without clear surface-form similarity.
4 Values Represent Distributions
Values are converted into vocabulary distributions and compared with the next tokens associated with their keys, revealing stronger key-value alignment in upper layers.
- Values are converted into probability distributions over the vocabulary by multiplying each value vector by the output embedding matrix and applying softmax.
- 3.5% agreement emerges by layer 11 after near-zero agreement in layers 1–10, compared with 0.0004% for random vocabulary prediction.
- Upper-layer values increasingly rank the next token from a key’s top trigger example highly, indicating that values complement corresponding key patterns.
- Higher maximum probabilities identify values more likely to agree with their key’s top trigger example, although the distributions are not calibrated.
- 46 out of 100 high-probability values have at least one trigger example whose next token agrees with the value’s top prediction.
- Upper-layer values tend to assign higher probability to tokens appearing next in examples that trigger their corresponding keys, unlike lower-layer values.
5 Aggregating Memories
Feed-forward layers aggregate many active memories into compositional distributions, then residual connections refine predictions across layers. This process often preserves residual information while progressively shaping the model’s final output and confidence.
- 5 Aggregating Memories: Each feed-forward layer combines multiple memories into a distribution qualitatively different from individual value distributions.The analysis uses 4,000 validation prefixes to characterize inference-time behavior.
- 5.1 Intra-Layer Memory Composition: 10%-50% of 4096 dimensions are active per layer on typical examples, while most memory cells remain inactive.The number of active memories drops toward layer 10, where semantic patterns become more prevalent than shallow patterns.
- 5.1 Intra-Layer Memory Composition: At least ∼68% of examples in every layer have a final layer prediction different from every memory’s prediction.This indicates that layer-level outputs typically arise from composing multiple memories rather than one dominant cell.
- 5.1 Intra-Layer Memory Composition: 60% of cases where a memory agrees with the layer prediction involve common stop words, and 43% involve prefixes shorter than 5 tokens.These cases suggest that very common patterns may be cached in individual memory cells without requiring compositionality.
- 5.2 Inter-Layer Prediction Refinement: Roughly a third of model predictions are determined in the bottom few layers, with the proportion growing rapidly from layer 10 onward.The residual’s influence is refined through successive layers, and Figure 10 indicates that confidence is refined alongside top-token identity.
- 5.2 Inter-Layer Prediction Refinement: When residual predictions change, composition usually produces a compromise prediction different from both the residual and feed-forward predictions.In most examples, the residual’s top prediction remains the model’s prediction; last-layer changes were semantically distant in 66 of 100 cases and related in 34.
6 Related Work
Prior work has examined neuron functionality, self-attention, inter-layer differences, and feed-forward-layer importance, but the role of feed-forward layers remains under-explored. This paper instead studies transformer mechanisms rather than explaining individual predictions.
- Neuron and model interpretability: Earlier studies characterized neuron functionality across NLP and vision models, including extraction of key n-grams in CNN text classifiers.These approaches examined properties captured by individual neurons or neuron subsets.
- Transformer architecture: Transformer research has focused on self-attention, differences between lower and upper layers, and the importance of feed-forward layers.Despite this attention, the role of feed-forward layers remains under-explored.
- Interpretability scope: Interpretability methods that explain predictions differ from this paper’s focus on understanding the mechanism of transformers.The paper does not aim to interpret individual predictions.
- Memory-cell analysis: Characterizing memory-cell functionality through examples that trigger maximal activations has precedent in both NLP and vision.This prior strategy is related to the paper’s analysis of feed-forward memory patterns.
7 Discussion and Conclusion
The paper frames feed-forward layers as key-value memories and identifies implications, open questions, and future research directions for understanding transformers.
- Feed-forward layers account for most transformer parameters, yet their network function remains poorly understood.This motivates analyzing them alongside the better-studied self-attention layers.
- The paper proposes that feed-forward layers emulate key-value memories and experimentally links keys to interpretable patterns, values to output distributions, and model predictions to aggregated memory outputs.The proposed account includes residual connections that refine layer outputs throughout the model.
- Layer embedding space: The increasing correlation between value distributions and key patterns from lower to upper layers raises whether transformations in layer output space involve self-attention.The paper notes that feed-forward lookups and weighted aggregation alone would more naturally imply one shared embedding space.
- Beyond language modeling: The key-value-memory formulation is presented as applicable beyond language modeling, but verification across encoders and translation models remains future work.The authors specifically mention BERT encoders and neural translation models as settings for future verification.
- Practical implications: The findings motivate future interpretability, training-data privacy, and architectural research around pattern identification, memory cells, and suppressed correct patterns.The paper connects these directions to automated pattern identification, white-box membership inference, and aggregation behavior.
- The work advances understanding of transformer internals and opens new research threads in modern NLP models.
A Pattern Analysis
The section presents a fully annotated example of trigger prefixes associated with memory cell k5.
- Table 3 annotates 25 trigger prefixes for the memory cell k5.
B Implementation details
The appendix describes the experimental implementation and identifies Table 3 as a pattern-annotation table for memory cell k5.
- The implementation-details section documents the experimental setup for reproducibility.
- Experiments use a 247M-parameter Baevski–Auli language model trained on WikiText-103 with fairseq.
- WikiText-103 contains over 100M tokens extracted from Wikipedia, and spaCy segments examples into sentences for analysis.
- The appendix provides the fairseq and WikiText-103 reference links used for the experimental setup.
- Table 3 is a pattern annotation of trigger examples for memory cell k5.
- The annotated trigger examples classify repetitive patterns as shallow or semantic.