Transformer Feed-Forward Layers Are Key-Value Memories

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 5484–5495 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 5484 Transformer Feed-Forward Layers Are Key-Value Memories Mor Geva1,2 Roei Schuster1,3 Jonathan Berant1,2 Omer Levy1 1Blavatnik School of Computer Science, Tel-Aviv University 2Allen Institute for Artificial Intelligence 3Cornell Tech {morgeva@mail,joberant@cs,levyomer@cs}.tau.ac.il, rs864@cornell.edu Abstract 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 tex- tual 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 pat- terns, while upper layers learn more semantic ones. The values complement the keys’ in- put 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 Transformer-based language models (Vaswani et al., 2017) are at the core of state-of-the-art natu- ral language processing (Devlin et al., 2019; Brown et al., 2020), largely due to the success of self- attention. While much literature has been devoted to analyzing the function of self-attention layers (Voita et al., 2019; Clark et al., 2019; Vig and Be- linkov, 2019), they account for only a third of a typ- ical transformer’s parameters (4d2 per layer, where d is the model’s hidden dimension). Most of the parameter budget is spent on position-wise feed- forward layers (8d2 per layer), yet their role re- mains under-explored. What, if so, is the function of feed-forward layers in a transformer language model? We show that feed-forward layers emulate neural memories (Sukhbaatar et al., 2015), where the first k1 k2 kdm v1 v2 x .2 1.5 0 self-attention layer Stay with you for while ... it will take a ... every once in a ... , and for a FF layer a 5 x4 x3 x2 x1 x’5 x’4 x’3 x’2 x’1 vdm Transformer layers Transformer layers Figure 1: An illustration of how a feed-forward layer emulates a key-value memory. Input vectors (here, x5) are multiplied by keys to produce memory coefficients (e.g., the memory coefficient for v1 is 0.2), which then weigh distributions over the output vocabulary, stored in the values. The feed-forward layer’s output is thus the weighted sum of its values. parameter matrix in the layer corresponds to keys, and the second parameter matrix to values. Figure 1 shows how the keys (first parameter matrix) inter- act with the input to produce coefficients, which are then used to compute a weighted sum of the val- ues (second parameter matrix) as the output. While the theoretical similarity between feed-forward lay- ers and key-value memories has previously been …

Content truncated to 500 words