Recurrent Memory Transformer
Authors: Aydar Bulatov
Yuri Kuratov
Mikhail Burtsev
Venue: NeurIPS 2022 Accept (2022)
Keywords: memory augmented models
transformers
self-attention
recurrence

Abstract:   Transformer-based models show their effectiveness across multiple domains and tasks. The self-attention allows to combine information from all sequence elements into context-aware representations. However, global and local information has to be stored mostly in the same element-wise representations. Moreover, the length of an input sequence is limited by quadratic computational complexity of self-attention.
  In this work, we propose and study a memory-augmented segment-level recurrent Transformer (RMT). Memory allows to store and process local and global information as well as to pass information between segments of the long sequence with the help of recurrence.
  We implement a memory mechanism with no changes to Transformer model by adding special memory tokens to the input or output sequence. Then the model is trained to control both memory operations and sequence representations processing.
  Results of experiments show that RMT performs on par with the Transformer-XL on language modeling for smaller memory sizes and outperforms it for tasks that require longer sequence processing. We show that adding memory tokens to Tr-XL is able to improve its performance. This makes Recurrent Memory Transformer a promising architecture for applications that require learning of long-term dependencies and general purpose in memory processing, such as algorithmic tasks and reasoning.

Recurrent Memory Transformer
Aydar Bulatov' bulatov. wwwps/asssphhshcce..e.m edu
Yuri Kuratov 1.2 yurii www.kooo@pysteeh..eeee
Mikhail S. Burtsev 1,2
burtcev.. pertcev.ms@miptt mw@mimt. ru
Neural Networks and Deep Learning Lab, Moscow Institute of Physics and Technology, Dolgoprudny, Russia
2 AAIRI, Moscow, Russia
Abstract
Transformer- models show their effectiveness across multiple domains and tasks. The self-attention allows to combine information from all sequence elements into context-aware representations. However, global and local information has be stored mostly in the same element-wise representations. Moreover, the length an input sequence is limited by quadratic computational complexity of self- attention. In this work, we propose and study evemory-augmente recurrent Transformer (RMT). Memory allows to store and process local and global information as well as to pass information between segments of the long sequence with the help of recurrence. We implement a memory mechanism with no changes to Transformer model by adding special memory tokens to the input output quence. Then the model is trained to control both memory operations and sequence representations processing. Results of experiments show that RMT performs par with the Transformer-XL on language modeling for smaller memory sizes and outperforms it for tasks that require longer sequence processing. We show that adding memory tokens to Tr-XL is able to improve its performance. This makes Recurrent Memory Transformer promising architecture for applications that require learning of long-term dependencies and general purpose in memory processing, such as algorithmic tasks and reasoning.
segment-level
1 Introduction
Transformers Vaswani al. 2017) have been widely adopted across multiple domains and tasks (Radford
ctt al. 2018: Dong et al. 2018: Devlin et al. 20199 Dosovitskiy et al. 202. Ramesh et al. 2222 Jacgle ct al. 2021). The key component of Transformer layer self-attention Self-attention allows up- date each sequence element representation with in- formation from all other elements in the sequence. As result, rich contextual representation for every element is generated at the end of encoding. This way, global sequence-level and local information are stored in single representation. However, this mix- ing of two types of information in a single represen- tation has limitations. Distributed storage of global features across all sequence elements results in global features "blurring" and makes it harder to access them. Another well-known deficiency of Transformers poor scaling of self-attention with
Transfornet Layers
Tranaforner Layera
segrent
segnent
Figure Recurrent Memory Transformer. Memory is added as tokens to the input sequence and memory output is passed to the next segment. During training gradients flow from the current segment through memory to previous segment.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

input sequence length that hurts its applications to long imputs (Child et al., 2019 Guo et al.
2019
Dai et al.] 2019 Beltagy et al. 2020 Ainslie et al., 2020 Zaheer et al. 2020; Wang et al. Choromanski et al. (0020).
2020
segment-level recurrent Transformer named Recurrent Memory Transformer (RMT). RMT uses a memory mechanism based on special memory tokens (Burt: sev al. 2020) added to the input sequence. Memory tokens provide additional reserved capacity to the model that could be used to process information which is not directly representing any element in the input sequence. To process long sequences, we split them into segments and pass memory states from a previous to current segment. This memory passing makes the model recurrent and removes the input sequence length limitations. RMT model can theoretically work with infinite lengths but, in practice, it is limited by memory capacity and the efficiency of memory access/update operations. Our implementation of both memory and recurrence in RMT requires no changes to the Transformer model because modifications are made only to the input and output sequences of the model.
Our work introduces a memory-augmenni
We tested RMT on the tasks that require global information about the whole input sequence to be solved. We use copy, reverse, and associative retrieval tasks in the setting where the input sequence split into segments. RMT and Transformer-XL. perfectly solve these tasks, but exceeding some value of sequence length, RMT starts to outperform Transformer-Xe XL. Also, we experimentally show that the proposed Recurrent Memory Transformer requires less memory size to perform closely to Transformer-X XL on language modeling tasks. RMT code and experiments are available"
Contributions
In this study we augment Transformer with token based memory storage and segment-level recurrence.
We experimentally evaluate proposed architecture as well as vanilla Transformer and Transformer- XL on memory-intensive tasks such as copy, reverse, associative retrieval, and language modeling. We show that RMT outperforms Transformer-Xi XL for sequence processing tasks and on par with Transformer-a XL on language modeling but requires less memory.
3. We show that Tr-XL cache could be combined with RMT leading to better performance on language modeling.
4. We analysed how the Transformer model learns to use memory. Specific interpretable memory read-write patterns of attention are shown.
2
Related work
In our study we add memory to general purpose attention based neural architecture. Memory is a recurrent topic in neural networks research. It had started from the early works (McCulloch and
Pitts 1943] Stephen 1956) and significantly progressed in 90's with introduction of Backprepagation Through Time learning a algorithm [Werbos 1990) and Long-Short Term Memory (LSTM) (Hochreiter) neural networks (MANNs) usually rely on some kind of recurrent external-memory which is separate from the model's pa- rameters. Neural Turing Machines (NTMs) (Graves et al. 2014) and Memory Networks (Weston et al, 2014) are equipped with a storage for vector representations that can be accessed with an attention mechanism. Memory Networks (Weston al. 20144; Sukhbaatar et al. 2015) were designed to enable reasoning by sequential attention over to the content of memory. NTMs followed by Differentiable Neural Computer (DNC) (Graves et al., 2016) and Sparse DNC (Rae et al. 2016) are implemented as recurrent neural networks able to write to memory storage over time. All these models are differentiable and can be trained via backpropagation through time (BPTT). Parallel line of research extends recurrent neural networks such as LSTM with data structures like stacks, lists, or queues (Joulin and Mikolov 2015; Grefenstette et al. 2015). MANN architectures with a more advanced addressing mechanisms such as address-content separation and multi-step addressing were proposed in (Gulcehre et al. 2016 2017 Meng and Rumshisky 2018). The Global Context Layer model (Meng and Rumshisky 2018) uses the idea of address-content separation to solve the difficulty of training content- based addressing in the canonical NTM.
and Schmidhuber 1997) neural architecture. Today memory-augmentee
htttps://gtthi wwtp/oooda////L.---.... RMT The code, results of the raw experiments and hyperparameters are provided in the supplementary materials and GitHub.

---

The recent rise of Transformer models also resulted in introduction of a number of new memory architectures. Transformer-Xl (Dai al. 2019) introduces a segment-level recurrence at the level of hidden representations. These representations of sequence are computed and stored in the cache to be reused as an extended context for the next segment. Compressive Transformer (Rae et al.) 2019) adds the second layer of memory to Transformer-XX XL. This memory compresses and stores information from the cache. oo-former (Martins et al. 2021) utilizes continuous- space attention and represents input sequence a continuous signal to make long-term memory unbounded. Memory Layers (Lample et al. 2019) model has product key memory layer instead of a feed-forward layer within Transformer block to increase model capacity.
In many variations of Transformer different sorts of global representations are added. Among them are Star- Star-Tranformer (Guo et al. (019), Longformer (Beltagy et al. 20200, GMAT (Gupta and Berant 2020), Extended Transformer Construction (ETC) (Ainslie al., 2020) and Big Bird (Zaheer et al., 2020). All these architectures re-design self-attention mechanism to reduce it computational com- plexity with and ensure input coverage with the help of global representations. Memory Transformer (Burtsev et al. 2020) keeps Transformer model intact and adds memory by extending input sequence with special memory tokens. Perceiver IO (Jaegle et al. (021) maps an entire arbitrary input to the fixed number of latent representations. Transformer layers do further processing over latent memory representations only.
recurrence in Transformers is actively explored in number of studies. Transformer- XL, Compressive Transformer keep previous states and re-use them in subsequent segments. Ernie- Doc (Ding et al.) 2021) improves processing by using same-layer recurrence instead of attending to previous layer outputs of a precedent segment. Memformer (Wu et al.) 2020) introduces a dedicated memory module to keep previous hidden states in summarized representations. Memformer uses two special layers added to the Transformer model. Memory cross-attention layer reads from memory and memory slot attention layer updates it. MART (Lei et al., 2020) has similar approach as Memformer but uses memory update rules analogous to LSTM (Hochreiter and Schmidhuber, 1997) 2020) goes further with full, and not
Segment-level
and GRU (Cho et al., 2014). FeedBack Transformer (Fan et al. segment-level. recurrence. FeedBack Memory merges past hidden representations from all layers into a single vector and makes it accessible to the computations any layer. The disadvantage of full recurrence is that it is less parallclizable. FeedBack Memory requires every sequence element to be processed sequentially. In segment-level recurrent models, all elements of segment processed by Transformer layers in parallel. Only segments are processed sequentially. Staircase Transformer (Ju et al.) 2021) combines segment-level recurrence and depth recurrence. Staircase models the output for previous segments and pass them as input for the next segment. Our Recurrent Memory Transformer is based on special memory tokens similar Memory Transformer, segment-level recurrence as in Transformer-Xi and depth-recurrent mechanism for memory processing similar to Staircase.
3 Recurrent Memory Transformer
Transformer-X XL (Dai et al. 2019) extends Transformer model with state re-use cache mechanism for segment-leve recurrence and relative position encoding. Input sequence is split on segments processed sequentially. Hidden states computed for the previous segment M are cached for each transformer layer n. The input of the layer consists of the last m states from the eached memory and output of previous Transformer layer for the current segment
HP- [SG(M"m)) H;
here, SG stands for moop-gradient, o denotes concatenation. Cached states allow to increase effective context size of Transformer model and save on compute operations.
Then, Am-1 goes to Transformer layer TL to produce layer outputs for segment
H TI(Q", K" V.", Q = W W.
In Transformer-Xe XL, self-attention layers are modified to use relative position encodings to improve generalization to longer attention lengths. The overall architecture shown in the Figure
Memory augmented Transformers such as GMAT, ETC, Memory Transformer (Gupta and Berant 2020; Ainslic et al. 2020: Burtsev al. 2020) proposed to use special global tokens storage
3

---

Fransformer-Xi
Recurzent Mamozy Transformer
Forne layaz
Actuntion
Rffeclive
Forwaed pesa
Stop Gradient
Figure 2: Comparison of Recurrent Memory Transformer (RMT) and Transformer-Xi XL architectures. Recurrent Memory Transformer augments Transformer with global memory tokens and passes them to allow a segment-level recurrence. Special read/write memory tokens are added to the input sequence. Multiple memory tokens can be used in each read/write block. Updated representations of write memory are passed to the next segment. During training, RMT uses BPTT to propagate gradient to previous segments through memory tokens representation. Effective context length recurrence with memory which is case the cache Transformer-XT.
limited by the depth of network
for representations. Usually, memory tokens are added to the beginning of the input sequence. However, in decoder-only architectures the causal attention mask makes impossible for memory tokens at the start of the sequence to collect information from the subsequent tokens. On the other hand, if memory tokens are placed at the end of the sequence then preceding tokens unable to access their representations. To solve this problem we add recurrence to the sequence processing. Representations of memory tokens placed at the end of the segment are used as an input memory representations at the start well as the end of the next segment.
In the Recurrent Memory Transformer input augmented with special [mem] tokens, processed in standard way along with the sequence of tokens. Each memory token is real-valued vector. m memory tokens are added at the beginning of the segment tokens representations H and the same m tokens are added at the end:
H = [Hme H H. H. = Transformer(H!), [Hread o H Hurite] == H.N
here N number of Transformer layers.
The starting group of memory tokens functions as a read memory that allows sequence tokens to attend to memory states produced at the previous segment. The ending group works write memory that can attend to all current segment tokens and update representation stored in the memory. As result, H Hurste contains updated memory tokens for the segment т.
Segments of the input sequence are processed sequentially. To enable recurrent connection between segments, we pass outputs of the memory tokens from the current segment to the input of the next segment:
H H.rrite ft. [H#4 H H
Both memory and recurrence the RMT are based only on global memory tokens. allows to keep the backbone Transformer unchanged and make RMT memory augmentation compatible with any model from the Transformer family. Memory tokens operate only on the input and output of the model. In this study implement RMT on top of the original Transformer-Xi code. Both architectures are shown in Figure 2
Recurrence in the RMT is different compared to the Transformer-XI XI because the former stores only m memory vectors per segment. On the other hand, the Transformer-N XL stores m N vectors per segment. Also, in the RMT model memory representations from the previous segment are processed by Transformer layers together with the current segment tokens. This makes memory part of RMT effectively deeper in number of applied Transformer layers N. Additionally, we allow all memory tokens in the read/write block to access all other tokens in the same block. The causal attention mask is applied only to tokens of the input sequence (Figure (Figure((d)).
We train the RMT with Backpropagation Through Time (BPTT). During backward pass, unlike in Transformer-XX XLL memory gradients are not stopped between segments. The number of previous segments to backpropagate is hyperparameter of a training procedure. We vary BPTT unroll in our experiments from 0 to previous segments. Increasing this parameter is computationally expensive

---

and requires lot of GPU RAM. However, techniques such as gradient checkpointing could be used alleviate this problem.
4
Experiments
We designed our experiments to evaluate the ability of Recurrent Memory Transformers to preserve long-term dependencies across multiple imput segments. The first set of experiments includes copy, reverse, associative retrieval, and quadratic equations tasks. The second one addresses language modeling task for word-level WikiText- 103 (Merity et al. 2017) and for character-level on enwik8 (Mahoney] 2006). We compare Recurrent Memory Transformer with Transformer and Transformer-XI/ models.
Our RMT implementation is based on Transformer-XI. repository D The full set of hyperparameters
available in our repository well in supplementary materials. Language modeling experiments follow the same model and training hyperparameters as Transformer-XL. WikiText- 103 experiments use 16-layer Transformers (10 heads, 410 hidden size, 2100 intermediate FF), enwik8 - 12 layer Transformers (8 heads, 512 hidden size, 2048 intermediate FF). We used Adam optimizer] Kingma and Ba (2015) with linear schedule learning rate starting from 0.00025 for 200,000 steps for WikiText- 103 and 400,000 steps for enwik8. We refer to Transformer XL with memory size equal to zero as a Baseline. With this experimental setup were able to reproduce results for the Transformer-Xi model close to the original paper.
Algorithmic Tasks. We evaluate RMT on algorithmic tasks that require information about the whole input sequence to be solved successfully. In a recurrent setting, the model has to keep information about all previous segments to make predictions.
In the Copy task, an input sequence should be replicated twice after special sart-to-generate token. In the Reverse task, an input sequence should be generated in reverse order. Input for the Associative Retrieval task consists of N key-value pairs. Then one key is randomly selected, and the task is to produce an appropriate value for the selected key. Another task is to solve quadratic equations. One example consists of equation, its solution with discriminant, and an answer. The task to generate solution and answer, while only answer quality is evaluated.
For all tasks, input and output sequences are split into segments and processed by models sequentially. Datasets for algorithmic tasks were randomly pre-generated, the same data was used in all experiments, and character- tokenization was used. Because Transformer-Xl and RMT are decoder-only Transformer models, we don't 't compute loss over the input sequence before the start-to-generate token. The loss is computed over target sequence segments only (see Appendix |A.] for details).
Language Modeling and NLP. We use two standard benchmarks for language modeling: WikiText- 103 and enwik8. 103 (Merity et al. (017)) is used for word-level language modeling and contains 103M words from English Wikipedia articles. Enwik8 (Mahoney) (006) is used for character-level and consists of 10 first bytes of XMI. text dump of the English Wikipedia. Vocabulary contains 267735 words and 204 characters for Wikitext- 103 and enwik8 tokenizers accordingly.
We compare Recurrent Memory Transformer with decoder-only Transformer and Transformer-Xi as baselines. Model size and training parameters are selected to match Transformer-Xi paper. For Wikitext-103 an input context length was set to 150 tokens, and for enwik8 it was set to 512 characters. Another set of experiments inspected how RMT handles long-term dependencies and recurrence. We increased the number of segments and recurrent steps by making segments smaller (50 tokens for WikiText- 103, 128 characters for enwik8). The increased number of recurrent steps makes language modeling tasks harder for RMT because information has be stored in the same amount of memory for more steps.
As testbed for the real-life application scenario we select popular long-text classification benchmark Hyperpartisan news (Kiesel et al.] 2019). Instead of pre-training RMT from scratch we add recurrent memory mechanism to the most widely adopted models from HuggingFacc Transformers (Wolf et al.) 2020). Specifically, we augment 500 input tokens of already pretrained BERT-base, RoBERTa-base. DeBERTa-base and base with the recurrent memory of size 10 and fine-tune the target task.
5

---

5 Results
Baseline, Transformer-Xe XL (Tr-XL) and RMT perform perfectly in the single segment setting on copy and reverse tasks (Figure 3). In this case, the models do not need recurrence because the whole sequence available. When the number of segments is larger than one, non-recurrent baseline struggles to solve tasks, but both memory models demonstrate ability to retain required information from the previous segments in memory.
Hasslo
hamtourmeoL
Mermury Tartdleeree
HMI 1P11-0
Cony
Reverse
Associalive retrieval
1.0
1.0
1.0
Accurecy
0.8
0.8
0.6
0.6
0.4
0.2
0/2
0.2
3
Number of sogments [2]
Number of sogments (b)
Number of segment (c)
Figure 3: RMT outperforms Transformer-XL on Copy and Reverse tasks as a number of segments increases. Panels show test set per-character accuracy copy, reverse, and associative retrieval tasks (from left right). Memorylcache size equals the length of segment for both models. RMT does not pass gradients between segments in this experiment. MT results are the same as for the Baseline. Source/target sequence lengths for copy, reverse, and associative retrieval tasks: 24/48, 24/24, 10/1.
On Copy and Reverse tasks as a number of segments increases, RMT starts to outperform Transformer- XL with memory sizes less than the number of all previous tokens. With the number of segments up to 6 mean accuracy of Transformer-X XL drops by up to 0.2 points, and with 9 segments plunges close to the baseline without memory. Associative Retrieval results are similar with the number of segments up to 4. RMT manages to solve the task with Transformer-XL. closely behind. However, in the setting with segments, RMT performance slightly decreases and Transformer-Xi average accuracy rises higher.
We analyze how number of segments, sequence length, a length of training context, and memory size affect models' performance on Copy task (Figure As we split a sequence into more segments it becomes more crucial to be able to pass information between segments. We split 360 tokens of source + target sequence into multiple segments. In Figure 4a we observe that Transformer-A performance starts to degrade and eventually falls to the baseline model performance as the number of segments increases. In contrast, RMT continues to solve the task perfectly. In a more extreme setting, when we keep memory size fixed, but increase the total length of sequence to copy Transformer-XL. fails shortly, while RMT starts to gradually degrade only after the length of 720 tokens (Figure (b))
On the Quadratic Equations task (Table I we have checked that it is possible to solve the task with the Transformer baseline and no segmentation used. The baseline in this case defines upper bound for this task. With multiple segments recurrency RMT solves the task perfectly, while Transformer-Xi XL finds the task challenging.
The results of experiments on word-level lan- guage modeling on WikiText- 103 are shown in Table In the first section with a segment length of 150, Tr-XL and RMT outperform the base- line and Memory Transformer (Mem'Tr) by large margin. It shows the significance of in- creased effective context length by Tr-XL cache or RMT memory for language modeling. RMT improves over MemTr memory mechanism with read/write blocks. The best RMT models with memory size 10 and 25 show similar performance Transformer-Xi with memory size equal to 75. RMT learns to use smaller memory more effectively than Transformer XL. Additionally, the smaller memory size of RMT leads to reducing required GPU memory for running the model.
Table 1: Quadratic equations task. Sequence of 180 tokens consists of quadratic equation, a solution, and an answer. It is split into number of segments with an answer in the last segment. Accuracy equals 1.0 if the full answer is predicted correctly.
MODEL
MEMORY
SEGMENTS
ACC
BASELINE
1
0.99 1: 0.01
TRANSPORMER-XI
30
6
0.93 1 0.02
RMT
0.99 0.002
To force models to process longer recurrent dependencies the size of a segment is set to 50, so the number of recurrent steps increases. RMT with memory size I shows similar results to Transformer-

---

1.0
1.0
Accuracy
Accuracy
0.8
Transfomer-AL
RMT with single memory vector could be trained to achieve lower perplexity as Transformer-Xi with memory size 10. This means that RMT can learn to compress in- formation from the previous observations bet- ter. Another observation is that RMT with RIT BPTT-6
0.4
0.4
2
6 Number of segments
200
400
600
1000
Source
+ Target sequence length
(b)
Figure 4: RMT scales better with number of segments and sequence size, (a) RMT is able to solve copy lask perfectly up to 9 segments for fixed sequence length of 360, while Tr-XL fails. (b) RMT learns to use memory of the same fixed size (60 tokens) more effectively than TR-XL sequence length to copy inereases segment size is 120 for the both models).
XL with memory size 10. It is worth noting that Transformer-Xi XL memory consists of hidden representations from all layers (in this case, is 10 16 vectors) when RMT memory only memory_size vectors. Transformer-Xi with memory size 50 and RMT with memory size 5 show similar perplexity values (see Appendix A.5).
RMT could be combined with Tr-XL cache. In this case Tr-XL cache could be seen as short-term memory keeping the nearest context and RMT memory long-term memory. Such combination leads to the best results on WikiText-105 improving over Tr-XL.
Table 2:
Language modeling on WikiText-103. Av- erage perplexity for the best performed variations of RMT models reported (see full results in appendix(A.5] Underlined values show Tr-XL and RMT models with close results. RMT models with smaller memory sizes achieve similar scores to Tr-XL models with larger mem- ory. Combination of cache with recurrent memory (Tr- XL RMT) shows the best performance.
On enwik8 RMT models with memory size 5 and Transformer-Xl. with memory size 40 show similar results. Confirming that RMT learns to use smaller amounts of memory representation more effectively. All results for enwik8 dataset are shown in Appendix A.4
Recurrent Memory Transformer learns to make predictions depending on #BPTT_unrolL. previous segments +1 current segment. Transformer-XLL does not use BPTT and relies only nemory_size cached states and cur- rent segment making in total: memory _size segment taanentt_enghhhhhh tokens. In Figure 5a compare RMT and Tr-LL according to the de- seribed value of visible context training time.
MODEL
MEMORY
SEGMENT LEN
PPL
TR-XI. (PAPER)
150
150
24.0
BASELINE
150
29.95
0.15
MENTR
10
150
29.63 :: 0.06
(OURS)
TR-XL
150
25.57
0.02
TR TR-XL
75
150
24.68
RMT BPTT-3
10
150
25.04
0.07
RMT BPTT-2
TR. TR-XL. RMT
75+5
150
24.47
0.05
TR-XL RMT
150+10
150
23.99
0.09
BASELINE
50
39.05 + 0.01
TR-XL
TR TR-XL
50
50
26.54 + 0.01
TR-XI.
27.57
0.09
TR-XL
10
50
28.98
1 0.11
memory sizes 10 and 25 performs only a
RMT BPTT-1
50
28.71
0.03
bit weaker compared Transformer-XL even when Transformer-XL has access to more non- compressed states (50, 100, 200) from previous segments. In general, training RMT with unrolling gradients in earlier segments drastically improves scores thus showing the importance of BPTT training but, observe instabilities and out-of-memory issues during RMT training for a larger memory sizes with deeper BPTT unrolls.
RMT 3
10
50
26.37
RMT wins lot when only memory token is added but then the effect from increasing memory size from 5 to 50 fades (Figure 55). Still, RMT with memory size have performance on par with Transformer-X XL with cache 50, confirming that RMT leams to store more compact representations. The results suggest that there some optimal memory size for RMT to solve the task, and further increase does not add much.
Proposed recurrent memory mechanism affects only input and gradient flows of the augmented core model. This might be important advantage because the memory can be added to already pretrained
7

---

40.0
32
Tansformer-AL
Baseline
RMT BPITO
37.5
TransformerX
31
BPTTC
RMT BPTT-0
RMT BPTT-
35.0
30
RMT BPTT
BPTT-1
BPTT-4
32.5
RMT BPTT-2
29
1est
Test
30.0
26
27.5
27
50
25.0
26
100
200
25
50
75
100
125
150
175
200
225
250
Visible context while training
Memory size
(b)
Figure 5: Deeper BPPT unrolling improves RMT scores Wikifexx-1033 (a) Visible context training
time be increased by deeper BPTT unrolls for RMT enlarging cache for Tr-XL. Larger visible context leads to lower perplexity both models (marker size corresponds to memory size). (h) Recurrence improves performance of RMT compared to Tr-XL for the same memory sizes.
model. Evaluation results for four memory augmented language models fine tuned for long text classification are presented in the Table Incorporation of 10 memory tokens in the input sequence of 512 allows to encode longer stretches of a text up to 2000 tokens and significantly improve metrics for the majority of models. Moreover, combination of recurrent memory with RoBERTa-bare results in state of the art performance for the Hyperpartisan news classification task (Kiesel et al.) 2019). Interestingly, many competing models have input size of 4096 that is at least twice longer compared to RMT extended counterparts but still lag behind.
To get an understanding of memory operations, learned by RMT for algorithmic tasks we visualise atten- tion maps for copy and re- verse tasks (Figure (6). In each RMT attention map quence tokens are preceded by read memory, located the top left corner, and fol- lowed by write memory the bottom right. Diago- nal the central part of the G(a) (top) shows classic attention of token sequence to itself, but the bottom diagonal represents the operation of writing of sequence tokens to memory in straight order. When completing reverse (fig 6[a) bottom) the model learns to write the sequence to the memory in the reversed order, which is in line with common sense.
Table 3: Hyperpartisan news detection. Models starting with RMT are taken from HuggingFace Transformers and augmented with 10 memory tokens and recurrence before fine-tuning. Train/valid/test split as in (Beltagy et al. 2020) and metric is Fl.
MODEL JINPUT SIZE]
3
BIG BIRD [4096] ZAHEER
2020
92.20
LONGFORMER [4096] BEETTGG FT
2020
94,80
GRAPB-ROBERTA T012
96.15
ERNIE-DOS DOC LARGE [640] (DING Er AL.
2021
96.60
FRNIF-SPPRRS [4096] (Lau
92.81
RMT DEXT-BASE-CASI [512]
91.60
94.12
93.06
94.34
RMI ROBERTA-BASE [512]
94.87
97.20
96.72
98.11
RMT DEBERTA-V3-BASH [512]
94.17
96. 78
94,80
94.80
15-BASE [512]
94.99
When it comes reproducing the target sequence, the model accesses memory (fig() b)) and writes to the output sequence. Another operation (figlGc)) is rewriting from read memory to write memory. It is commonly used by RMT in settings with larger number of segments to keep information about recent segments longer.
Transformer-Xi
mechanism of accessing memory (f) does not allow straightforward writing to memory without changing sequence token representations. Sequential reading from cache is represented by diagonals harms model performance in tasks with larger number of segments. For reverse task with 4 segments Transformer- XL with limited memory size 6 (Appendix Figure D attempts to mix representations of tokens and read multiple symbols from one cached state in the next segments giving average accuracy of 0.8 the target task. Despite having the same memory size, RMT manages compress the whole segment in memory tokens (Appendix B
-X.
Transformer-XL. attention maps. Using token representations as storage
Figure [a)) and achieve mean accuracy 1.
Visualizations from Figure and Appendix B Figure provide evidence to support our hypotheses that TL- has to mix representations from previous current segments the same hidden states to pass information between segments. Also, visualizations show how memory tokens in RMT help

---

mitigate such kind of mixing. RMT ability of sequence compression to memory is illustrated in Appendix Figure 8
For copy with 6 segments RMT compresses and then reads the sequence of 12 tokens with just 6 memory tokens. For Transformer-X XI. decreasing memory size harms the accuracy score significantly with number of segments larger than 2.
FECESENE
update
seq
read
write to
Copy
Layer
Layer
resrite FXn
wrace 1903
HEad
188088800
Chunc
Copy
Reverse Layer
Lover
Layer
Reverse
Heeed
Heac
Tear
(c)
Figure 6: Selected attention map patterns of memory models. (color intensity corresponds attention score) RMT with segment length=24, memory size=24 (a) write to memory, (h) read from memory. (c) RMT, segment length=8, memory size=8, rewrite from read memory to write memory. (d) Transformer- XL, segment length=24, memory size=24 read from the previous hidden states.
Conclusions
In this paper we introduced Recurrent Memory Transformer simple recurrent memory augmentation of Transformer model. RMT is implemented by extension of an input sequence with special global memory tokens and segment-level recurrence. Importantly, our method allows learn more compact sequence representations and improve existing pretrained models without extensive additional com- pute, thus making practical machine learning applications more energy efficient and environmentally friendly.
In experiments we compared RMT with Transformer baseline and Transformer-Xl. which is well-known modification of Transformer for long sequences. RMT almost perfectly solves Copy, Reverse well quadratic equations tasks for sequences consisting of multiple segments outperforming Transformer-XX XL.. It also demonstrates quality for associative retrieval task on par with Transformer As expected, baseline Transformer fails solve these tasks for selti-segment settings.
RMT trained language model performs significantly ahead of Transformer baseline and shows quality metrics similar to Transformer-X but for up to 10 times smaller memory size. Experimental results demonstrate that for fixed memory size backpropagating gradients for segments improves performance of RMT. Proposed approach to memory augmentation is quite universal and might be easily applied to any pretrained transformer based model as demonstrated by achievement of state of the art results for long text classification task by fine tuning a combination of RoBERTa and RMT.
Analysis of attention maps suggests that better RMT performance be related more effective storage of input representations in dedicated memory tokens compared to mixing representations storage in Transformer-Xt XL. RMT could be combined with Pransoomer-XL cache and improve the performance of both models.
Overall, results of the study show that dedicated memory storage and recurrence provided by Recurrent Memory Transformer make it a promising architecture for applications that require learning of long-term dependencies and general purpose in-memory processing, such as algorithmic tasks and reasoning. Furthermore, we believe that RMT could open the way for adding memory and recurrence to other models in the Transformer family.

---

Acknowledgments and Disclosure of Funding
This work was supported by a grant for research centers in the field artificial intelligence, provided by the Analytical Center for the Government of the Russian Federation in accordance with the subsidy agreement (agreement identifier (100000333211550000
and the agreement with the Moscow Institute of Physics and Technology dated November 1, 2021 No. 70-2021-00138.
References
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Philip Pham, Anirudh Ravula, and Sumit Sanghai. Ete: Encoding long and structured data in transformers, 2020.
Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, Catalin Ionescu. Using fast weights attend to the recent past. In D. Lee, M. Sugiyama, U. Luxburg. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URI, URI,https: h/pppe/ee//ccceingg/// Papas---pppr... happrr.dff
neurips. heurips.cc/paper/ /016 (ttp:///4444444422226e6666666e60
Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/55 2020.
Mikhail S Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V Sapunov, Memory transformer. arXiv preprint arXiv:2006. 1527, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers, 2019.
Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decode approaches. In Proceedings of SSS7-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation, pages 103-111, Doha, Qatar, October 2014. Association for Computational Linguistics. doi: 10.3115/ 1/111/4012144
URL https:://aclaatbooogg wwwpa/ooooooooooooeooooooe wwwp:///////000//
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Zihang Dai, Zhilin Yang. Yiming Yang. Jaime Carbonell, Quoe V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context, 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lec, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume / (Long and Short Papers), pages 4171-4186. 2019. URL htttpp/://acleeb..oggaaatttooooo haapeooooogyppppe///////// 1423/
SiYu Ding, Junyuan Shang. Shuohuan Wang. Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. ERNIE-DoG retrospective long- document modeling transformer. Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language
Processing (Volume I: Long Papers), pages 2914-2927, Online, August 2021. Association for Computational Linguistics. doi: 10. (0.18653/w1/2001.acchhe 227. URL httph:t/////e..t. https:://aclanthology/ 00rg:///00.0ooooo.022211121 aalalclllnn
227
Linhao Dong, Shuang Xu, and Bo Xu. Speech-transformmer: A no-recurrence sequence-to--sequecc model for speech recognition. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 5884-5888, 2018. doi: 10.1 1L.1109//CASSP2200 8.8462506.
Alexey Dasovitskiy, Lucas Beyer, Alexander Kolesnikov. Dirk Weissenbom. Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkorcit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogmition at scale. In International Conference on Learning Representations, 2021. URL URL
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Addressing some limitations of transformers with feedback memory. arXiv preprint arXiv:2002.09402, 2020.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines, 2014.
10

---

Alex Graves, Greg Wayne, Malcolm Reynolds. Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwinske Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, Adrià Puigdomenech Badia, Karl Moritz, Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukeuoglu, and Demis Hassabis. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626):471-44 476, October 2016. ISSN 00280836. URL http:
7/dx. do1. org/ 10. 1038/ 1038/naturee/eeeee nature20101
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 arXiv:1607.00036 2016.
Caglar Gulcehre, Sarath Chandar, and Yoshua Bengio. Memory augmented neural networks with wormhole connections, arXiv preprint orXiv:1701.065 18, 2017.
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star-transformer, 2019.
Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for transformers. arXiv preprint crXiv:2006.03221 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Newral Comput. 98):1733-1788. November 1997. ISSN 0899-7667. doi: 10.1162/neco.1999 199..9...99999.... httpa////// org/ 1162/
neco beco.1997.33 ..8.7.98..11 1735
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu. David Ding. Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer. al. Perceiver io: general architecture for structured inputs & outputs. arXiv preprint arXiv:2107.14795 2021.
Armand Joulin and Tomas Mikolov. Inferring algorithmic pattems with stack-augmented recurrent nets, 2015.
Da Ju, Stephen Roller, Sainbayar Sukhbaatar, and Jason Weston. Staircase attention for recurrent processing of sequences. arXiv preprint (et:::2000044799. 2021.
Johannes Kiesel, Maria Mestre, Rishabh Shukla, Emmanuel Vincent, Payam Adineh, David Corney, Benno Stein, and Martin Potthast. Semeval-2019 task 4: Hyperpartisan news detection. In Proceedings of the 13th International Workshop on Semantic Evaluation, pages 829-839, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poxter), 2015. URL
http://arxiv. org/ abs/ 1112.111.6601 6980
Guillaume Lample. Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys, 2019.
Jie Lei, Liwei Wang, Yclong Shen, Dong Yu, Tamara L. Berg, and Mohit Bansal. Mart: recurrent transformer for coherent vidco paragraph captioning. 2020.
Yang Liu, Jiaxiang Liu, Li Chen, Yuxiang Lu,, Shikun Feng, Zhida Feng, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. Ernie-sparse: Learning hierarchical efficient transformer through regularized self-attention. arXiv preprint arXiv:2203. 12276, 2022.
Matt Mahoney. Large text compression benchmark, 2006. URL html
Pedro Henrique Martins, Zita Marinho, and André FT Martins. co-former: Infinite memory transformer. arXiv preprint 0rXiv:2109.00333 2021.
Warren S McCulloch and Walter Pitts. A logical calculus of the ideas immanent in nervous activity. The bulletin of mathematical biophysics. 1000:115-133, 1943.
Yuanliang Meng and Anna Rumshisky. Context-aware neural model for temporal information extraction. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 527-536, 2018
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. OpenReeview.net, 2017. URL /attps://openree net/forum? id B2yxedx

---

Alec Radford, Karthik Narasimhan, Tim Salimans. and Πlya Sutskever. Improving language understand- ing by generative pre-training. 2018. URL httts://www.cce ca/ ammbaao0//L1115533/ppprrss
rradford2088mmpooiingg pdf
Jack W Rae. Jonathan Hunt, Tim Harley, Ivo Danihelka, Andrew Senior, Greg Wayne, Alex Graves, and Timothy P Lillicrap. Scaling nemory-augmennt neural networks with sparse reads and writes, 2016.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, and Timothy P. Lillierap. Compressive transformers for rong-range sequence modelling. 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and 1lya Sutskever. Zero-shot text-to-image generation. Marina Meila Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8821-8831. PMLR, 18-24 Jul 2021. URL http::t///// ptttps://proceedingg mlr ppesspressseeee v339/ ramesh2la. html
C Steepenn. Kleene. representation of events in nerve nets and finite automata. Automata studies, 1956.
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 Gomez, Lukasz Kaiser, and Illia Polosukhia. Attention is All you Need. In Advances in neural information processing syatems, pages 5998-6008, 2017. URL http::/ppper:s mips httpp///c//pppp//tl-aaatttttooon
you- need
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020,
Paul J Werbos. Backpropagation through time: what does and how to do it. Proceedings of the IEEE, 78(10): 1550-1560, 1990.
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, et al. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38 45, 2020.
Qingyang Wu, Zhenzhong Lan, Jing Gu, and Zhou Yu. Memformer: The memory- augmeneed transformer. arXiv preprint arXiv:2010.06891, 2020.
Peng Xu, Xinchi Chen, Xiaofei Ma, Zhiheng Huang, and Bing Xiang. Contrastive document representation learning with graph attention networks. arXiv preprint arXi1:2110.1077,, 2021.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. arXiv preprint arXiv:2007. 14062, 2020.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We mention training instabilities
and GPU RAM issues in Section
(c) Did you discuss any potential negative societal impacts of your work? [No] The proposed model and method do not have any specific impacts. All general negative societal impacts applicable to the field could be potentially relative.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material code, training scripts, and raw experimental data in the supplementary material. The supplemental materials would be published on github with the final version of the paper. Instructions for language modeling data&experiments are taken from Tr-XL repo
URL)? [Yes] We include
(b) Did you specify all the training details (e.g., s.. data splits, hyperparameters, how they were chosen)? [Yes] See Section Appendix A and provided supplementary material. (c) Did you report error bars (c.g., with respect to the random seed after running exper- iments multiple times)? [Yes] All the key experiments results are reported with std. Furthermore, we provide raw experimental data in the supplementary materials.
(d) Did you include the total amount of compute and the type of resources used (e.g.. type of GPUs, internal cluster, cloud provider)? [Yes] We used different GPUs depending on the task: 1080Ti, V100, A100. We provide this information in Appendix A for each task.
4. If you using existing assets (e.g., code, data, models) or curating/releasing
assets... (a) If your work uses existing assets, did you cite the creators? [Yes] We refer to the original Tr-XI code and Tr-XI. paper. We use for establishing baselines and setting our methods. See Section 4
(b) Did you mention the license of the assets? [No] Tr-XL license is Apache 2.0 and available at its github repo.
(c) Did you include any new assets either in the supplemental material or as URL? [Yes] Our code is the supplemental material and on GitHub: https: //gi thub com/ booydar/LM-RR RMT
(d) Did you discuss whether and how consent obtained from people whose data you's 're using/curating? [No] We used publicly available Tr-XL code (Apache 2.0) and datasets. (c) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [No] We use either synthetic data datasets collected from the Wikipedia (Wikitext- 103, enwik8).
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "qw9EWphwGEG", "reviewer": "Reviewer_VBgN", "summary": "In this paper, authors propose a recurrent Transformer including the memory to store and process local and global information in a input segment and transfer info between segments. The add special tokens at the start and end of each input segment and define the output of these special tokens as READ and WRITE. The representation of the WRITE will be fed to the next segment. They claim that their model can achieve similar results with Transformer-XL with smaller smaller memory sizes. ", "questions": "Please refer the above section.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The idea is clever and easy to follow. To resolve the issues brought by the autoregressive property, READ token can be the memory for its following tokens and WRITE token can be the summary of this segment which can be used for the next segments. \n2. The experiment is sufficient and comparison to the Transformer-XL is also sufficient to me.\n\nWeakness:\n1. My main concern is the training process of the RMT. To me, since the gradient will be back propagated to the very beginning, is there any forgetting for the memory token? So the token in the last segment will forget the contextual information of the first segment.\n2. Since Transformer-XL will stop the gradient for each segment, I think the training process will be more stable than RMT. Can you confirm this?\n3. Since you claim that RMT can benefit the longer sequence modeling due to the quadratic computational complexity, it would be reasonable to include some other efficient transformer and its alternatives, such as sparse transformer, Luna, S4 in the Table 2 for comparison. ", "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": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "WEP0iwWmBhW", "reviewer": "Reviewer_hj1Y", "summary": "This paper shows a new method for adding recurrence to Transformers via special overlapping special memory tokens between segments, hence, the new Transformer architecture is called Recurrent Memory Transformer (RMT). A special type of tokens, aka memory tokens are prepended and appended to each segment, and the number M for adding how many memory tokens is a new hyperparameter. The gradients can propagate to past segments through memory tokens, which is the distinct feature from previous architecture such as Transformer XL. The paper uses Transformer XL as the baseline.", "questions": "Most of the details about the Transformer architecture and its training procedure are not given in the paper. The width of the Transformer used for the experiments is not given at all. You can't just link to the open-source library, and explain that you used the same configuration defined there, for the following reasons: 1) The GitHub repository may not exist in the future. 2) The configs can change by the owners of the library. 3) It's not a recommended way to communicate experimental details. Although, you would like to contain more interesting results from your findings, but the details about the model architecture, optimization, data processing (tokenization), training procedure are important for reproducibility. And we should all care about reproducibility.\n\nI think the expressiveness of memory doesn't solely depend on the number of memory tokens (M steps), but also by the depth and the width of the Transformer. Did the authors ablate these as well?\n\nWhat are the solutions in practice to tackle training instability and OOM issue? Does this mean that, for large-scale models, TrXL would be more advantageous than RMT? ", "limitations": "The paper didn't properly discussed about the limitations of the proposed method rather than adding only one sentence in Section 5. \"Furthermore, we observed instabilities and out-of-memory issues during RMT training for a larger number of BPTT unrolls and memory sizes\". I suggest being more clear about this and how it can affect the scaling property of the proposed method.\n\nThe negative societal impact was not discussed in the paper, so N/A.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Transformer is easy to perform parallel training, and it suffers less from vanishing gradients compared to recurrent neural networks (RNNs), which made it as popular nowadays. However, Transformer is more expensive than RNNs due to its extensive use of attention mechanism, especially during inference. Another downside is increasing the sequence length modelled by Transformers is difficult as the state-space of the memory is much larger than RNNs. And again, the computational cost for attending to previous context increase linearly as the context length increases. \n\nPrevious works have addressed this issue by caching the Transformer state of immediate previous segment and inventing a new type of positional embedding that only cares about relational difference, or compressing the states of previous context. The former (TrXL) has been often considered as the baseline for Transformers introducing recurrence. Although, TrXL can arrange longer-term contexts for sequence modelling tasks, there are discontinuities in the backward pass so the model doesn't effectively model long-term temporal dependencies. Also, the Transformer state of previous segment is directly given as the memory, which can increase the computational cost and memory requirement for the forward pass. This works provides a way to 1) pass the gradients from the future segments, which is presumably more principled way of modelling long-term dependencies, 2) connect consecutive segments with small overheads (M memory tokens, where M is much smaller size than the segment size, N). \n\nI think the claim \"Recurrence with memory has no limitations on effective context length\" written in the caption of Figure 2 doesn't sound right. It is difficult to increase the context length without trading off other factors such as the size of model parameters due to limited memory space of the hardwares. Which is mentioned in the Section 5, the authors wrote that RMT can suffer from instabilities and OOM issues when using larger number of BPTT unrolls. This could affect the scaling property of RMT. \n\nOverall, the writing is good, it conveys the idea very clearly. Their experiments look convincing, which I will leave my comments below.\n\nComments about the experiments:\nThey results from synthetic tasks looks good (Figure 4). However, the performance of the RMT on WikiText-103 shows that TrXL still performs the best, and RMT can achieve the performance similar to TrXL variant which uses half-sized memory. I think the choice of segment length (only 50 and 150) is quite limited, and also the length of memory in this particular experiment. And 150 tokens sounds quite short as a sequence length in in the first place. It would be nice to add extra longer length such as 300, 600 for this task as well. However, the experiments on enwiki8 (Figure 5) show nicely that the baseline Transformer and TrXL both shows worse results as the sequence length increases, but RMT is more robust and preserves its performance.\n\nIt would be nice if you can write more details about the models, optimizer, tokenization and training procedure. These details should be contained in the main manuscript for reproducibility and readability.\n\nOne thing that I would like to propose to the authors as a reviewer to make the paper to look more stronger, is to add a bit more baselines that tried to compress the previous contexts, e.g., compressive Transformer. Also, RMT and Feedback Transformer has a close relationship, the RMT can be seen as block-level recurrent version of feedback Transformer. One could argue that instead of providing one memory token per each step, RMT provides M memory tokens as the context for processing a segment of N tokens. Maybe adding Feedback Transformer and its variants (providing M memory state similar to RMT for processing arbitrary number of tokens in the new segment).\n\nMinor comment:\nFigure 7 referenced in the main paper is not included in the main, but in the appendix. It would be better to make this more clear", "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": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "TXhmduzzLJN", "reviewer": "Reviewer_ien2", "summary": "The author proposed a memory-augemented transformer where special memory tokens are added to the attention mechanism and test its performance on language modeling and sequence modeling tasks. The author shows that their technique requires less memory in the language modeling tasks to achieve the same accuracy comparing to state-of-the-art model and outperformed baselines where processing long sequence is required.", "questions": "1. How does the method perform on other LM related downstream tasks?\n2. Can the method match the baseline's accuracy and how will that hurt the time and memory requirement for the method?\n\n------------------------------Update------------------------------------\n\nAfter reading the author's response, I believe that the author addressed my second concern. My first concern was partially addressed and further effects on other downstream tasks maybe of independent research interests. The author added discussions regarding the wider impact of the method. In the experiment section, the author included more details about their experiment settings and hyper-parameter selection, when another reviewer raised concerns. This made the experiments much clearer and therefore I would like to raise my rating.", "limitations": "The author may use a separate paragraph to explain further the societal impact of their work as it is related to large language models, which may raise environmental and other concerns.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The method that the author proposed is simple and their experiments showed strong performance of their method. However, the experiments focuses mainly on the copy and reverse tasks and it remains to be seen how the method performs on other tasks. Also, the perplexity that the author shows is slightly worse than that of baseline. It is remained to be shown whether their method can match the perplexity of the baselines and whether that would mean higher training cost (time and memory), etc.", "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": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept