Augmenting Self-attention with Persistent Memory

1 INTRODUCTION . Transformer networks ( Vaswani et al. , 2017 ) are sequence models that rely on the attention mechanism ( Bahdanau et al. , 2015 ) to capture long term dependencies . Since their introduction in the context of machine translation , they have been applied to many natural language processing tasks , such as language modeling ( Al-Rfou et al. , 2019 ) or sentence representation ( Devlin et al. , 2019 ) . On most of them , they are now surpassing the former state-of-the-art models based on recurrent ( Hochreiter & Schmidhuber , 1997 ) or convolutional networks ( Dauphin et al. , 2017 ) . At their core , transformers use a self-attention layer that forms a representation of the current input by gathering the most relevant information from its context . This layer is repeated along the network depth , allowing for information to flow for long distances and to form rich sequence representations . The self-attention mechanism is often considered as the key component of their success and many have worked on improving transformers by increasing the size of the context captured by those layers ( Wu et al. , 2019 ; Dai et al. , 2019 ; Sukhbaatar et al. , 2019 ) . However , self-attention layers are not the only component of transformer networks and they do not explain the effectiveness of transformers by themselves . Each of these layers is followed by a feedforward layer . These feedforward layers contain most of the parameters of the model . This suggests that their role is probably as important as the self-attention mechanism . In fact , the transformer layer , i.e. , the sequence of self-attention and feedforward sublayers , should be regarded as a single mechanism that gathers information from the context and transforms it into a rich representation . Having such two different layer types of at the core makes Transformer models harder to analyse and understand . In particular , there are not many works exploring the properties of feedforward layers . In this work , we simplify the transformer architecture by revisiting its mechanism , while keeping its properties . We introduce a new layer that merges the self-attention and feedforward sublayers into a single unified attention layer , as illustrated in Figure 1 . As opposed to the two-step mechanism of the transformer layer , it directly builds its representation from the context and a persistent memory block without going through a feedforward transformation . The additional persistent memory block stores , in the form of key-value vectors , information that does not depend on the context . In terms of parameters , these persistent key-value vectors replace the feedforward sublayer . This modification dramatically simplifies the structure of the network with no loss of performance . We evaluate the resulting architecture on standard word level and character level language modeling benchmarks and report performances that are competitive with transformers . Self-attention Feedforward All-attention 2 RELATED WORK . Neural language modeling . Different network architectures have been proposed for language modeling , such as feed-forward networks ( Bengio et al. , 2003a ) , recurrent networks ( Mikolov et al. , 2010 ) , gated convolutional networks ( Dauphin et al. , 2017 ) and transformer networks ( Vaswani et al. , 2017 ) . Of particular interest , Al-Rfou et al . ( 2019 ) apply deep transformers to character level language modeling . Dai et al . ( 2019 ) introduces a caching mechanism , relying on the relative position embeddings from Shaw et al . ( 2018 ) , which makes inference in these models much more efficient for unbounded sequences . More recently , Sukhbaatar et al . ( 2019 ) add a learnable self-attention span to extend the size of the context . Word level language models deal with large vocabularies and computing the most probable word is computationally demanding . Solutions are to either replace the softmax loss with an approximation ( Goodman , 2001 ; Morin & Bengio , 2005 ) , to sample from the vocabulary during training ( Bengio et al. , 2003b ; Jozefowicz et al. , 2016 ) or to include subword units ( Sennrich et al. , 2016 ) . A simple yet effective solution is to replace the loss by a hierarchical softmax designed to better take advantage of the GPU specificities ( Grave et al. , 2017a ) . Finally , many works focus on the regularization of large language models . In particular , Zaremba et al . ( 2014 ) show that dropout ( Srivastava et al. , 2014 ) is effective for recurrent networks . More recently , Press & Wolf ( 2017 ) show that tying the embedding and classifier weights significantly improves generalization . Baevski & Auli ( 2019 ) further show that combining this regularization technique with the adaptive softmax of ( Grave et al. , 2017a ) reduces the memory footprint of a transformer while improving its performance . Attention based models . The attention mechanism was first introduced in the context of mixture of experts by Jordan & Jacobs ( 1994 ) . It is only recently that Bahdanau et al . ( 2015 ) have shown their potential when used in neural networks in the context of machine translation . Since then , this mechanism is commonly incorporated within many models , with applications in natural language processing and computer vision , besides transformers . Sukhbaatar et al . ( 2015 ) apply the attention mechanism on the same sequence , i.e. , the so-called self-attention , in an auto-regressive model called end-to-end memory network . They show their potential in the context of language modeling . Graves et al . ( 2014 ) use the attention mechanism for reading from and writing to internal memory for solving algorithmic tasks . Vinyals et al . ( 2015 ) combine this self-attention mechanism with a recurrent network to solve simple algorithmic problems . Later , Merity et al . ( 2017 ) show that these networks can be used as language models if combined with a cache mechanism ( Grave et al. , 2017b ) . The attention mechanism has been also applied to question answering ( Miller et al. , 2016 ) and image captioning ( Xu et al. , 2015 ) . Finally , Shazeer et al . ( 2017 ) uses the attention mechanism as a mixture of experts in a recurrent network . 3 TRANSFORMER LAYER . A transformer model is made of a stack of identical layers , called transformer layers . Each layer is composed of a multi-head self-attention sublayer followed by a feedforward sublayer . Each sublayer is also followed by an add-norm operation , i.e. , a skip-connection ( He et al. , 2016 ) , and layer normalization ( Lei Ba et al. , 2016 ) . In this section , we review the structure of the transformer layer and refer the reader to Vaswani et al . ( 2017 ) for additional details of the overall model . Multi-head self-attention sublayer . A core mechanism of a transformer network is the multi-head self-attention layer , which consists of multiple attention heads applied in parallel . Each attention head applies the attention mechanism of Bahdanau et al . ( 2015 ) on an input sequence of vectors . More formally , given a sequence x1 , ... , xT of d-dimensional input vectors , each head applies two linear transformations to these vectors to form the key and value vectors : kt = Wkxt , ( 1 ) vt = Wvxt , ( 2 ) where Wk and Wv are the “ key ” and “ value ” matrices of a size dh × d , where dh = d/H is the dimension of a head and H is the number of heads . The key vectors are then used to compute a similarity score between an element t of the input sequence and all the elements of its context Ct . The context can be , for instance , the elements of the sequence that precede t in the case of language modeling , or the whole sequence in the encoder for machine translation . The similarity score between t and an element c of its context Ct is defined as stc = x > t W > q ( kc + p ( t , c ) ) , ( 3 ) where Wq ∈ Rdh×d is the “ query ” matrix , and p ( t , c ) is a position encoding function . There are several ways to encode positions : fixed absolute ( Vaswani et al. , 2017 ) , learned absolute ( Al-Rfou et al. , 2019 ) , and learned relative ( Sukhbaatar et al. , 2015 ; Shaw et al. , 2018 ) . The relative position encoding function improves the efficiency for unbounded sequences , making them useful for language modeling ( Dai et al. , 2019 ) . In this paper , we thus use the relative position encoding defined as p ( t , c ) = ut−c , where ui are position embeddings learned during training . The head then outputs a vector yt by taking the average of the context representations weighted by attention weights atc obtained by applying a softmax function to the similarity scores : yt = ∑ c∈Ct atc ( vc + p ( t , c ) ) and atc = exp ( stc/ √ dh ) ∑ i∈Ct exp ( sti/ √ dh ) . ( 4 ) Note that one can use different position encoding functions for the key and value sides . Finally , the outputs from the different heads are concatenated for each timestep t and multiplied by the d × d “ output ” matrix Wo . The final output of this sublayer is thus a sequence of T vectors of dimension d. Feedforward sublayer . The second element of a transformer layer is a fully connected feedforward layer . This sublayer is applied to each position t in the input sequence independently , and consists of two affine transformations with a pointwise non-linear function in between : FF ( xt ) = U σ ( Vxt + b ) + c , ( 5 ) where σ ( x ) = max ( 0 , x ) is the ReLU activation function ; V and U are matrices of dimension d×df and df × d respectively ; b and c are the bias terms . Typically , df is set to be 4 times larger than d. Add-norm . Both the multi-head self-attention and the feed-forward layer are followed by an addnorm operation . This transformation is simply a residual connection ( He et al. , 2016 ) followed by layer normalization ( Lei Ba et al. , 2016 ) . The layer normalization computes the average and standard deviation of the output activations of a given sublayer and normalizes them accordingly . This guarantees that the input yt of the following sublayer is well conditioned , i.e. , that yTt 1 = 0 and yTt yt = √ d. More precisely , the AddNorm operation is defined as : AddNorm ( xt ) = LayerNorm ( xt + Sublayer ( xt ) ) , ( 6 ) where Sublayer is either a multi-head self-attention or a feedforward sublayer . Transformer layer . The overall transformer layer has the following set of equations : zt = AddNorm ( MultiHead ( xt ) ) , ( 7 ) yt = AddNorm ( FF ( zt ) ) , ( 8 ) where MultiHead is the multi-head self-attention sublayer . This is shown on the left panel of Fig . 1 . 4 OUR APPROACH . In this section , we first show that a feedforward sublayer can be viewed as an attention layer . Then , we take advantage of this interpretation of a feedforward model to concatenate it with the self-attention layer , forming a novel layer that relies solely on a multi-head attention layer without the need for a feedforward sublayer .

Summary: This paper proposes a simple modification to the ubiquitous Transformer model. Noticing that the feed-forward layer of a Transformer layer looks a bit like an attention over "persistent" memory vectors, the authors propose to explicitly incorporate this notion directly into the self-attention layer. This involves concatenating the contextual representations with global, learned memory vectors, which are attended over.