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 considers an architecture change to the transformer in which they swap the feedforward subcomponent of the standard transformer with an "attention only" variant that includes persistent "memory" vectors. The model is evaluated against a suite of baselines on the tasks of character- and word-level language modeling. Combining this "all attention" approach with adaptive span yields results about equivalent to the SOTA, in some cases with fewer parameters than existing models. The authors do a nice job of presenting ablation results. A key finding here, for example, is that the a model stripped of both persistent vectors and the feedforward sublayer performs poorly.