Memory Representation in Transformer
Authors: Mikhail Burtsev
Yurii Kuratov
Anton Peganov
Grigory V. Sapunov
Venue:  (2021)
Keywords: transformer
memory augmented networks

Abstract: Transformer-based models have achieved state-of-the-art results in many natural language processing tasks. The self-attention architecture allows transformer to combine information from all elements of a sequence into context-aware representations. However, information about the context is stored mostly in the same element-wise representations. This might limit the processing of properties related to the sequence as a whole more difficult. Adding trainable memory to selectively store local as well as global representations of a sequence is a promising direction to improve the Transformer model. Memory-augmented neural networks (MANNs) extend traditional neural architectures with general-purpose memory for representations. MANNs have demonstrated the capability to learn simple algorithms like Copy or Reverse and can be successfully trained via backpropagation on diverse tasks from question answering to language modeling outperforming RNNs and LSTMs of comparable complexity. In this work, we propose and study few extensions of the Transformer baseline (1) by adding memory tokens to store non-local representations, (2) creating memory bottleneck for the global information, (3) controlling memory update with dedicated layer. We evaluate these memory augmented Transformers and demonstrate that presence of memory positively correlates with the model performance for machine translation and language modelling tasks. Augmentation of pre-trained masked language model with memory tokens shows mixed results for tasks from GLUE benchmark. Visualization of attention patterns over the memory suggest that it improves the model's ability to process a global context.

Under review as a conference paper at ICLR 2021
MEMORY REPRESENTATION IN TRANSFORMER
Anonymous authors
Paper under double-blind review
ABSTRACT
Transformer-based models have achieved state-of-the-art results in many natural language processing tasks. The self-attention architecture allows transformer to combine information from all elements of sequence into context-aware repre- sentations. However, information about the context is stored mostly in the same element-wise representations. This might limit the processing of properties related to the sequence as whole more difficult. Adding trainable memory to selectively store local as well as global representations of a sequence promising direc- tion to improve the Transformer model. Memory-augmented neural networks (MANNs) extend traditional neural architectures with general-purpose memory for representations. MANNs have demonstrated the capability to learn simple algorithms like Copy or Reverse and can be successfully trained via backpropaga- tion diverse tasks from question answering to language modeling outperform- ing RNNs and LSTMs of comparable complexity. In this work, we propose and study few extensions of the Transformer baseline (1) by adding memory tokens store non-local representations, (2) creating memory bottleneck for the global in- formation, (3) controlling memory update with dedicated layer. We evaluate these memory augmented Transformers and demonstrate that presence of memory posi- tively correlates with the model performance for machine translation and language modelling tasks. Augmentation of pre- trainerie masked language model with mem- ory tokens shows mixed results for tasks from GLUE benchmark. Visualization of attention patterns over the memory suggest that it improves the model's 's ability to process global context.
INTRODUCTION
Transformers (Vaswani et al., 2017) are extremely successful in a wide range of natural language processing and other tasks. Due to the self-attention mechanism transformer layer can be trained to update vector representation of every element with information aggregated over the whole sequence. As a result, rich contextual representation for every token is generated at the end of encoding. However, combination of local and global information in the same vector has its lim- itations. Distributed storage of global features results in "blurring" and makes harder to access them. Another well-known deficiency of Transformers is poor scaling of attention span that hurts its applications to long sequences.
In our work, we propose and study a simple technique augment Transformer with memory rep- resentation Memmraansormmr). We extend the Transformer baseline by adding [mem] tokens at the beginning of the input sequence and train the model to if it is able use them universal memory storage. To assess the capacity of proposed memory augmentation, we additionally applied number of other architectures. In the MemCtrl model update of [mem] tokens is controlled by dedicated Transformer layer. MemBottleneck model has removed attention between sequence el- ements, thus making memory the only channel to access global information about the sequence. We also tested memory augmented BERT (Devlin et al. 2019) and Transformer XL models.
(Dai et al., 2019)
Our work lies at the intersection of two research directions Memory-augmented neural networks (MANNs) and Transformers. The history of memory augmentation in neural networks is pretty long. Classic LongrShort Term Memory (LSTM) (Hochreiter & Schmidhuber 1997) can be seen as a simple yet powerful form of fine-grained memory augmentation with single memory value per LST cell and memory control logic implemented by internal learnable gates. Thus, in LSTMs,

---

Under review as a conference paper at ICLR 2021
computations are heavily intertwined with memory. In contrast that, memory- neu- ral networks incorporate external-memory, which decouples memory capacity from the number of model parameters. Neural Turing Machines (NTMs) (Graves et al. 2014) and Memory Networks (Weston et al., 2014) are among the best-knows MANNs that provide powerful random access op- erations over external memory. Memory Networks (Weston al. 2014, [Sukhbaatar al., 2015) are trained to iteratively reason by combining sequence representation and embeddings in long- term memory with the help of attention. NTMs, and their successors Differentiable Neural Com- puter (DNC) (Graves et al. 2016) and Sparse DNC (Rae et al. 2016) are recurrent neural networks equipped with content-addressable memory. similar to Memory Networks, but with the additional capability to write to memory over time. The memory is accessed by controller network, typi- cally an LSTM. The full model is differentiable and can be trained via back-propagatin
through time (BPTT). There is also a line of work to equip neural networks (typically, LSTMs) with data structures like stacks, lists, or queues (Joulin & Mikolov) 2015 Grefenstette et al., 2015). MANN architectures with a more advanced addressing mechanisms such address-content separation and multi-step addressing were proposed in (Gulcehre et al. 2016 2017 Meng & Rumshisky 2018).
Family of Transformer models have been recently applied to many deep learning tasks and proved to be very powerful for the language modeling tasks. The core element of Transformers is self- attention that allows updating representation for every element with information aggregated over the whole sequence. Self-attention scales as O(N2) with a sequence length, and as a result, it is severely limited in application to long sequences.
There is a separate line of work dedicated to reducing the computational cost of the transformer attention to O(N VN) using sparsity (Child et al., 2019). O(N log N) with local-sensitive hashing (Kitaev et al. 2020) even O(N) with low-rank approximations (Wang et al.) 2020), kernel-based
formulation Katharopoulos et al. 2020), or sparse attention with randomness ZZhherr et al. 2020)
Several recent approaches try to solve this problem by adding some kinds of memory elements to their architecture. Transformer-XL (Dai et al. (019) adds segment-level recurrence with state reuse, which can be seen as a sort of memory. During training, the hidden state sequence computed for the previous segment is fixed and cached to be reused as an extended context when the model processes the next segment. Compressive Transformer (Rae et al.] 2019) extends the ideas of Transformer- XL by incorporating the second level of the memory into the architecture. Memory on the second level stores information from the short-term memory of the first level in compressed form. Memory Layers (Lample et al. 2019) replace forward layer with product key memory layer, that can increase model capacity for negligible computational cost.
Some transformers introduce different sorts of global representations. Among the most recent archi- tectures with global representations are Star-Trans Tran-rranfrrmer (Guo et al. 2019). Longformer (Beltagy et al. 2020), Extended Transformer Construction (ETC) (Ainslie
2020 and its successor Big Bird (Zaheer et al., (020). All these architectures reduce full self-attention to some local or patterned attention and combine it with sparse global attention bottleneck. For example, Longformer uses selected tokens such as [CLS] ] or tokens for question marks to accumulate and redistribute global information to all other elements of the sequence. Among these, the BigBird-ETC with dedicated "global" tokens is the most similar to our MemTransformer approach.
MemCtrl and MemBottleneck Transformer models can be seen as more gen- eral limit cases for this class of models. They have dedicated general purpose [mem] tokens that can be used by the model as a placeholders to store and process global or copy of local representations. Mem Transfoorrre has full self-attention over the memory+input sequence. In contrast, MemBottle- neck has full both-way attention between the input sequence and memory but no attention between sequence tokens.
Our MemTransforme
2 MEMORY IN TRANSFORMER
2.1 BACKGROUND: TRANSFORMER ARCHITECTURE
The process of calculating single Transformer self elffaateenioon layer can be seen as a two-step process- ing flow (see fig. [p).
2

---

Under review as a conference paper at ICLR 2021
b. Mem Transformer layer
mam
vectors
[sequence] updiate
Mem Bottleneck Transformer layer
Transformer layer
[sequence] attention
mem
token vectors
[sequence]
[sequence] update
mem toklnss
[sequence] attention
[sequence] attention
MemCtri Transformer layer
[mem] update
mem
token vectors
[mom] attention
tokions
[mem] update
[sequence] updiate
[mom] attention
[requence] attention
sequarce
Figure I: Memory modifications of Transformer architecture. (a) Transformer layer. For ev- ery element of sequence (solid arrow), self-attention produces aggregate representation from all other elements (dashed arrow). Then this aggregate and the element representations are combined and updated with fully-connected feed-forward network layer. (b) Memory Transformer (Mem- Transformer) prepends input sequence with dedicated [mem] tokens. This extended sequence is processed with standard Transformer layer without any distinction between [mem] and other ele- ments of the input. (c) Compared to Mem Transforrrr MemCtrl Transforemer has dedicated mem- ory controller sub-network. (d) Memory Bottleneck Transformer (MemBottleneck Transformer) uses [mem] tokens but separates memory and input attention streams. At the first step, representa- tions of [mem] tokens are updated (2) with the attention span (I) covering both memory and input segments of the sequence. Then representations of input elements are updated (4) with memory attention (3) only. Thus information flow is distributed representations of elements only through the memory.
1. Self-attention. Calculate normalized sum of input X with multi-head attention (Q, V) between all elements of the sequence:
= LN + XXXX,
(I)
2. Update. For every element of the sequence update aggregated representation A with FF feed-forward sub-layer then add skip connection and normalize:
H = LN(AA++F(AA))
(2)
2.2 SIMPLE MEMTRANSFORMEE
The first proposed model is simple extension of baseline Transformer we call MemTransformer. The idea to add m special [mem] ory tokens to the standard input (see fig. [T) then process them in standard way. So, the input vectors X became the concatenation of the memory token vectors X mem and the original input token vectors Xseq;
meem+ mem X=e9]= E I (m+x)) X E RRmx xeey E Raad
This modification can be applied independently to encoder and/or decoder. The rest of the Trans- former stays the same with the multi-head attention layer processing the extended input.
3

---

Under review as a conference paper at ICLR 2021
2.3 MEMCTRL TRANSFORMER
In the simple Mem Transformer tokens of the memory and the sequence are processed by layers with the same parameters. Thus memory and sequence updated in similar way. To test if dedicated sub- network for memory update might improve performance we introduce a separate memory control layer (see fig. Te.. Thus, memory representation of MemCtrl Transformer is updated as:
A mem
LN (X mem + M mmmm (X " X mem + seg
X Xmee-ssey
Hmem LN (A mem F pmem Amem))
Sequence representation is updated as: A = LN (X seg MH (X (x x memt+ee X Xmem+ser))
Heay (N seg F Feee A ))
2.4 MEMBOTTLENECK TRANSFORMER
In the Tranrfosfrrmer input and [mem] tokens are updated inside the same traditional self-attend and update processing flow. In this case, representations of the input sequence elements potentially might be updated "as usual" without attending to the content of the memory. Here, global infor- mation can propagate in a "peer to peer" manner. To block this distributed information flow and separate storage and processing of global and local representations, The resulting MemBottleneck Transformer has two-staged processing flow (see fig. It).
add memory bottleneck.
1. Memory update. First, calculate attention between every memory token and full sequence of memory Xmmem and input X00q (see Step on the fig. [d), and update memory token representations (see Step 2 on the fig. TT):
MM (X Xme, m x mem++sg seg X X m FF pmem Amem))
A Amem
= LN (X mem
mmmm=se))),
mem LN mem
Sequence update. Calculate attention between sequence and memory (Step 3 on the fig. [.) and update sequence token representations (Step 4 on the fig. It):
A seg
LN(X #aa + .M Haeq (X (xaeq, mmm Hmmm)),
Heeg LN(A + FF
In other words, the memory "attends" to itself and a sequence, and the sequence "attends" only to the memory. This should force the model to accumulate and re-distribute global information through memory. Computations for MemBottleneck Transformer scales linearly with the size of the input sequence or memory (M) when the traditional transformer scales as O
For all encoder-decoder variants of the memory transformers the decoder part was the same in the baseline. Output of the last encoder layer [H mem, [=*0]] passed to the decoder layers.
RESULTS AND DISCUSSION
As a reference model for machine translation task use vanilla Transformer from official TensorFlow I Two model sizes were studied for machine translation task small with N 4 and base with N = 6 layers the encoder. The decoder has the same number of layers as the encoder. For language modeling task we augmented Transformer XI. (Dai et al. 2019) base with 20 [mem] ] tokens. For a masked language model memory augmentation we used pre-trained BERT
(Devlin et al. 2019). All values reported in the paper averaged over 3 runs if otherwise stated.
dmodel = 128, d = 512, 8, Parop = 0.1, batch = 64, wormupptepe 4000
dmedel = 512. 2048, h 8. Parop 0.1. batch 64, warmuppstess 32000
https://github: com wwwnuggg////oooooooo t r rans former-xi
checkpoint from HuggingFace Transformers Wolf al. 2020) was trained with
DeepPavlov (Burtsev al. 2018) GLUE tasks.
4

---

Under review as a conference paper at ICLR 2021
Table I: Performance of baseline and memory models on WMT-14 DE-EN translation. Values represent an average BLEU scores for runs every model evaluated on 2000 samples from WMT-14 DE-EN validation set.
Small models
Base models
layers per encoder/decodee, 20 epochs
layers per encoder/decoder, 10 epochs
Transformer (baseline)
19.01
Transformer (baseline)
24.65
Mem LommT'rann Transformerr
19.17
Mem nem'ransformer 10
19.15
McmTransformer 10
25.07
Mem Transforrrr 20
19.14
Mem'rransformrr 20
25.58
MemBottleneck Traantormer 10
11.20
McmCtrl Transformer 20
24.13
MemBottleneck Transformer 20
10.41
MemCtrl Shared Transformer 20
25.73
MemBottleneck Skip Transformer
16.45
3.1 PERFORMANCE METRICS
The main hypothesis of the study says that adding memory to multilayered encoder-decoder archi- tectures should result in better performance for sequence processing tasks such as machine transla- tion. BLEU scores for WMT-14 DE-EN translation task (Bojar et al. 2014) are presented in Table After 20 epochs of training, small MemTransformer models have similar scores and clearly outper- form the Transformer baseline. Base 6-layer MemTransffrmer with 10 memory tokens improves the baseline, and doubling the memory up to 20 tokens results in an even higher score of 25.58. This is modest but solid performance given hyperparameters fine tuning and beam search were used. Transforrree results supports our intuition that self-attention could be trained to utilize representations of extra memory elements that not related to the input sequence to improve the quality of encoding. Surprisingly, adding separate layer for memory control decreases scores below baseline (see MemCtrl 20 in Table T.. On the other hand, memory controller with shared parameters for all 6 encoder layers (MemCtrl Shared 20 in Table I demonstrates the best performance among modifications we studied for this task.
The Mem Transformer results suggest that if memory extends but not intervene in the Transformer sequence processing, then it is beneficial. But to what extent processing of relations between ele- ments of the sequence can be abstracted to memory? Experiments with MemBottleneck Transformer (Table [ shows that it is possible, but performance suffers. This can be due the more complex architecture of the MemBottleneck that has twice more layers in the encoder part (see Pp.). So, it more difficult and longer to train compared to baseline. On the other hand, degraded performance can also be attributed to the insufficient throughput of the memory bottleneck. Then, there might be a trade-off between the size of the bottleneck and the complexity of learning a deeper network. From the experiments, we see that MemBottleneck 10 learns faster and has lower loss compared to MemBottleneck 20, which points to the complexity of training but not the bottleneck width as a major factor limiting performance.
The limit scenario for the MemBottleneck model is when only memory representations are pro- cessed. In MemBottleneek Skip modification of Transformer, representations for sequence tokens are not updated at all (steps and 4 on the fig. [T are skipped) and encoder output consists of input sequence embeddings and memory representations. Quite unexpectedly, leaving only in memory processing in encoder significantly improves 10.41 BLEU score of MemBottleneck 20 to 16.45 (MemBottlenncck Skip in T..
Memory models have better scores after training, but do they require memory for inference? If the performance of trained MemTransformer will stay the same for the inference without [mem] tokens, then memory was only needed to improve training and not used for the processing of input sequence. Results of memory lesion experiments presented Table 2 demonstrate that removing [mem] tokens from to 11.75 for the TraTransformer 10 and from 25.58 to 3.87 for Transforrme 20 (both models have 6-layers in the encoder). This is an indicator that the presence of [mem] tokens is critical for MemTransformer during inference.
input Icads to dramatic drop BLEU score, from 25.07
Another important question is related the universality of the learned memory controller. Is it able to utilize memory of arbitrary size, or can work only with the memory capacity was trained? Memory lesions data (see Table suggest that MemTransformer learns solution that is partially
5

---

Under review as a conference paper at ICLR 2021
Table 2: Memory lesions. Performance of the models trained with memory gradually degrades if the memory size is changed during inference.
memory size
inference
5
10
20
30
MeaTransfommer 10
11.75
15.91
18.22
25.07
12.39
7.87
Mem' I Transformer 20
3.87
8.38
9.75
14.51
25.58
7,51
Table 3: Memory extension. Values represent an average of BLEU 4 scores for rums of every model.
20 epochs +5 epochs
+10 epochs +13 epochs
mem
10
mem
20
Mem'rransformer (small)
19.17
19.18
19.19
19.41
Table 4: Memory augmentation for language modeling task. Average performance after training on WikiText-103 (Merity et al. 2016) over 3 runs.
Transformer-Xi.
20 mem fixed pos. emb.
20 mem rel. pos. emb.
bpc ppl 24.09
3.182
3.179
3.176
24.02
23.95
robust to the variations in the size of the memory. BLEU score of MemTransformnn 10 with [mem] tokens shows it is still able to produce translation that makes sense On the other hand, we add 20 more [mem] tokens to the same model, it will have scores that are lower even compared to the case when the model is evaluated without memory at all. Interestingly, the model trained with a larger memory size of 20 has weaker generalization abilities. It is evident from the more steep decline of performance with the deviation memory size from the one used during training.
As from memory ablation results (Table 2 increasing memory without fine tuning hurts performance. But, what will happen the model will be fine-tuned after memory extension? To answer this question memory in a few stages up to 20 [mem] tokens. On every stage 5 [mem] tokens were added and the model was fine tuned for 5 epochs. Results presented in Table 3 Extension of the memory followed by fine tuning proved to be beneficial and resulted in the model with the highest BLEU score among all small sized modifications.
take small Mem'Transformer pre-trained for 20 epochs and grow it's
To test an effect of mem tokens on performance in a language modelling task we trained Trans- former XL base augmented with memory size 20. Original Transformer XL has fixed and relative positional encodings, so results for the both options and the baseline are presented in the Table Memory augmentation allows the model achieve better perplexity.
Positive memory extension results suggested experiments with memory augmentation of an already pre-trained encoders. We took BERT-base model and augmented it with a memory of different sizes. The model was trained on datasets from the GLUE (Wang et al., (018) benchmark. Adding [mem] tokens to BERT-base model improved performance on 6/9 tasks as shown the Table
Table 5: Results on GLUE dev set with [mem] tokens added only for end task fine-tuning, Each [mem] was randomly initialized and trained only on the GLUE task. All runs were repeated times and avorage scores
reported. +pool stands for using concatenation of max and avg pooling over outputs for [mem] tokens
instead of the output from [CLS] token for classification.
CoLA
SST-2
MRPC
STS-B
QQP
MNLI-m/mm
QNLI
RTE
BERT-base
62.9
92.7
90.2/85.8
86.0/85.8
86.6/89.8
83.0/83.5
90.5
65.0
5mem
61.3
92.4
90.4/86.4
86.0/85.8
86.87 86.8/
82.7/83.3
90.7
68.0
Smem+pool
62.1
89,4/84.8
85.8/85.6
86.9/90.2
83.3/83.3
90.8
60.2
10mem
60.6
92.5
91.3/87.6
86.6/86.4
86.4/89.8
82.8/83.3
90.5
66.8
10mem+pool
62.6
92.6
90.2/86.0
86.7/86.5
87.1/90.2
83.1/83.0
90.7
61.2
20mem
60.9
92.4
91.2/87.5
86.4/86.2
8698/90.1
82.8/83.1
90.7
65.3

---

Under review as a conference paper at ICLR 2021
Verbessenung
Enttwicktung
könstiiche
Intelliganz
wesentich
meigetragen
Land
Schtrii
ist
aima
Technik
pur
der
von
omemm
<mamz
cmemm
4mems
<mom>
<mems
<mema
<mama
estart
Lams
<mem>
Long
<mem>
neighbour
representations
short
<mem>
<mem>
<mem>
write
<mem>
<mem>
15
<start>
Lang
technique
that
contributed
significantly
read from memory
<mem>
<mem>
neighbour
<mom>>
<memm
fuse neighbour representations
Improving
representations
development
<mem>>
<mem>>
store
of
<mem>
copy
<mem>
<mem>
<mem>
<mem>>
representations
Tuse neighbour
Intelligence
Figure 2: Operations with memory learned by Mem Transformer 10. (a) The pattern of self- attention in the 3rd encoder layer. Here, [mem] tokens in the central segment of memory (on the left) attend to the vector representations of tokens Technik, Entwicklung Intelligenz (and some others). This attention values are consistent with the writing of selected token vectors to the [mem] tokens. Activity in the left top corner that involves first four tokens might indicate fusion of neighbour vectors by pairwise summation of [mem] tokens. (b) In the next 4th layer of the same encoder similar fusion operation with the same [mem] 's is repeated. A parallel diagonal activity (c) Another attention head in the same encoder layer demonstrates combination of fusion and store operations. Sharp self-attention of three tokens in the middle results in adding vectors to themselves. (d) Attention pattern in decoder layer over the output of 6th encoder layer suggest that vectors [mem] tokens are selectively read and added to the output token representations during decoding.
just below the fusion pattern can be attributed to copy operation.
3.2 ATTENTION PATTERNS IN MEMORY
Generic system with memory relies on three types of operations, such as writing, reading and pro- cessing. In this section we present results of analysis of the inner workings of memory augmented transformers to localize these operations. Following previous studies (Kovaleva et al.) 2019 Clark et al. 2019), visually explored attention patterns. Kovaleva al. (2019) introduced five cate- gories of self-attention and suggested that only "heterogeneous" patterns that spread attention across all input tokens might extract non-trivial information about the linguistic structure. Numerous ob- servations of attention maps across bascline Transformer and Mem MemTransformer models allow us to conclude that the overall structure and distribution pattern types in the sequence sequence part of the attention mechanism are similar. Thus, for further analysis, we skip sequence to sequence attention and focus on memory to sequence, memory to memory, and sequence memory attention patterns. All attention maps for selected models are presented in the Appendix.
Memory sequence attention makes it possible to selectively update vectors stored in [mem] token positions with representations of input sequence elements. Such an update is a form of soft write to memory operation. Indeed, we found many patterns consistent with writing from sequence to memory in all MemTransformer models. One of them is shown in Figure 2.. Write memory type of attention more frequent in the first layers and almost absent in the deeper part of the encoder.
Memory to memory attention allows recombining vectors of [mem] tokens. We found a few common patterns related to in memory processing. The most common arrangement of activity is diagonal. The diagonal can blurred (or "soft"), making local fusion of the neighboring memory representations possible. Examples of this operation be in the left top corner of the figures

---

Under review as a conference paper at ICLR 2021
P.. and If diagonal attention is sharp (see fig. 2 in the middle), then corresponding memory vectors are added to themselves, so their content is amplified and became more error-tolerant. This can be seen a store operation. Another possible operation is block copy (examples can be found in the Appendix). It is usually manifested as a vertical block of attention. number of consequent [mem] vectors are updated with the same values aggregated from some another sequence of [mem] vectors. A copy operation can also be performed in a [mem] to [mem] manner with preserving source order as in figure Pp. or with reversing the order (see Appendix for examples).
this case,
Sequence memory attention implements read from memory operation and can be found in the first layers of the encoder, but it is more pronounced in the middle layers of the decoder. A typical example of the memory "reading" presented in Figure 21. Note, that during decoding token representation is updated by reading from a block of subsequent [mem] tokens.
The overall pipeline of memory processing similar for the different runs and sizes of MemTrans- former. It consists of writing some information from the input sequence to the memory in the first layers of the encoder, then followed by memory processing in the intermediate layers and amplifica- tion in the output layers of the encoder. During decoding, information is read from memory. Here, the highest "reading" activity is commonly observed in the intermediate decoder layers. Interest- ingly, memory-relater attention patterns usually have block structure. For example, patterns form the particular BemTTansformee 10 presented in Figure suggest that the model had learned to split memory into three blocks. Commonly, the same memory operation is applied to all [mem] of the same block by one particular head. During memory processing, the model can operate in block- wise manner, as in Figure 2 where the "block(1-3)" copying to the "block(4-6)". We speculate that the block structure of memory processing might reduce error rate because the memory repre- sentation is "averaged"" over the [mem] s of the block during reading (see fig. D1). Experiments with MemBottleneck architecture show that the model might be able to learn how to copy representations of the input sequence into the memory of the fixed size and use only this memory during decoding.
CONCLUSIONS
We proposed and studied a series of memory augmented transformer based architectures Mem'rrans- former, MemCtrl and MemBottleneck transformers. Qualitative analysis of attention patterns pro- duced by the transformer heads trained to solve machine translation task suggests that both models successfully discovered basic operations for memory control. Attention maps show evidence for the presence of memory read/write as well as some in-memory processing operations such as copying and summation.
A comparison of machine translation quality shows that adding general- memory in Mem- Transformer improves performance over the bascline. Moreover, the final quality positively corre- lates with the memory size. On the other hand, MemBottleneck Transformer, with all self-attention restricted the memory only, has significantly lower scores after training.
Memory lesion tests demonstrate that the performance of the pre-trained MemTransformer model critically depends on the presence of memory. Still, the memory controller learned by the model degrades only gradually when memory size is changed during inference. This indicates that the controller has some robustness and ability for generalization. We also found, that extension of memory followed by fine tuning leads to better performance.
Application of proposed technique to language model training as well as fine-tuning of BERT based encoder for battery of GLUE tasks further demonstrated beneficial effect of memory augmentation. This suggests that simple addition of [mem] tokens can extend almost any incoder-decoder with attention" framework. It can be also applied to the tasks that depend on the multi-hop reasoning or planning. In this cases memory should help to store and process representations for the intermediate stages of the solution.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Philip Pham, Anirudh Ravula, and Sumit Sanghai. Etc: Encoding long and structured data in transformers, 2020.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long long-document transformer, 2020.
Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow. Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ale Tamchyna. Findings of the 2014 workshop on statistical machine trans- lation. In Proceedings of the Ninth Workshop Statistical Machine Translation, pp. 12- 58, Baltimore, Maryland, USA, June 2014. Association for Computational Linguistics. URL www.waallweb.. org/ antaanathoooggoo /w//////w h/W14/W14 W14 3302
Mikhail Burtsev, Alexander Seliverstov, Rafael Airapetyan, Mikhail Arkhipov. Dilyara Baymurzina, Nickolay Bushkov, Olga Gureenkova, Taras Khakhulin, Yurii Kuratov, Denis Kuznetsov, et al. Deeppavlov: Open-source library for dialogue systems. In Proceedings of ACL 2018, System Demonstrations, pp. 122-127, 2018.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers, 2019.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning, What Does BERT Look at? An Analysis BERT's Attention. Proceedings of the 2019 ACI. Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, 276-286, Florence, Italy, 2019, Association for Computational Linguistics. doi: 10.18653w1//199-4833. URL
9912419--9989888 4828
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context, 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Hu- man Language Technologies, Volume (Long and Short Papers), pp. 4171-4186, 2019. URL
htttss:/acclwan orq/ http:///tt..ooooo///////////////
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines, 2014.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiáska, Sergio Gómez 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 Hass- abis. Hybrid computing using a neural network with dynamic external memory. Nature, 538 (7626):471-476, October 2016. ISSN 00280836. URL http://dx doi wrt/oorg11111... 1038/ nature201011
Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom. Learning to transduce with unbounded memory, 2015.
Caglar Gulcehre, Sarath Chandar, Kyunghyun Cho, and Yoshua Bengio. Dynamic neural turing machine with soft and hard addressing schemes. arXiv preprint 1rXiv:1607.0000 2016.
Caglar Gulcehre, Sarath Chandar, and Yoshua Bengio. Memory augmented neural networks with wormhole connections. arXiv preprint arXiv: /7Xv0:/701.087/18 2017.
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star- transformer, 2019.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8): 1785-1780, November 1997. ISSN 0899-7666 7667. doi: 10. (116////oo.19979... 1735. URL
https://ddi.org/10 10....111/neeeee 1997. 9. 1735

---

Under review as a conference paper at ICLR 2021
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets, 2015.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention, 2020.
Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer, 2020.
Goro Kobayashi, Tatsuki Kuribayashi, Sho Yokoi, and Kentaro Inui. Attention Module Not Only Weight: Analyzing Transformers with Vector Norms. arXiv:2004.10102 [cs]. 2020. URL
http:://axiss beeg/abs//00 org/abss/200..1110
Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the Dark Se- crets of BERT. In Proceedings of the 2019 Conference on Empirical Methods in Natural Lan- guage Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4356-4365 Hong Kong, China, 2019. Association for Computational Linguisties. doi: 1/186653/1/D11-1444 URLhttps. www. www.wwwwwwccocccc org/ /atghoootth/////
D19-1445
Guillaume Lample, Alexandre Sablayrolles, Marc' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys, 2019.
Yuanliang Meng and Anna Rumshisky. Context-aware neural model for temporal information traction. In Proceedings of the 56th Annual Meeting of the Association for Computational Lin- guistics (Volume I: Long Papers), pp. 527-536, 2018
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint (003:v166900788800 2016.
Jack W Rae, Jonathan J Hunt, Tim Harley, Ivo Danihelka, Andrew Senior, Greg Wayne, Alex Graves, and Timothy Lillicrap. Scaling memory-augmented neural networks with sparse reads and writes, 2016.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling, 2019.
Sainbayar Sukhbaatar, Arthur Szlam, Jason Weston, and Rob Fergus. End-to-end memory networks, 2015.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in neural informa- tion processing systems, pp. 9998,6008, 2017. URL httt:://pappp:: nips cc/ /ppper// 71el-attentiooniss-ei you--need
Alex Wang, Amapreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neu- ral Networks for NLP, pp. 353-355, Brussels, Belgium, 2018. Association for Computational Linguistics. URL http: ac web. wrw.org//aaooo hwwgg//athllggy////oooo W1 5444
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks, 2014.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. HuggingFace's Trans- formers: of-te-- Natural Language Processing. 19X1::1910.0377) [es], February 2020.
URL http:://rx//ii.. org/
03771
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences, 2020.
10

---

Under review as a conference paper at ICLR 2021
A ATTENTION MAPS FOR MEMORY AUGMENTED TRANSFORMERS
In this section we present attention maps for two representative cases of NeemTransfornee and Mem- Bottleneck transformer models. For both models we use the same imput sequence.
Input sequence: Langes Kurzzeingedüchhnis ist eine Technik, die zur Verbesserung der Entwicklung von künstlicher Intelligenz, wesentlich beigetragen hat
Predicted translation Mem Transformer 10: Long, short-term memory is a technique that has contributed significantly to improving the development of artificial intelligence.
Predicted translation MemBottleneck 20: The short time memory is technique that has helped improve the development of artificial intelligence in lot of sense.
Reference: Long-term short-term memory is a technique that has contributed significantly to im- proving the development of artificial intelligence
A short guide to help with interpretation of attention maps is shown on the Figure
input (source) elements
Entwiikhuugg
intolliganz
wssontich
begetragen hat
comeent
Technik
<and>
Lang
Rurz
ene
ist
<mem>
<mem>
[memory]
[sequence]
<mem>
<mem>
to
<mems
[memory]
[memory]
<mem>
<mem>
<mems
<mem>
Lang
output (target) elements
[memory]
[sequence] to [sequence]
zeitg
ed
[sequence]
achtnis
1st
ene
Technik
Verbesserung
Entwicklung
von
kinstliche
Intelliganz
wesentlich beigetragen
hat
<end>
Figure 3: How to read Memory Transformer attention map. Attention values indicate how ele- ments of input sequence (on the top) contribute the update of representation for specific output ement (on the left). Attention map for memory augmented transformer can be split into four blocks: (I) update [sequence] to [sequence]; (2) write [sequence] to [memory]; (3) read [memory] to [sequence]; (4) process- [memory] to [memory].
A.1 MEMTRANSFORMER ATTENTION MAPS
Visualisation of attention maps for Trarrsooorerrr 10 (see Section 2.2 with a memory size of 10 presented the Figure 4 for 6 layers of encoder and the Figure for 6 layers of decoder. Every transformer layer has attention heads. The model was trained for 10 epochs on WMT-14 DE-EN (Bojar et al. 2014) dataset.
https://translate..oom wow.://lllsat...ggee.ooooooo.

---

Under review as a conference paper at ICLR 2021
MemTransformer encoder
(self-attention maps)
layer
layer
Hit Hul
Witbord Bitkk
Aull badiin
Alll handingg
Anill tedt
Aillle
aillket Aitii
Hillieed
pilllerd Aitif
Jrill Min
mith
a
Mill Jnn
dull and
autbal Haila
autheta
Layer
Layer
Layer 6
minal llata
a
4 Aith
liut
Milll In
Jtill ma
In Ial m
Iilll to H
Buk
a. Hath.
a Bitk.
Millles
mill
Mill mnd
authela
a Aillfi
Aill Sedttn
Anil aillastddii te
Hill halnma
a
ammat
Figure 4: MemTransforrrr 10 encoder attention maps. As the model encodes input quence the change of attention patterns related to memory can be interpreted as a read-process-stonn storess-store pipeline. Heads in layers to 3 have many read to memory patterns. Patterns consistent with in ory processing are more frequent in layers 3-6. The last layer is dominated by diagonal attention that can be seen amplification of calculated representations.
12

---

Under review as a conference paper at ICLR 2021
MemTransformer decoder
(attention maps output layer of encoder)
Layer 1
Layer 2
Layer
Matila
Malla
Matl
Jummite Lod Matila
Hitilllsdd ald Mati.
mumm Matll
Maril
millilggddd Matia
mumntnt Maria
4
Layer
Layer
Malila
Mathl
of
0000000 mc MALE
mom Maril
millihee Marik
Imiililend Mitkk
jumii Mart
mumtaatt Mari
Figure 5: MemTransformer 10 decoder attention maps. Every layer of the decoder has heads with signs of memory reading activity. Reading patterns suggest that the representations memory locally grouped in 3 blocks.
13

---

Under review as a conference paper at ICLR 2021
A.2 MEMBOTTLENECK TRANSFORMER ATTENTION MAPS
Attention patterns generated by MemBottleneck Transformer architecture (see Section 2.4) strongly suggest that the model learned copy given sequence into memory, process it and only this representations of input for decoding. The main idea of MemBottleneck is a restriction of global information exchange to memory. Therefore, update for representations of the input sequence elements can access representations of other elements only by writing into and then reading them from memory. To do that, MemBottleneck uses two different transformer sub-layers each with its' own set of attention heads (see fig. T.)
Encoder attention maps (see fig. 6 suggest that, expected, representations for the input elements are copied into memory in layers and 2. Surprisingly, after that they are not properly updated anymore and the decoder mostly attends the content of memory (see fig. 7 This impressive outcome shows that transformer can be trained to read and process all the information about the input sequence in memory only.
14

---

Under review as a conference paper at ICLR 2021
MemBottleneck Transformer encoder (attention maps)
Layer
Layer
Imemory sequence] to (memoryl
[memory] [requence]
[memory sequence] [memory] Iliik aumuttheed MuL.
[memory] to [sequence] Hangaaae
Jun
Marla
Jink
Murt.
linn
Layer
Layer
(meeorr sequence] lo (memmry] Bull
[memory] (sequence)
[memorr sequence) [memoryy JArr.
[memmrr]] to [eeqeecee]
litk
Jun.
Bink
Mut.
link
Figure 6: MemBottleneck 20 encoder attention maps. In the 1st layer, all attention heads of memory sub-layer ([memory+scquence] to [memory]) read from the input sequence. Only heads of memory sub-layer in the layer reads from the input, but all others are diagonal to amplify content of the memory. No more input reading is present in layers and 4. Notably, all heads of the 1st layer memory attention have patterns that split into three blocks. The top block has sparse attention over the whole sequence without preserving the order. The middle block reads the first half of the sequence in the reverse order, and the bottom block reads the rest in the proper order. This suggests encoding of global information in the top block and local information in the middle and bottom blocks. Layer memory sub-layer has sharp amplifying diagonals. and something like shifting operations represented by broken diagonals. Layer 4 of memory sub-layer demonstrates mainly heterogeneous patterns which indicates in memory processing. Maps of attention which belongs to the sequence sub-layer ([memory] to [sequence]) of MemBottleneck layer degrade to vertical lines layers 3 and This a sing that these attention heads are bypassed as follows from (Kobayashi et al. 2020).
15

---

Under review as a conference paper at ICLR 2021
MemBottleneck Transformer 20 decoder
(attention to the output layer of encoder)
Layer
Layer
al cand
Mad
| All
Layer
Layer
ii al litt B
al Hit
pom at Ala Bb
iiliiiiiii al ditt B.
Figure 7: MemBottleneck 20 decoder attention maps. the decoding phase, almost all heads attend to the content of memory but not on the representations of sequence elements.
16


Reviews:
[{"review_id": "sfmSW-_b85H", "reviewer": "AnonReviewer1", "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": "The paper presents transformer extensions with global memory.  The basic idea is to augment the input sequence with special memory token [mem].  The authors explore three variants: (a) MemTransformer - just the augmented sequence, (b) MemCtrlTransformer - separate params for memory update, and (c) MemBottleneck - Bottleneck of memory used to update token representations. The authors show gains over vanilla transformer/pretrained baseline for machine translation, language modeling, and the GLUE setup. \n\nI think the idea in general is interesting but it's not particularly novel. The authors cite a lot of recent works in this area, and a particularly similar approach has been explored by Gupta et al in \"GMAT: Global Memory Augmentation for Transformers\" (authors should cite this work). The main problem with the paper is that the empirical results are quite weak and don't show a consistent trend with regards to utility of memory. \n\nApart from the machine translation results, it's hard to make a case that the gains on GLUE and language modeling are significant. For the MT results, the baseline results look too weak and the authors don't cite any prior work. Through a quick scan of MT papers, I found papers reporting +6 BLEU scores on the same task and the same architecture (see https://arxiv.org/pdf/1904.09324.pdf).  This is concerning because the baseline model seems highly undertuned. \n\nAmong the three variants, the authors only provide results for MemTransformer on language modeling and GLUE. \nFor machine translation, some of the entries in Table 1 are missing, and the small and base model don't have the same entries.\n\nOther minor comments:\n* What do the authors mean by 2000 validation \"segments\"? Are the segments just sentences? Is it the standard dev set? If not, then why not? If it's the standard dev set, why not refer to it at as WMT-14 DE-EN dev set.\n* The authors don't specify the number of attention heads in main text. I presume it's 8, as mentioned in appendix.\n* In Table 1, the authors add a result with \"MemCtrl Shared Transformer\" which they say is the MemCtrl model with shared parameters. Isn't this the same as MemTransformer?\n* The authors say in conclusion that \"the speed of training positively correlates with the memory size\". They don't provide any evidence for this, and if they train for the same number of epochs then increasing memory size should only increase the training time. \n* Again in conclusion, the authors say that \"the memory controller learned by the model degrades only gradually when memory size is changed during inference.\" That is clearly not true given the results in Table 2. \n\nSuggested text edits:\n* Section 1 - \"They have unspecific [mem] tokens that can store global or copy of local information.\" -> Rephrase this/Expand on this. \n* Section 2.2 - \"It minimally changes the original model architecture.\" -> There's no architecture change. Get rid of this line.\n* Section 3 - \" if the opposite is not stated.\" -> \"if otherwise stated.\"\n* Section 3.2 - \"Kovaleva et al. (Kovaleva et al., 2019) \" -> Use \\citet\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "_0S1DsO1_kU", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "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: The paper proposes to study three formulations (MemTransformer, MemCtrl and MemBottleneck) of memory-augmented self-attention transformers, and investigate the influence of adding memory tokens to the model to its overall performance. The authors claim via some experiments on MT and LM that memory augmentation is able to improve the canonical transformer models.\n\n-------------------------\n\nPost-rebuttal thoughts:\n\nWhile the rebuttal of the authors did clarify some points, I think it is undeniable that the empirical results presented in the current version of the paper are still far from being convincing enough to claim this form of memory augmentation is a useful addition to the Transformers. The visualization of the memory augmentation via attention map is also too handwavy. \n\nWhile I appreciate the authors' effort to address my questions in the rebuttal phase, I will keep my current score.\n\n-------------------------\n\n**My general opinion**: \n\nI personally think that memory augmentation is a simple but potentially useful technique that explicitly adds memory to a deep network, which, as the authors mentioned in their paper, is not a new topic. Memory augmented networks have been around in many deep models, especially sequence models. However, I find this paper still not strong enough for me to recommend its acceptance in its current form, mainly because 1) the formulations are rather incremental; and 2) the empirical results are too weak to demonstrate (convincingly) the usefulness; and 3) I found some experimental settings uncommon and empirical analysis too handwavy. My detailed question and comments are below.\n\n-------------------------------------\n\nI have some detailed questions/comments.\n\n1. In the MemBottleneck formulation, does the update rule of $A^\\text{seq}$ take $X^\\text{mem}$ or $H^\\text{mem}$? Figure 1d seems to suggest that the memory update at a layer $i$ precedes the update on the sequence tokens. In addition, what is the difference between $H$ and $X$? Are they not the same sequence (i.e., the $H$ from a layer is passed in as $X$ of the next layer)?\n\n2. In the MemBottleneck subsection, the paper claims that its cost \"scales linearly with the size of the input sequence $O(N)$\" instead of $O(N^2)$. But if the hypothesis of the authors are correct, that these memory modules serve to collect and \"re-distribute\" the sequence information, shouldn't one generally expect the memory size to be proportionate to the sequence length? For example, I certainly would expect that a sequence of length 8000 would certainly require a different memory size than a sequence length 40. The point is, one cannot simply assume \"the size of the memory is constant\", just like one cannot assume \"the size of the sequence is constant\". The more accurate way, for instance, is to say the complexity is $O(NM)$ with $M$ being the memory size.\n\n3. My major concern with this paper is with its various empirical results and experimental settings:\n   - i) In machine translation, how exactly does the encoder-decoder self-attention work in the three proposed Mem settings? E.g., does the decoder queries attend to $X^\\text{mem}$ only (of the encoder output), or also the $X^\\text{seq}$?\n   - ii) How does MemCtrl models perform in the small setting of Table 1? (They seem to be slightly better on the larger setting)\n   - iii) Overall, I found the numbers reported for WMT'14 de-en in Table 1 **too low** for me to convince of anything. Even though the *much more commonly used setting is WMT'14 en-de*, it should be overall easy to get a >29 BLEU score with a base Transformer model (e.g., [1,2] got >31 for the base Transformer). I am not sure why the authors halted their experiment at 20 epochs and 10 epochs (how many training steps are these though?). Hence, it's not clear to me whether these minor improvements in the below-expectation BLEU scores are truly indicative of the final performance of these models. \n    - iv) The paper says no beam search was used. Why? What are the results if you use beam_size = 5?\n    - v) The MemBottleneck Skip Transformer is really just a weak form of a typical Transformer, except that you initialized the first layer of the hidden units to 0, and downsample the sequence from the length of $X^\\text{seq}$ to the length of $X^\\text{mem}$. So it's not a surprise to me that \"MemBottleneck Skip Transformer 20\" is significantly better than MemBottleneck 10/20. \n    - vi) The various ablation studies the authors made in the paper, in my opinion, exactly suggest that the usefulness of these memory augmentations are in doubt. For example, the huge gap between MemBottleneck and canonical transformer, as well as the negative gap between MemBottleneck 10 and MemBottleneck 20, I think is good evidence of this. Another example is the Table 2, where inference time with a larger memory size even degrades the performance. I didn't find a satisfactory explanation from the authors on this.\n    - vii) What are the values in Table 3? Are they BLEU scores? Why are they <1?\n    - viii) The paper says that \"BLEU score of MemTransformer 10 with 5 `[mem]` tokens shows it is still able to translate with acceptable quality\". This is false. 7 BLEU score is **A LOT** of difference. If you look at the generation result, they are qualitatively very different.\n    - ix) How exactly did you use the memory augmentation for a causal/decoder Transformer for the language modeling task in Table 3? This is not explained in detail in the paper, but how did you handle the potential information leakage problem (i.e., we don't want future tokens to flow back to the past)? If you simply do it in the same way as in encoders of MT by appending it to the sequence, then you can either only *read* from the memory without writing (if you append to the start of the sequence), or only *write* to it without reading (if you append to the end of the sequence). I might be missing something here, and hopefully the authors can clarify.\n    - x) The fact that there's no consistent improvement in Table 5 of the memory-augmented Transformer on the base Transformer is very concerning to me. Apparently, sometimes BERT-base is better than many of the MemTransformers.\n\n4. One potentially interesting study that I think the authors can look into is the memory slot permutation. In a certain sense, the memory slots have the same representation embedding `[mem]`, and are mainly differentiated (among themselves) by their respective position embedding. However, it also seems from Figure 2 that when writing to memory, the original order of the sequence tokens is not necessarily preserved. So what if, in the intermediate layers of the MemTransformer, you permute the order of the memory slots? Do you expect that to affect the performance of the MemTransformer?\n\n5. The other issue I found about the paper is that, although this is an empirical paper, there are still a bunch of not well-supported claims (i.e., handwavy claims) which the authors should have looked into. For example:\n    - i) In section 3.1: \"This can be due to the more complex architecture of the MemBottleneck that has twice more layers in the encoder part\". But there's no formal investigation of this in the paper... for example, does an 8-layer MemTransformer (which has \"twice as many layers as the 4-layer MemBottleneck) also perform badly?\n    - i) In the caption of Figure 2: \"Activity in the left top corner that involves first four tokens might indicate fusion of neighbour vectors by pairwise summation of `[mem]` tokens.\" This sounds like a posterior explanation... and it's not clear to me what this \"fusion\" is doing and why it's doing this. How is this helping MemTransformer? Does it occur in every head in every layer? Does it happen even if you only use memory size 5? These questions should be answered if you would like to claim this phenomenon.\n    - ii) In the \"memory to memory attention\" paragraph: \"If the diagonal attention is sharp, then corresponding memory vectors are added to themselves, so their content is amplified and became more error-tolerant.\" Do you have support for the error tolerance claim? Is it still \"amplified\" even if there is a LayerNorm immediately afterward, which would shrink things back anyway (i.e., $\\text{LN}(x)$ and $\\text{LN}(x+x)$ should be the same)?\n\n-------------------------------------\n\nSome minor points that didn't impact the score:\n\n1. In the intro section, 'results in \"blurring\"' ---> it's not clear to me what this means. Please clarify.\n2. In the very beginning of section 3, the authors used $N=4$ to refer to number of layers. But $N$ has been used to refer to sequence length before (e.g., when saying $O(N)$). \n\n- [1] https://www.ijcai.org/Proceedings/2020/0534.pdf\n- [2] https://arxiv.org/pdf/2010.07638v1.pdf\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "7b_VWzn-bf", "reviewer": "AnonReviewer3", "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 brought up an interesting limitation of a Transformer network, that information about the context is stored mostly in the same element-wise representation, which might limit the processing of properties related to global information. This work proposed adding memory tokens to store non-local representations and creating memory bottleneck for the global information. The evaluation shows positive results adding memory to a Transformer network.\n\n#### Weakness:\nThe idea is not entirely new; Memory augmented networks have been proposed in previous work and the paper does not differentiate the method from related work very well. \n\n-The paper is not clear about how to obtain the memory input token. According to the paper, the memory input token seems like a static vector; however, a memory augmented network can retrieve a dynamic vector for global information, with the ability to read and write to the memory. \n\n-The naive Mem Transformer Layer is counter intuitive, as there is little differentiation between the memory layer and the sequence layers. How the memory layer contains additional global information is not explained. \n-Figure 1c is confusing, are the arrows crossed or not?\n\n-The formulation of memory representation and sequence representation are exactly the same for the MemCtrl Transformer, which again is very counter intuitive. It is intuitive that we should augment the sequence transformer layer with additional global information, however, we might not want the other way. The memory layers should be updated at coarser granularity. \n\n-The paper is poorly written with many errors in the paper. Please proof read the paper before submission. \n\n#### Detailed feedbacks:\nThe reviewer finds this paper hard to read as there is not a flow of story in the paper. The layout of the paper and the design of experiments seems very arbitrary. \n\n-The motivation for a memory network like proposed in the paper is not clearly stated. How the memory network can capture long-term dependencies and information is not clearly explained. The design of variants of Memory Transformers seems very arbitrary. \n\n-The paper also lacks a more detailed comparison with related work, like Transformer-XL, Star-Transformer, Longformer, ETC, etc.\n\n-\"global\"->''global'' (make sure it looks correct in latex). \n\n-\".Surprisingly\"->\". Surprisingly\"\n\n-Please improve the quality of writing and ask native speakers to proof read the paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "3glSHz0IRDb", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "The authors propose augmenting transformer architectures by adding trainable memory via introducing special memory tokens to the input sequence. They also explore different architecture extensions (memory bottleneck, dedicated layer for memory updates). The authors’ main hypothesis is that adding memory to Transformer-based architectures should result in better model performance. They present BLEU scores on WMTF-14 DE-EN translation and GLUE benchmark results to support that.\n\nPros:\n+ Good qualitative analysis of attention patterns that emerge when memory tokens are added\n+ A practically viable memory extension training scheme that can be used to fine-tune models with added memory tokens\n\nCons:\n- The main claim of the paper is that presence of memory tokens positively correlates with the model performance. At the same time, adding memory tokens increases the model’s capacity, making the comparisons with the baseline unfair. The more rigorous approach would be to compare to a baseline with proportionally increased dimensions as having memory tokens may be akin to having additional space to store sequence token representations. Some of the attention patterns (write-to-memory -> store -> read-from-memory) indicate that this may be the reason for the improved model scores.\n- The use of global tokens (like [CLS] token in Longformer as mentioned in the paper) is a well-known technique that undermines the paper’s novelty.\n- The authors claim that models’ quality positively correlates with the memory size. However, this statement contradicts the findings in Table 1 and Table 5. The authors briefly discuss the lack of trainability in section 3.1. Still, the authors do not specify conditions under which adding more memory will benefit the model quality.\n \nGeneral comments/questions:\n* Memory lesions experiments (table 2). Memory size change during inference seems to lead to drastic drops in model quality. At the same time, the models are not trained to perform in the absence of memory. What if the models are trained on different memory sizes from the beginning with memory size sampled per-batch?\n* Memory extension experiments (table 3). Does the model forget how to work with smaller memory sizes during the memory extension process? E.g., what will memory lesions results look like for the memory extension model?\n* Are any of the models for the GLUE benchmark trained using a memory extension training scheme?\n* It’s not clear if all mem tokens use the same embeddings or if they use separate embeddings (e.g., [mem1], [mem2], etc)\n* MemBottleneck results lack comparison with other extensions that lower the computational cost (O(N), O(N log N) extensions, etc)\n* Relative positional embeddings seem to work better with memory tokens, are there any more experiments with relative positional embeddings except for the ones in table 5?\n* Not all results in Table 1 are symmetrically available for both small and base models. This table lacks the results for MemCtrl Transformer for small models and lacks the results for MemBottleneck Transformer for base models which makes it harder to assess the viability of those architecture extensions. \n* No results for MemBottleneck and MemCtrl in table 5. At the moment, both MemBottleneck and MemCtrl look weak; having more results for those extensions may help.\n* The authors can add FLOPS/params count to tables 1 and 5 to make a comparison of the models' more convenient.\n \nOn rating:\n \nAlthough extending transformer architectures with memory tokens looks like a practically viable way to improve models’ quality, the experiments are not rigorous enough to confirm this is due to the memory mechanism and not because of the increased model capacity, which can be achieved by a simple model scaling. The well-documented use of global tokens in literature also diminishes the novelty of this paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "2TdaDWGUP77", "reviewer": "AnonReviewer4", "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": "This paper proposes to augment transformer architectures with memory components. The high-level idea is to use multiple special “mem” tokens as additional inputs. Depending how the “mem” tokens’ representations interact with the true input sequence, three variants are studied: (a) in MemTransformer [“mem”, input] attends to [“mem”, input]; (b) MemCtrl is similar, but uses a separately parameterized module to calculate the “mem” representations; (c) in MemBottleneck, “mem” attends to [“mem”, input], and [input] attends only to “mem.” Experiments with machine translation, language modeling, and fine-tuning on the GLUE benchmark are conducted.\n\nOverall the presentation of the paper could be significantly improved. For example, it lays out the technical sections in a way that assumes the readers have hands-on expertise with transformer architectures. Further, the motivation is not clear to me. It seems from the experiments that little practical gains can be expected, since MemTransformer never outperforms the transformer baseline in terms of accuracy or efficiency. This is probably fine if the paper studies a set of clear research problems with well-designed experiments and answers some interesting questions. But it is not the case. Last but not least, the experimental setting is flawed: e.g., the MT experiment cuts training at a suspiciously early stage, presumably contributing to the bad performance. In sum I do not think the paper is ready for ICLR.\n\nPros:\n- Studying the transformers through the memory perspective is interesting.\n\nCon:\n- Writing can be significantly improved.\n- Motivation is not clear: the proposed approach doesn’t seem to bring any interesting practical gain, nor does it answer any interesting research question.\n- Experimental design is flawed, making the conclusions less convincing.\n\nDetailed comments:\n- I’m not against discussing related works in the introduction. But the current version is probably trying to scramble too much stuff into intro, making it tedious to read.\n- I suggest walking through the transformer architecture in a more self-consistent way. I would be very surprised if one without much hands-on experience of transformers can understand what’s going on in section 2.\n- The baselines’ performance in the MT experiments is far worse than what we usually see. I’m guessing this is because the training stops too early. Can the authors explain why they choose to do so? If limited computation is the concern, I recommend the authors to work with smaller datasets instead but use a more convincing settings. Same for other experiments.\n- The experiments do not show any practical gain from MemTransformer. The paper talked about its linear complexity in input sequence length. This paper would be stronger if it can present convincing experiments and show that MemTransformer can have a better efficiency in practice.\n- Section 3.2 would be much more clear and interesting if it opens with the research questions it’s trying to answer.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies three kinds of memory-augmented Transformers, focusing on one (the MemTransformer, which adds [MEM] tokens to a document.)  This is a nice clean extension of Transformers and a topic well worth investigating.  Unfortunately, the experimental results were considered unconvincing:

 - The baselines were relatively weak
 - The experimental setting was unusual (eg only 10 epochs)
 - The experiments did not show consistent improvement

Overall the paper was considered below acceptable quality for ICLR.

Decision: Reject