Memformer: The Memory-Augmented Transformer

1 INTRODUCTION . Memory has a fundamental role in human cognition . Humans perceive and encode sensory information into a compressed representation in neurons , and later our brains can effectively retrieve past information to accomplish various tasks . The formation of memories involves complex cognitive processes . Modeling and studying the behavior of human memory is still a challenging research problem in many academic areas . Many researchers have attempted to incorporate memory systems in artificial neural networks . Early works like recurrent neural networks ( RNN ) ( Rumelhart et al. , 1988 ) , including LSTM ( Hochreiter & Schmidhuber , 1997 ) model temporal sequences with their internal compressed state vector as memory . Although RNNs are theoretically Turing-complete , they are limited in preserving the longterm information due to the memory bottleneck . To alleviate the limitation , more powerful memory network architectures such as Neural Turing Machine ( NTM ) ( Graves et al. , 2014 ) , Differential Neural Computer ( DNC ) ( Graves et al. , 2016 ) have been proposed by leveraging a large external memory . However , due to their complex memory addressing mechanism , they are not widely used in NLP . More recently , Vaswani et al . ( 2017 ) proposes Transformer by ditching the use of memory and recurrence . Instead , it maintains all O ( N2 ) dependencies in the sequence with self-attention ( Bahdanau et al. , 2015 ) . Transformer and its followers have achieved great success in various NLP tasks . Nevertheless , the quadratic complexity can be extremely costly when the input sequence is long . Some works address the limitations of self-attention , including Reformer , Sparse Transformer , Longformer , Linformer , etc ( Child et al. , 2019 ; Kitaev et al. , 2020 ; Wang et al. , 2020 ) . They successfully reduce the complexity of self-attention and can process longer sequences . However , the space cost still scales with sequence length , which can not be fully eliminated without memory and recurrence . Transformer-XL ( Dai et al. , 2019 ) re-introduces the concept of memory and recurrence . It caches each layer ’ s hidden states of self-attention into a fixed size queue and re-uses them in the later attention computation . However , the memory as raw hidden states can not effectively compress high-level information . Transformer-XL in practice needs a huge memory size to perform well . Compressive Transformer ( Rae et al. , 2020 ) improves upon Transformer-XL by further compressing its memories into fewer vectors via a compression network . However , as mentioned in the papers , both Transformer-XL and Compressive Transformer still have a theoretical maximum temporal range due to the uni-directional self-attention constraint . In this work , we propose Memformer , which includes a more efficient memory system with a Transformer encoder-decoder architecture . The resulting model has a theoretically unlimited temporal range of memorization . We also improve the relative positional encoding in Transformer-XL with a simplified version . As the traditional back-propagation through time ( BPTT ) has an unaffordable memory cost for our model , we introduce a new optimization scheme , memory replay backpropagation ( MRBP ) , to significantly reduce the memory cost of training recurrent neural networks with large memory . We show that Memformer is compatible with different self-supervised tasks and can further improve its performance on language modeling . Our main contributions can be summarized as follows : ( 1 ) We introduce a new optimization scheme for training recurrent neural networks with large memory and long temporal range . ( 2 ) We propose Memformer , a Transformer-based model , which outperforms the previous Transformer-XL and Compressive Transformer on WikiText-103 language modeling . ( 3 ) We show that Memformer is compatible with a wide range of self-supervised tasks other than autoregressive language modeling . 2 METHODS . 2.1 SIMPLIFIED RELATIVE POSITIONAL ENCODING . The standard attention mechanism involves the dot product between the query vector qi and the key vector kj , where Wq , Wk , Wv are the projection matrices to produce the query , key , and value . TransformerXL proposes a new type of relative positional encoding method . The attention computation is decomposed into four parts : ( a ) content-based addressing , ( b ) content dependent positional bias , ( c ) global content bias , and ( d ) global positional bias . The relative positional embedding Ri−j provides the positional information between every pair of xi and xj . The equation is defined below . u and v are trainable parameters . Ai , j = E > xiW > q WrExj︸ ︷︷ ︸ ( a ) +E > xiW > q WrRi−j︸ ︷︷ ︸ ( b ) +u > WkExj︸ ︷︷ ︸ ( c ) + v > WrRi−j︸ ︷︷ ︸ ( d ) . ( 1 ) However , we observe that ( c ) and ( d ) can be simplified by introducing a bias term to the original query and key projection . Thus , we re-formalize the self-attention , as shown in Eq . 3 . The product of bq and Kx is equivalent to the term ( c ) global content bias . For the term ( d ) , since v , Wr , and Ri−j are all trainable parameters , it can be simplified into the product between bq and bk , which has a similar effect to the global attention bias . Different from Transformer-XL that only injects positional information in the attention computation , our attention mechanism shown in Eq . 4 attends over the positional information and accumulate the results to have more robust output representations . Qx = WqEx + bq ; Kx = WkEx + bk ; Vx = WvEx + bv ( 2 ) Ai , j = Q > xiKxj +Q > xiRi−j ( 3 ) Hx = ∑ j Ai , j ( Vxj +Ri−j ) ( 4 ) 2.2 MEMFORMER . This section explains the details of Memformer . We first talk about the language model background and a new way of formulating language generation with text continuation . Then we describe an instance of such formulation , which is our proposed Memformer model . After that , we introduce the multi-task training setting . Finally , we describe the newly proposed optimization scheme , memory reply back-propagation to tackle the memory cost problem . 2.2.1 BACKGROUND : STANDARD LANGUAGE MODEL . To understand Memformer better , we first study the standard language model . Given a document of N tokens x = ( x1 , x2 , . . . , xN ) , an standard language model learns the joint probability of the document by taking the product of each token ’ s probability conditioned to the previous tokens , which is defined as P ( x ) = ∏ t P ( xt|x1 : t ) . Figure 1a and 1b are the standard language models . They autoregressively predict the next token by feeding the previous generated tokens into the model . An extension of Figure 1a is to incorporate relative positional encoding and cache the past hidden states . Then this model would be equivalent to Transformer-XL . Figure 1b is an assumed language model with memory . Self-attention module now attends not only to its token inputs but also to the memory Mt at time t. After all the tokens in the segment are processed , the model summarizes the computed hidden states in the segment and produce the next timestep ’ s memory Mt+1 . Each layer has its own individual memory representation . One limitation for this model is that the read and write operations on memory may not have enough capacity to retain important information due to the uni-directional attention . 2.2.2 ENCODER-DECODER LANGUAGE MODEL . To address this capacity issue of uni-directional attention , we introduce a more powerful architecture shown in Figure 1c , where we have an encoder-decoder and a memory system . If a document is split into T segments of length L , for each segment st , we define st = [ xt,1 , xt,2 , . . . xt , L ] . The encoder ’ s role is to encode the segment st and inject the information into the memory Mt , while it also retrieves past information from the previous timestep ’ s memory Mt−1 . The final output of the encoder will be fed into the decoder ’ s cross attention layers to predict the token probabilities of the next timestep ’ s segment st+1 as standard language modeling . The definition is as below : Mt = Encoder ( st , Mt−1 ) ( 5 ) P ( st ) = ∏ n=1 : L PDecoder ( xt , n |xt , < n , Mt−1 ) ( 6 ) P ( x ) = ∏ t=1 : T PModel ( st|s < t ) ( 7 ) At each timestep , the process can be deemed as a text continuation task . Given a text segment as the input , the model needs to continue that segment by generating the next text segment . Since the memory stores all the past information , we can autoregressively generate all the text segments in a document . In this fashion , the model can behave as a language model . 2.2.3 MEMFORMER ENCODER-DECODER . To implement the encoder-decoder language model , we propose Memformer Encoder-Decoder . The model incorporates a Transformer encoder-decoder and a memory system . The encoder is equipped with two new modules : Memory Cross Attention ( Figure 2b ) and Memory Slot Attention ( Figure 2c ) to read from or write to the memory respectively . The encoder is fully responsible for encoding and retrieving past information via memory . The decoder then takes the last layer ’ s outputs from the encoder and feeds them into the cross attention module similar to the standard Transformer . For the text continuation task , we let the encoder take the input of the current timestep ’ s text segment , and let the decoder generate the next timestep ’ s segment tokens . Figure 2a shows the detailed structure . Figure 2b demonstrates how Memory Cross Attention module extracts information from the memory Mt with the current segment ’ s tokens X . Each input token ’ s hidden state is projected into queries , while the memory hidden states are projected into key-value pairs . Then the input hidden states will attend over the projected memory key-value pairs to produce the final outputs . This module can effectively retrieve past information from memory Mt given the current text segment . Memory Slot Attention in Figure 2c produces the next timestep ’ s memory Mt+1 . This module takes the inputs of the previous timestep ’ s memory Mt and the encoder ’ s final hidden states . It then projects the memory into queries , keys , and values , while the encoder outputs are into keys and values . Since each memory slot should not be interfering with other memory slots , we design a special type of sparse attention pattern ( details shown in Figure 2c ) . Thus , each slot in the memory can only attend over itself and the encoder outputs . This is to preserve the information in each slot longer over the time horizon . For example , if one slot only attends itself , then the information in that slot will not change in the next timestep . 2.3 MULTI-TASK SELF-SUPERVISED LEARNING . Unlike existing models built either for denoising objectives or language modeling , Memformer can accomplish both types of tasks . This flexibility helps the model learn better representations of the document and strengthen the memory of past information . To avoid conflicts of different tasks , we use separate special tokens for each task . In this work , we only experiment with three self-supervised tasks . We believe that our model is flexible with many other self-supervised tasks to further improve performance . We randomly sample the following three tasks with a probability [ 0.6 , 0.3 , 0.1 ] during training . Text Continuation This is the primary task , as our goal is for language modeling . Given the current timestep ’ s text segment , the model needs to generate the tokens in the next timestep ’ s segment . Text Infilling This task is inspired by BART ( Lewis et al. , 2020 ) . We mask some text spans in a document . The span length is drawn from a Poisson distribution ( λ = 3.5 ) . The span is replaced with a “ [ mask ] ” token . The model needs to predict these masked tokens . Text Recall Reverse of the text continuation task , Text Recall needs to predict the previous text segment given the current timestep ’ s segment . This task aims to directly help the model to better preserve the past information .

Summary: The paper presents a new model for the task of language modeling especially suited for longer sequences. This new model dubbed as Memformer consists of Transformer encoder-decoder and a memory module to store the past information from the encoder outputs. The encoder bidirectionally attends to the immediate previous sequence/segment information and to the memory module, which is designed to capture useful information from the past history of the full sequence. The idea is that by bidirectionally attending simultaneously to the previous input segment and to a memory module, the decoder should be able to improve its generation capabilities.