Transformer Feed-Forward Layers Are Key-Value Memories (2021)

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 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.

## introduction transformer-based language models @xcite are at the core of state-of-the-art natural language processing @xcite , largely due to the success of selfattention. while much literature has been devoted to analyzing the function of self-attention layers @xcite @xcite , they account for only a third of a typical transformer's parameters (4d 2 per layer, where d is the model's hidden dimension). most of the parameter budget is spent on position-wise feedforward layers (8d 2 per layer), yet their role remains under-explored. what, if so, is the function of feed-forward layers in a transformer language model? figure 1 : an illustration of how a feed-forward layer emulates a key-value memory. input vectors (here, x 5 ) are multiplied by keys to produce memory coefficients (e.g., the memory coefficient for v 1 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. we find that each key correlates with a specific set of human-interpretable input patterns, such as n-grams or semantic topics. for example, k 2 in figure 1 is triggered by inputs that describe a pe-riod of time and end with "a". simultaneously, we observe that each value can induce a distribution over the output vocabulary, and that this distribution correlates with the next-token distribution of the corresponding keys in the upper layers of the model. in the above example, the corresponding value v 2 represents a distribution that puts most of its probability mass on the word "while". lastly, we analyze how the language model, as a whole, composes its final prediction from individual memories. we observe that each layer combines hundreds of active memories, creating a distribution that is qualitatively different from each of its component memories' values. meanwhile, the residual connection between layers acts as a refinement mechanism, gently tuning the prediction at each layer while retaining most of the residual's information. in conclusion, our work sheds light on the function of feed-forward layers in transformer-based language models. we show that feed-forward layers act as pattern detectors over the input across all layers, and that the final output distribution is gradually constructed in a bottom-up fashion. foot_0 2 feed-forward layers as ## keys capture input patterns we posit that the key vectors k in feed-forward layers act as pattern detectors over the input sequence, where each individual key vector k i corresponds to a specific pattern over the input prefix x 1 , . . . , x j . to test our claim, we analyze the keys of a trained language model's feed-forward layers. we first retrieve the training examples (prefixes of a sentence) most associated with a given key, that is, the input texts where the memory coefficient is highest. we then ask humans to identify patterns within the retrieved examples. for almost every key k i in our sample, a small set of well-defined patterns, recognizable by humans, covers most of the examples associated with the key. ## values represent distributions after establishing that keys capture patterns in training examples, we turn to analyze the information 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 stored in their corresponding values. we show that each value v ℓ i can be viewed as a distribution over the output vocabulary, and demonstrate that this distribution complements the patterns in the corresponding key k ℓ i in the model's upper layers (see figure 1 ). casting values as distributions over the vocabulary. we begin by converting each value vector v ℓ i into a probability distribution over the vocabulary by multiplying it by the output embedding matrix e and applying a softmax: foot_3 the probability distribution p ℓ i is uncalibrated, since the value vector v ℓ i is typically multiplied by the input-dependent memory coefficient m ℓ i , changing the skewness of the output distribution. that said, the ranking induced by p ℓ i is invariant to the coefficient, and can still be examined. this conversion assumes (naïvely) that all model's layers operate in the same embedding space. value predictions follow key patterns in upper layers. for every layer ℓ and memory dimension i, we compare the top-ranked token according to v ℓ i , (argmax(p ℓ i )) to the next token w ℓ i in the top-1 trigger example according to k ℓ i (the example whose memory coefficient for k ℓ i is the highest). figure 4 shows the agreement rate, i.e. the fraction of memory cells (dimensions) where the value's top prediction matches the key's top trigger example (argmax(p ℓ i ) = w ℓ i ). it can be seen that the 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 layer 0 5000 10000 15000 20000 25000 30000 rank distribution figure 5 : distribution of the rank of the next-token in the top-1 trigger example of k ℓ i (w ℓ i ), according to the ranking induced by the value vector v ℓ i . we cut the tail of the distribution, which stretches up to the vocabulary size (∼270k tokens). next, we take the next token of k ℓ i 's top-1 trigger example (w ℓ i ), and find where it ranks in the value vector's distribution p ℓ i . figure 5 shows that the rank of the next token of a trigger example increases through the layers, meaning that w ℓ i tends to get higher probability in the upper layers. ## aggregating memories so far, our discussion has been about the function of a single memory cell in feed-forward layers. how does the information from multiple cells in multiple layers aggregate to form a model-wide prediction? we show that every feed-forward layer combines multiple memories to produce a distribution that is qualitatively different from each of its component memories' value distributions (section 5.1). these layer-wise distributions are then combined via residual connections in a refinement process, where each feed-forward layer updates the residual's distribution to finally form the model's output (section 5.2). value prediction precision@50 trigger example v 15 222 each 68% but when bees and wasps resemble each v 16 752 played 16% her first role was in vijay lalwani's psychological thriller karthik calling karthik, where padukone was cast as the supportive girlfriend of a depressed man (played v 13 2601 extratropical 4% most of the winter precipitation is the result of synoptic scale, low pressure weather systems (large scale storms such as extratropical v 15 881 part 92% comet served only briefly with the fleet, owing in large part v 16 2070 line 84% sailing from lorient in october 1805 with one ship of the line v 12 3186 jail 4% on may 11, 2011, four days after scoring 6 touchdowns for the slaughter, grady was sentenced to twenty days in jail ## related work considerable attention has been given to demystifying the operation of neural nlp models. an extensive line of work targeted neuron functionality in general, extracting the properties that neurons and subsets of neurons capture @xcite @xcite @xcite , regardless of the model architecture or neurons' position in it. @xcite analyzed cnn architectures in text classification and showed that they extract key n-grams from the inputs. the study of the transformer architecture has focused on the role and function of self-attention layers @xcite @xcite and on inter-layer differences (i.e. lower vs. upper layers) @xcite . previous work also highlighted the importance of feed-forward layers in transformers @xcite @xcite . still, to date, the role of feed-forward layers remains under-explored. also related are interpretability methods that explain predictions @xcite , however, our focus is entirely different: we do not interpret individual predictions, but aim to understand the mechanism of transformers. characterizing the functionality of memory cells based on examples that trigger maximal activations has been used previously in nlp @xcite . ## discussion and conclusion understanding how and why transformers work is crucial to many aspects of modern nlp, including model interpretability, data security, and development of better models. feed-forward layers account for most of a transformer's parameters, yet little is known about their function in the network. in this work, we propose that feed-forward layers emulate key-value memories, and provide a set of experiments showing that: (a) keys are correlated with human-interpretable input patterns; (b) values, mostly in the model's upper layers, induce distributions over the output vocabulary that correlate with the next-token distribution of patterns in the corresponding key; and (c) the model's output is formed via an aggregation of these distributions, whereby they are first composed to form individual layer outputs, which are then refined throughout the model's layers using residual connections. our findings open important research directions: thus, by illuminating the role of feed-forward layers, we move towards a better understanding of the inner workings of transformers, and open new research threads on modern nlp models.