Memformer: The Memory-Augmented Transformer
Authors: Qingyang Wu
Zhenzhong Lan
Jing Gu
Zhou Yu
Venue:  (2021)
Keywords: transformer
language model
memory networks

Abstract: Transformer models have obtained remarkable accomplishments in various NLP tasks. However, these models have efficiency issues on long sequences, as the complexity of their self-attention module scales quadratically with the sequence length. To remedy the limitation, we present Memformer, a novel language model that utilizes a single unified memory to encode and retrieve past information. It includes a new optimization scheme, Memory Replay Back-Propagation, which promotes long-range back-propagation through time with a significantly reduced memory requirement. Memformer achieves $\mathcal{O}(n)$ time complexity and $\mathcal{O}(1)$ space complexity in processing long sequences, meaning that the model can handle an infinite length sequence during inference. Our model is also compatible with other self-supervised tasks to further improve the performance on language modeling. Experimental results show that Memformer outperforms the previous long-range sequence models on WikiText-103, including Transformer-XL and Compressive Transformer.

Under review as a conference paper at ICLR 2021
MEMFORMER: FORMER
THE MEMORY-AUGMENTED TRANS-
Anonymous authors Paper under double-blind review
ABSTRACT
Transformer models have obtained remarkable accomplishments in various NLP tasks. However, these models have efficiency issues on long sequences, as the complexity of their self-attention module scales quadratically with the sequence length. To remedy the limitation, we present Memformer, a novel language model that utilizes a single unified memory to encode and retrieve past information. It includes new optimization scheme, Memory Replay Back- PropaPProa,atton, which promotes long-range back-propagation through time with a significantly reduced memory requirement. Memformer achieves O(n) time complexity and ((1) space complexity in processing long sequences, meaning that the model can han- dle infinite length sequence during inference. Our model is also compatible with other self-supervised tasks to further improve the performance on language modeling. Experimental results show that Memformer outperforms the previ- ous long-range sequence models on WikiTexxt-009 103, including Transformer-XL. and Compressive Transformer.
INTRODUCTION
Memory has fundamental role human cognition. Humans perceive and encode sensory infor- mation into 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 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 & Schmidhuberr 1997) model temporal sequences with their internal compressed state vector as memory. Although RNNs are theoretically Turing-complet, they are limited in preserving the long- term 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 al., 2016) have been proposed by leveraging large external memory. However, due to their complex memory addressing mechanism, they are not widely used NLP.
More recently, Vaswani et al. (017)) proposes Transformer by ditching the use of memory and recurrence. Instead, it maintains all (N2) dependencies in the sequence with self-attention (Bah- danau 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 al., 2020). They suc- cessfully reduce the complexity of self-attention and process longer seguences. However, the space cost still scales with sequence length, which cannot fully climinated without memory and recurrence.
Transformer- XL Dai et al. 2019) re-introduces the concept of memory and recurrence. It caches each layer's 's hidden states of self-attention into fixed size queue and re-uses them in the later atten- tion computation. However. the memory as raw hidden states cannot effectively compress high-level information. Transformer-XI in practice needs huge memory size perform well. Compressive

---

Under review as a conference paper at ICLR 2021
Transformer (Rae et al. 2020) improves upon Transformer-XL by further compressing its mem- ories into fewer vectors via a compression network. However, as mentioned in the papers, both Transformer-XL XL and Compressive Transformer still have theoretical maximum temporal range due to the uni-directional self-attention constraint.
In this work, propose Menformer, which includes more efficient memory system with Trans- former 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 unafford- able memory cost for our model, we introduce new optimization scheme, memory replay back- propagation (MRBP), to significantly reduce the memory cost of training recurrent neural networks with large memory. We show that Memformer compatible with different self supervisee tasks and can further improve its performance on language modeling.
Our main contributions can be summarized as follows: (1) We introduce new optimization scheme for training recurrent neural networks with large memory and long temporal range. (2) We pro- pose Memformer, a Transformer-basee 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 sulf-supe tasks other than autoregressive language model- ing.
METHODS
2.1 SIMPLIFIED RELATIVE POSITIONAL ENCODING
The standard attention mechanism involves the dot product between the query vector and the key vector kj. where W WE, W, are the projection matrices to produce the query, key, and value. TransformerXL proposes new type of relative positional encoding method. The attention compu- tation decomposed into four parts: (a) content-basee addressing, (6) content dependent positional bias, (c) global content bias, and (d) global positional bias. The relative positional embedding R provides the positional information between every pair of x; and xy.. The equation is defined below. and are trainable parameters.
Arj
E W W. W.E.
W W
WaE,
W WrRi- (a)
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 by and K2 is equivalent to the term (c) global content bias. For the term (d), since W,, and Ri-j are all trainable parameters, can be simplified into the product between by and by, which has similar effect to the global attention bias. Different from Transformer-y XL that only injects positional information the attention computation, our attention mechanism shown in Eq, | attendds over the positional information and accumulate the results to have more robust output representations.
Qu = W W.E.r bii K, W.E, + +bi V.= W.E, b,
(2)
Arj Q K2, + +Q R
(3)
H Σ Ars (Va, + +Ra-s)
(4)
2.2 MEMFORMER
This section explains the details of Menformer. 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 hack-propag tackle the memory cost problem.
2

---

Under review as a conference paper at ICLR 2021
X
Xs
Ng Ng
X6
7a
M2
Max3
M
Moxi
M2
Mork
Xg
Xg
X4
X6
% X4
X
X2
X2
Xg
(a) Language Model
Language model with memory
(c) Memformer Ene-Dec
Figure I: Illustrations of different language model architectures. (a) is the standard Transformer decoder. (b) a language model with memory. The decoder takes M, and outputs the next timestep's memory M,+1 at the end of segment. (c) is Transformer encoder-decoder model. The encoder is modified to read and write memory.
2.2.1 BACKGROUND: STANDARD LANGUAGE MODEL.
To understand Memformer better, first study the standard language model. Given document of N tokens x (I1, I2, ,,,..., N, 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) II. P(ra(=111).
Figure La and lb are the standard language models. They autoregressively predict the next token by feeding the previous generated tokens into the model. An extension of Figure [aa is to incorporate relative positional encoding and cache the past hidden states. Then this model would be equivalent to Transformer-XI.
Figure Tblis is an assumed language model with memory. Self-attention module now attends not only to its token inputs but also to the memory M, at time t. After all the tokens in the segment are processed, the model summarizes the computed hidden states the segment and produce the next timestep's memory Mr+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 the uni-directional attention.
2.2.2 ENCODER-DECOME LANGUAGE MODEL
To address this capacity issue of uni-directional attention, we introduce more powerful architecture shown in Figure Ic where we have an encoder-decoder and a memory system. If a document is split into T segments of length L, for each segment 8t, define s; [.t,1, 2f,,1,2r,2.... ... Lt,L). The encoder's s role is to encode the segment 8, and inject the information into the memory M. while also retrieves past information from the previous timestep's memory M,-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 s+41 as standard language modeling. The definition is below:
Mr Encoder(st, M,-1)
(5)
P(sc)
II n=1:L
PDecouerr (Je, [st,en, Mr-1)
(6)
P(x) II (s (=1:2
At each timestep, the process can be deemed as a text continuation task. Given 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 document. In this fashion, the model can behave language model.
3

---

Under review as a conference paper at ICLR 2021
Sugus
Merwry Set
Decoder
Memory Cross Attention
FeedForward
FendFormerd
KV
KV
KV
KV
Q
Q
Loyetteree
Ligethorm
+
Memory,
Menery Croys
(b) Memory Cross Attention
Lyyettom
Layettor
Memory
Memoryevt
+
Memory Slot Attention
Layertors
Layertor
QKV
QKV
QRV
QKC
KV
KV
Seeuus
Merooryy
Ouepue
(a) Memformer Decoderr
(c) Memory Slot Attention
Figure 2:
Illustration of Memformer Encoder- Encoder-Decoder Decoder and its sub-modules. (a) shows the over- all architecture of Memformer Secoder-Decoder (b) demonstrates the memory cross attention for reading from memory. (c) is the memory slot attention module to write the next timestep's memory.
2.2.3 MEMFORMER ENCODER-DECOME
To implement the encoder-decoder language model, we propose Memformer Encoder- Decoder. The model incorporates a Transformer encoder-decoder and memory system. The encoder is equipped with two new modules: Memory Cross Attention (Figure and Memory Slot Attention (Figure to read from 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 text continuation task, we let the encoder take the input of the current timestep's 's text segment, and let the decoder generate the timestep's segment tokens. Figure 2a shows the detailed structure.
attention module similar to the standard Transformer. For the
Figure{bjdemmooe demonstrasss how Memory Cross Attention module extracts information from the memory M, with the current segment's 's tokens X. Each input token's hidden state 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 produce the final outputs. This module can effectively retrieve past information from memory Mj given the current text segment.
Memory Slot Attention in Figure 2c produces the next timestep's memory M,+1. This module takes the inputs of the previous timestep's memory M, and the encoder's final hidden states. It then projects the memory into queries, keys, and values, while encoder outputs are into keys and values. Since each memory slot should not be interfering with other memory slots, we design 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, 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, use separate special tokens for each task. In this work, we only experiment with three self-supervised tasks. We believe that model is flexible with many other self 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 rent timestep's text segment, the model needs to generate the tokens in the next timestep's segment.
This is the primary task, as our goal is for language modeling. Given the cur-

---

Under review as a conference paper at ICLR 2021
Text Infilling This task is inspired by BART (eewis al. 2020). We mask some text spans in a document. The span length is drawn from a Poisson distribution (A = 3.5). The span is replaced with "[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.
2.4 MEMORY REPLAY BACK- POCC-PRRPAGATION
Memformer relies on the explicit memory to encode long-range document, At inference time, there no additional memory cost because of the single unified memory design. Nevertheless, during training, such design would require pack-propagatie timesteps so that the memory writer network can potentially learn retain long-term information. The problem with BPTT is that it unrolls the entire computational graph during the forward pass and stores all the intermediate activations. This process would lead to impractically huge memory consumption for Memformer, which causes training almost impossible.
through time (BPTT) over long range of
A favorable existing approach to climinate this problem is gradient checkpointing (Chen et al, 2016). The algorithm can significantly reduce the memory cost of large computational graph. However, the standard gradient checkpointing still needs to compute all the nodes in the compu- tational graph and store unnecessary hidden states during the forward pass. We propose Memory Replay Back-Propagation (MRBP), a more efficient variant of gradient checkpointing, by replaying the memory at each timestep accomplish gradient back pack-propagttion over long unrolls.
MRBP designed specifically for recurrent neural networks. The algorithm takes an input with rollout [ 1, with length T and the previous memory Mg- MRBP only traverses the critical path in the computational graph during the forward pass. It then obtains each timestep's memory and stores those memories in the replay buffer. During the backward pass, MRBP back- tracks the memories in the replay buffer from time T to 0 and recompute the partial computational graph for the local timestep. It continues the computation of the remaining graph with the output O, to get the loss for back-propagation. There are two directions of gradients for the model. One direction gradients comes from the local back pack-propagation of loss, while the other part comes from the back- propagation of the next memory's 's Jacobin VM The full algorithm is described in Algorithm 14
Algorithm Memory Replay Back- Back- Preck-Propggtiin
Input: rollout=[xo,,I1] .1, æT]: : a list containing each timestep t's input 2t. prevMemory: memory from the previous rollout
b initialize a list to store all the memories computed
replayBuffer=
2 replayBufferaappen((M))
D previous memory
b forward pass
3 for f t= 1,2,..... do
4 Mr+1. = Model(æ,, M,)
No gradient
replay (1+1)
end
& backward pass
VM++1
8 for f T,T I, 1,0
M++1, O (x., M,);
Recompute
10
loss = (tos.
loss.backward()
12
backardcy M++1);
Computes VM,
13 end
save Mr for next rollout's update
5

---

Under review as a conference paper at ICLR 2021
2.5 TEMPORAL RANGE ANALYSIS
We analyze the theoretical maximum temporal range here. Transformer-XL. and Compressive Trans- former store the past hidden states in FIFO queue their memories. However, they have the- oretical limitation for the maximum temporal range when modeling a sequence. Transformer-Xi has maximum temporal range of Nm L, where Nm is the memory size, and is the number of layers. Compressive Transformer extends the temporal range to L x (N + x Ne by compress- ing the memories in Transformer-X XI. into the new compressed memories with size of Nem and compression ratio If sequence is longer than the maximum temporal range, the model will lose information when stored memories are discarded. contrast, Memformer has a single unified memory system, which theoretically has maximum temporal range of infinity.
EXPERIMENTS
3.1 SETTINGS
We conduct all experiments on WikiText- (Merity et al. 2017). which is a popular long-range language modeling benchmark. It contains 28K artieles with average length of 3.6K tokens per article. We adopt byte pair encoding (BPE) (Sennrich et al. 2016) to avoid unknown tokens. The original Transformer-XL XL and Compressive Transformer set the attention length 1,600, close to the average length per article. We argue that such large attention length is against the purpose of using memory to compress information. Therefore, better demonstrate memory efficiency, we set the input size to 64, and restrict the memory size under 128.
Besides, we make the following changes to the baselines. We have Transformer-X XL model: base and large. Transformer-XL. base has 12 layers. Transformer-XL. large has 16 layers. Since Compressive Transformer does not have code released, we re-implement the model following the paper. Our Compressive Transformer has 12 layers. The compression ratio is set to 4. For fair comparison, Memformer Encoder-Decoder has a 4-layer encoder and an 8-layer decoder. For all baselines and our models, the hidden size is set to 512, the feed-forward hidden size to 2048, the number heads to 8, and the head size to 64. We disable dropout as it causes high variance in the final score due randomness, making fair comparisons impossible. We the simplified relative positional encoding for all models as it generally performs better under our setting.
3.2
MAIN RESULTS
Model
Memory Size
#Params
Speed (it/s)
Perplexity $
Transformer-2 XL. base
32
66.69M
1,136
23.52
Transformer-Xet base
64
66.69M
928
22.44
Transformer-XX XL base
128
66.69M
805
21.49
Transformer-Xe XI. large
64
80.35M
745
21.95
Transformer- XL large
128
80.35M
640
21.25
Compressive Transformer Compressive Transformer
16(16)
79.29M
746
22.79
32(32)
79.29M
702
21.98
Memformer Secoder-Decodder
32
76,21M
994
21.26
+ Text Infilling
32
76.21M
994
21.19
+ Text Recall
32
76.21M
994
21.53
+ All Multi-task
32
76.21M
994
21.11
Table 1: Results on WikiText-103. "16(16)" means that the XL memory size and the com-
pressed memory size is 16.
Table shows the results on WikiText-103. We report the number of parameters, the number of items per second as the training speed, and perplexity for comprehensive comparison. Memformer Gecooer--ecoode achieves the best perplexity score with efficient computation and memory trade- off.

---

Under review as a conference paper at ICLR 2021
When increasing Transformer-Xit XL's memory size, observe that the perplexity drops as expected, because the attention length is also increased. Note that the speed decreases with larger mem- ory size. After we have enlarged the memory size of Transformer-XL still worse than Memformer Encoder-Deecoder, and the speed is much slower. Since Memformer Encoder-Decodes has slightly more parameters, compare our model with Transformer-Xe XI. large, which has 16 layers. In Transformer-XL, the number of layers is important for performance, as the maximum temporal range scales with the number of layers. Transformer-Xi large indeed obtains better perplexity scores than Transformer XL-base models. However, our model still achieves better perplexity. Not mention that Memformer Encoder- Encoder-Decoder 55% faster than Transformer-a large. This suggests that Memformer Incoder-Decodee is more efficient in modeling the document than -
128, the perplexity is
Compressive Transformer is another baseline we report in the table. It introduces extra com- pression network to compress the memory hidden states in Transformer-XL. For a fair comparison, Compressive Transformer has half of the memory size for the compressed memory. With the same memory budget, Compressive Transformer performs better than Transformer-XL. However, the ex- tra compression network requires that Transformer- XL is more efficient in terms of the number of parameters and speed under our setting.
number of parameters and computation. We actually find
22.6 VINITRO-- 103 Wildittnn
Eflects
Effects
Soagle Bac
2:1
222
Taltive Tire
SacoPropagatoon Time Harize
Pearaay
(a) Multi-task learning.
(b) Effects of different time horizons (c) Effects of different memory sizes
Figure 3: Effects of different configurations. (a) shows how multi-task learning helps. (b) shows the effects of changing time horizon. (c) shows the effects of changing memory size.
3.3 ABLATION STUDY
We conduct ablation studies explore how each component contributes Memformer's good per- formance, by analyzing the performance improvement of the simplified relative positional encoding and memory replay back-propagation
3.4 MODEL
Effects of Multi-Task Training When combine the three tasks text continuation, text infilling and text recall, the model yields the best performance. We find that when applying only text con- tinuation and text recall, the performance drops. This drop might be because the model over-fits on the text recall task, which hurts the performance of text continuation task. Overall, the performance improvement of multi-task learning marginal. However, in Figure 3al we observe that models trained with multi-task learning have smoother validation curve and is less prone to over-fitting. This indicates that with multi-task learning, the model more robust and potentially learns better feature representations.
Effects of Time Horizon We test how the time horizon for back-propagaie affects the perfor- mance. The results are shown in Figure 3b We vary the back-propagation time horizon from 1 to 32. When the time horizon is set 1, proppaaapoan cannot pass gradients through memory to the previous timestep. Thus, we observe the performance is the worst when the time horizon is 1. As we increase the time horizon, the model achieves better perplexity scores. When the time horizon is increased to 32, observe the marginal improvement perplexity is almost gone.
Effects of Memory Size A large memory size ideally helps to store more information. From Table 3c we can see huge improvement when increasing the memory size from to 8. However, when we further increase the memory size, the perplexity stops decreasing. In future work, we will study how to gain more improvement with larger memory sizes.
7

---

Under review as a conference paper at ICLR 2021
3.4.1 SIMPLIFIED RELATIVE POSITIONAL ENCODING
To test the performance of our simplified relative positional encoding (RPE), only replace the self-attention layers in the original Transformer with the new module without changing other parts. The results in Table[showw that our proposed simplified relative positional encoding has much better performance than the original Transformer-X X RPE in all metrics.
Method
Time(ms)
GFLOPS
Parameters
Perplexity
Method
GPU Memory (MB)
Speed (relative)
XI. RPE 1.27
68.99
1.32M
23.32
BPTT GC MRBP
16.177
x1.00
9,885
x0.48
Ours
0.88
68.72
1.05M
21.74
7,229
x0.90
Comparison for relative positional encoding. The experiments are conducted on Transformer-XL.
Table 2:
Table 3:
Memory Replay Back- Propagation performance comparison.
3.4.2 MEMORY REPLAY
To test MRBP's effectiveness, we compare against the standard back-propagation through time (BPTT) and the standard gradient checkpointing (GC) algorithm. We use Memformer Decoder with 12 layers, 8 heads, 512 hidden size, and 32 memory size for all the experiments here. The time horizon for each truncated back-propagation update is set to 4.
The mack-propagation through time (BPTT) approach is the fastest because it does not need computation However, it costs the most amount of memory due to unrolling the entire computational graph. While gradient checkpointing other two methods (x0.48). In contrast, our MRBP saves more GPU memory with only slight speed degeneration (x0.90). When further increasing the time horizon to 16, only increases 62 MB, suggesting the sub-linear memory growth with the time horizon.
huge amount of memory, it is much slower than the
that the GPU memory
4 RELATED WORK
Optimizing the attention pattern of Transformer is one direction to process long sequences. Child
al. 2019) first proposes Sparse Transformer reduce the computation complexity O(N) with a sparse attention pattern for sequence modeling. Longformer (Beltagy et al.) 2020) and Big Bird (Zaheer al.) (020)) follow Sparse Transformer and explore the effectiveness of different sparsity patterns. Reformer (Kitaev et al. 2020) applies a multi-round locality-sensitiee hashing (LSH) to O(N log N). Linformer (Wang al. 2020) further reduces the complexity to O(N) by observing that self-attention is low-rank and can be approximated with linear attention. However, the memory cost of these approaches still scales with the sequence length.
reduce the computation complexity
Meanwhile, applying recurrence to Transformers is an orthogonal direction comparing to the ef- ficient attention approaches. Recurrence enables the model to have constant memory complexity O(1) during inference. There are mainly two works exploring this direction. TransformerXL. (Dai ct al., 2019) uses relative positional encoding and consists segment-level recurrence mechanism to encode beyond fixed-length context. Compressive Transformer (Rae et al. 2020) extends from Transformer XL by further compressing the previous segment information to achieve longer context. However, they have theoretical maximum temporal range of context related to the memory size and the number of layers. In practice, TransformerXL and Compressive Transformer needs huge memory size to achieve good performance and inefficient in their memory representations.
5 CONCLUSION
In this work, we present Memformer, which takes advantage of memory system to efficiently process long sequences with a linear time complexity and constant memory complexity. Along with Menformer, we introduce new optimization scheme, Memory Replay Back-propagation, which enables training recurrent neural networks with large memory. Our model achieves strong perplexity results on WikiText- 1033. is also flexible to wide range of self- self-supervised learning tasks. With the infinite temporal range capability, we believe Memformer can spark interesting works in domains such lifelong learning and memory-augmentted meta-learning.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio, Neural machine translation by jointly learning to align and translate. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Con- ference Track Proceedings, 2015. URL
Iz. Beltagy, Matthew F. Peters, and Arman Cohan, Longformer: The locgmmmt transformer. CoRR, abs/2004.05150, 2020.
05150
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Trah ory cost. CoRR, abs/1604.06174, 2016. URL ULL//tt:.//aarxxvv.cc/a////50...
with mem-
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequeencesssseeeeeeeuuusuusesee with sparse transformers. CoRR, abs/1904. 10509, 2019.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G. Carbonell, Quoc Viet Le, and Ruslan Salakhut- dinov. Transformer-xl: Attentive language models beyond fixed-length context. In Anna Ko- rhonen, David R. Traum, and Lluís Mårquez (eds.). Proceedings of the 57th Conference of the Association for Compututional Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume Long Pupers, pp. 2988. Association for Computational Linguistics, 2019. doi: 10.18653/v1/pp 1285. URL /dotps://oo....... org. /2//110110 10.116633v1/p19-133
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. CoRR, abs/1410.5401, 2014. (RLLattpp//axxx wrt. haa///ba//a00 /14...110..5500 5401
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwinska, Sergio Gomez Colmenarejo. Edward Grefenstette, Tiago Ramalho, John Agapiou, Adrià Puigdomènech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukcuoglu, and Demis Hassabis. Hybrid computing using neural network with dynamic external memory. Nat. 538(7626)47711 476, 2016. doi: 10. 00.1038//ature20101. URL httthttss://// //doi wwt.ddi..orr///// org/o0 100 /
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput. 178: 1780, 1997. doi: 10. (tt.////ccc..999999..0 1997. 9. 8 1735
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL httphhs://pppeeeieeeeeee net forum? i VK
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy. Veselin Stoyanov, and Luke Zettlemoyer. BART: denoising sequence-to-sene souueeee pre- training for natural language generation, translation, and comprehension. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pp. 7871 7880. Association for Computational Linguistics, 2020. URL https: / www www.ww..wwwbboooo. 2athology/2020.ac1
20tp://22...ccc. -main. 703/
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. https:: //openzeview. net wwtp:/wwwwee..t...eeeeeeee..aaee 772uxxx
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenRe- net wwt./eeuunnteeeeeeeee forummrumii// idd d-SSikikSYDD
view.net, 2020. httpttttps///pepeeeeeeeeee
David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning Representations by Back- Pock-Propggating Errors, pp. 696-699, MIT Press, Cambridge, MA, USA, 1988. ISBN 0262010976.

---

Under review as a conference paper at ICLR 2021
Rico Sennrich, Barry Haddow. and Alexandra Birch. Neural machine translation of rare words with subword units. Proceedings of the 54th Annual Meeting of the Association for Com- putational Linguistics. ACL 2016, August 7-12, 2016, Berlin, Germany, Volume I: Long Pa- doi: 10.18663/1/116-1162. URL
pers. The Association for Computer Linguistics, 2016.
https://doi http://ooo...g....//// 1ttp://0111155///1111111111111
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, UI- rike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: An- nual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 5998-6008, 2017. URL http://papers: wwtppprrr.iii..oom cttpapercppprr
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006. 04768, 2020. URL https://arxxv.orgg 2006. 04768
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. CoRR, abs/2007. 14062, 2020. URL URL https / arxiv wrw.v..rgg//b///007. 14062
10


Reviews:
[{"review_id": "oto51gg1Uog", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nThe 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.\n\nPros:\nThe motivation of the proposed model is interesting, which is to bidirectionally encode the past segment information and a memory module to capture rich signals from the entire history. \n\nCons:\n- The main drawback of this paper is the lack of controlled experiments in the results section. Due to this, fair comparisons to the previous language modeling results in the literature is not possible and thus merits of the approach are neither convincing nor clear.\n- To evaluate performance of the Memformer model, the authors present results on the Wikitext-103 dataset. However, to compare with the perplexity results from baseline models such as Transformer-XL and Compressive Transformer, the authors don't include the results from the original papers and instead re-compute it under simplified settings. However, to report progress in a widely studied task such as language modeling, it is not fair to not compare against highly-cited state-of-the-art results.\n- In Section 3.1, it is mentioned that byte-pair encodings are used to represent words. Are the perplexity results also computed over BPE tokens? If so, then I feel this evaluation scheme is inconsistent with the standard language modeling evaluation of tokens which is done over words (or linguistic units).\n- Although the results in Table 3.2 show that Memformer models obtains small performance gains in perplexity over baselines, but these results are rather marginal improvements and in the absence of statistical significance testing results, it can't be really understood if these are actual performance gains or due to randomness in the training process.\n- As the core of the experimental results are on Wikitext-103, it is unclear if the experimental findings would generalize. Currently, the paper lacks comparisons on other datasets such as PG-19 from Compressive Transformers.\n- The authors propose a multi-task training approach for language modeling. However, the motivation and benefits of why doing this is actually needed is not clearly illustrated and the perplexity gains seem to be small.\n- In Section 2.1, the motivation of simplified relative position encoding is not presented. \n- In paragraph 3, it is mentioned that due to uni-directional attention is Transformer-XL style language models, the memory may not have enough capacity to retain important information. However, this does not necessarily holds true for different tasks. Can the authors include a more detailed explanation or cite a prior work that illustrates this phenomenon?\n\n\nWriting Issues:\n- abstract 1st line: The mention of \"remarkable accomplishments\" is very vague in the context of applicability of Transformer models.\n- compatible with other self-supervised tasks: compatibility does not necessarily imply that such language models would be useful in self-supervised tasks.\n- Introduction section is not well-written. For example, the transition from first paragraph to second paragraph is rather abrupt, paragraph 3: Transformers and its followers: this is a very informal writing style, something that is not really suited when submitting to a publication.\n- Sec 2.2.1: The first sentence in third paragraph - \"Figure 1b is an assumed language model.\" is grammatically incorrect. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "_p6C97oMAhU", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is clearly written and explains the concept well. Keeping a memory part of transformers network enables Memformer to achieve O(n) complexity in time and O(1) in space. It is particularly helpful for handling long sequences. \nThe authors proposed Memory replay Backpropagation scheme which is more effective and efficient than gradient checkpointing. \nThe authors draw a fair comparison with long-sequence models (Transformer XL and Compressive Transformer) on Wikitext-103 dataset.\nThe concept seems to be a mixture of some work done in the past already. However, making it work and slot and cross attentions tricks seems to be a good choice.\n \nWeakness:\n\n1. Lack of experimentation: The authors compared other two transformer networks on just one dataset and presented their results with a hypothesis. The hypothesis is hard to be validated on just one task. The claims that the architecture can be scaled to any task is not justified with any empirical evidence. \n2. In continuation with previous point, the idea usage and applicability is not justified. There are sparse transformers already like Longformer, BigBird that can process long documents too. The lack of experimentation doesnot validate the usefulness of the architecture. \n3. The memory scaling seems to be a problem. Also clearly seen from the graph and the experiments. Effect of memory size should further improve the results which is not there. It is left for future work by the authors which doesnot go well with the hypothesis. Further, the restriction on the memory size is too small. With the advent of larger models, the usage of smaller memory seems implausible. \n\nMinor comments:\n* The multi task sampling ratio [0.6, 0.3, 0.1] during training is not justified. Is it randomly chosen or was different ratios tried?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Uzk9wYK2iwM", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a new style transformer with external memory, which is updated and used through an attention mechanism. They also propose a new algorithm to train the memory, Memory Replay Back-Propagation (MRBP). The memory consists of key-value pair data and is recurrently updated after the segment encoding. Through this memory, it can attend the past knowledge without the limitation of the maximum temporal range. The MRBP algorithm trains the memory through the local back-propagation of loss to reduce memory overhead.\n\nStrengths:\n1. They propose a new type of transformer augmented explicit memory, which contains the entire past knowledge, so it can attend the past knowledge without the computational overhead.\n2. They also propose a new training algorithm, MRBP to update the memory with the local back-propagation of loss, which makes sense and they showed empirical results about how much this algorithm reduces memory overhead efficiently.\n3. They evaluated their method with baselines in the aspect of the number of parameters and training speed and performance for various memory sizes and reported plentiful ablation study results too.\n\nWeaknesses:\n1. I think that section 2.1 requires an additional description: (e.g., in (a), what is the meaning of W_rE_{x_j}?, in (c), how W_k E_{x_j} can be the global content bias?)\n2. This Memformer requires recurrent processing segment by segment, which can be another limitation on training or processing (can require more time than naive Transformer, which can process parallelly). I think that it is a fundamental weakness on this type of Transformer and if there is the comparison with naive Transformer (if the device can train it), it would be interesting for me (I don't think that it is a critical weakness).\n\nThe correctness of their claim and Clarity:\n\nThis paper is well written and looks correct except for some parts like requiring more descriptions or related to the previous method.\n\nAdditional feedback:\n\nThank you for submitting it. I enjoyed reading it. \n\nHowever for me, section 2.1 is hard to understand, if there is more description, then it can be easier. \n\nAnd, you mentioned Transformer XL has a` theoretical maximum temporal range, but by using the hidden states of the previous segment (the authors called a recurrent connection between the segments),  Transformer XL can access the past knowledge. Could you explain more the theoretical maximum temporal range?\n\nThis Memformer looks like the combination of the Transformer and Recurrent module through the explicit memory. Like RNN, it doesn't need incremental memory overhead as the sequence is going on through the memory, and like the Transformer, it can predict the output through the attention modules. As I mentioned in the weakness part, exploring the best combination of RNN and the Transformer can be interesting.\n\nMinor things are\nOn page 6, (L x Nm + c x Ncm -> (L x Nm + c x Ncm)\nOn page 7, when increasing the memory size from 1 to 8 -> 128\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "SOzd6X-Q0eP", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper proposes an encoder-decoder memory-augmented language model, called Memformer. Similar to recurrent neural networks, the memory part is updated as a new input token comes in. The author presents a new optimization method, which is a variant of gradient checkpointing.\n\nReasons for score:\nContrary to the paper's argument, I am not sure that Memformer improves computational complexity or performance on long-range sequence modeling.\n\nConcerns/Questions:\n- Contributions of the newly proposed methods (increased model capacity from memory augmentation, specially designed attention, multi-task self-supervised learning, simplified relative position embedding) to the performance are mixed. A more thorough ablation study is required.\n- To evaluate the effectiveness of Memformer in long-range sequence modeling, I would recommend exploring other settings than WikiTxt-103 with Transformer-XL (or Compressive Transformer). I suggest following experimental settings in the Longformer paper [1].\n- I feel that the memory size of 32 is unsatisfactory. The author should provide a plausible explanation about why increasing memory size larger than 32 is not helpful. To add, Lample et al. propose product key memory using 512^2 memory slots [2].\n- I am curious whether the Memformer using a memory size of 32 is really doing well in long-term modeling even though memory vectors can memorize information from multiple tokens.\n- The setting of the main experiment is not valid. Comparison of Transformer-XL and Compressive Transformer with Memformer by reducing memory size (attention length) is not fair because they are designed to capture long-term context.\n- The number of parameters in Table 1 is quite far from the values in the original Transformer-XL paper. Could you elaborate on it? \n- Could you provide detail of speed measurement in Table 1?\n- Can we try a memory replacement policy like LRU (Least Recently Used) via a sparse update?\n\nMinor comment:\n- The name of memory *cross* attention and memory *slot* attention is not intuitive.\n- Mention of the text continuation task appears in Section 2.2.3 before the definition in Section 2.3.\n\n[1] Iz Beltagy, Matthew E. Peters, and Arman Cohan. “Longformer: The Long-Document Transformer”\n[2] Guillaume Lample, Alexandre Sablayrolles, Marc'Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. “Large Memory Layers with Product Keys”", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new model, called Memformer, that combines the strength of transformer networks and recurrent neural networks. While the reviewers found the idea interesting, they also raised issues regarding the experimental section. In particular, they found the results unconvincing, because of weak baselines, non standard experimental settings (eg. using reporting perplexity results on BPE tokens), or evaluating on only one dataset. These concerns were not well addressed by the rebuttal. For these reasons, I recommend to reject the paper.
Decision: Reject