Transformer Feed-Forward Layers Are Key-Value Memories

By Mor Geva, Roei Schuster, Jonathan Berant, Omer Levy. 2020

1,546 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

Citations in the corpus

Locating and Editing Factual Associations in GPT (2022, background): “Some research has been done for masked models (Petroni et al., 2019; Jiang et al., 2020; Elazar et al., 2021a; Geva et al., 2021; Dai et al., 2022; De Cao et al., 2021), but GPT has architectural differences such as unidirectional attention and generation capabilities that provide an opportunity…”
Mass-Editing Memory in a Transformer (2022, methodology): “Based on observations from Geva et al. (2021; 2022) that transformer MLP layers serve as key–value memories, we narrow our focus to them.”
Knowledge Neurons in Pretrained Transformers (2021, methodology): “The key-value-memory nature (Geva et al., 2020) inspires us to propose the knowledge attribution method, which identiﬁes knowledge neurons in feed-forward networks by computing the contribution of each neuron to the knowledge prediction.”
Towards a Unified View of Parameter-Efficient Transfer Learning (2021, background): “We hypothesize that this is because the FFN learns task-specific textual patterns (Geva et al., 2021), while attention learns pairwise positional interactions which do not require large capacity for adapting to new tasks.”
Transformer Memory as a Differentiable Search Index (2022, background): “(Geva et al., 2020) analyzes the knowledge encoded within Transformer feedforward layers.”
On Transferability of Prompt Tuning for Natural Language Understanding (2021, background): “The max(xW T 1 + b1, 0) can be regarded as the non-negative activation values for dm hidden neurons (Mor et al., 2021).”
Towards Automated Circuit Discovery for Mechanistic Interpretability (2023, background): “MLP layers can be modelled as memory that is able to represent certain properties of the network inputs (Geva et al., 2021).”
Eliciting Latent Predictions from Transformers with the Tuned Lens (2023, methodology): “…that “write" (in the sense of Elhage et al. (2021)) to intermediate hidden states of a language model (Millidge and Black, 2022; Dar et al., 2022; Geva et al., 2022; 2020), and even use this analysis to successfully edit model behavior (Geva et al., 2022; Millidge and Black, 2022; Dai et al.,…”
Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small (2022, positive): “Work in mechanistic interpretability aims to discover, understand and verify the algorithms that model weights implement by reverse engineering model computation into human-understandable components (Olah, 2022; Meng et al., 2022; Geiger et al., 2021; Geva et al., 2020).”
Kformer: Knowledge Injection in Transformer Feed-Forward Layers (2022, methodology): “In our paper, inspired by previous work [Dai et al., 2021; Geva et al., 2021] on the feed-forward layers, we propose a novel way to filter and incorporate external knowledge through the feed-forward layers in Transformer.”
Rethinking machine unlearning for large language models (2024, neutral): “Localization in the context of model editing has also been applied to various elements, including neurons (Dai et al., 2021), network layers (Meng et al., 2022; Gupta et al., 2023), and feed-forward components of LLMs (Geva et al., 2020; Li et al., 2023c).”
Dissecting Recall of Factual Associations in Auto-Regressive Language Models (2023, background): “Notably, these findings are consistent with the view of MLP sublayers as key-value memories (Geva et al., 2021; Dai et al., 2022) and extend recent observations (Meng et al., 2022a; Wallat et al., 2020) that factual associations are stored in intermediate layers, showing that they are spread across…”
Locating and Editing Factual Knowledge in GPT (2022, background): “Research has been done for masked models (Petroni et al., 2019; Jiang et al., 2020; Elazar et al., 2021a; Geva et al., 2021; Dai et al., 2021; De Cao et al., 2021), but GPT’s architectural differences (e.”
The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A" (2023, background): “Complementing this, Geva et al. (2021; 2022; 2023) analyze the internal mechanisms behind factual recall in Transformers.”
Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space (2022, background): “These ﬁndings expand on previous results by Geva et al. (2021), who observed that value vectors in the upper layers represent next-token distributions that follow speciﬁc patterns.”
A Survey on Hallucination in Large Language Models: Principles, Taxonomy, Challenges, and Open Questions (2023, neutral): “While significant efforts have been undertaken in exploring their knowledge storage (Geva et al., 2021; Meng et al., 2022)and probing (Petroni et al., 2019; Zhong et al., 2021; Yu et al., 2023c), the exact mechanism by which LLMs capture the factual knowledge remains elusive.”
Editing Large Language Models: Problems, Methods, and Opportunities (2023, positive): “There is existing work that examines the principles that govern how PLMs store knowledge (Geva et al., 2021, 2022; Haviv et al., 2023; Hao et al., 2021; Her-nandez et al., 2023; Yao et al., 2023; Cao et al., 2023; Lamparth and Reuel, 2023; Cheng et al., 2023; Li et al., 2023b; Chen et al., 2023; Ju…”