Recurrent Memory Transformer
Anonymous Author(s)
AffiliationAddress email
Abstract
Transformer-based models show their effectiveness across multiple domains andtasks.
The self-attention allows to combine information from all sequence elements into context-aware representations.
However, global and local informationhas to be stored mostly in the same element-wise representations.
Moreover, thelength of an input sequence is limited by quadratic computational complexity ofself-attention.
In this work, we propose and study a memory-augmented segmentlevel recurrent Transformer (Recurrent Memory Transformer).
Memory allowsto store and process local and global information as well as to pass informationbetween segments of the long sequence with the help of recurrence.
We implementa memory mechanism with no changes to Transformer model by adding specialmemory tokens to the input or output sequence.
Then Transformer is trained tocontrol both memory operations and sequence representations processing.
Resultsof experiments show that our model performs on par with the Transformer-XLon language modeling for smaller memory sizes and outperforms it for tasks thatrequire longer sequence processing.
This makes Recurrent Memory Transformer apromising architecture for applications that require learning of long-term dependencies and general purpose in memory processing, such as algorithmic tasks andreasoning.
Introduction
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 the previous segment.
Transformers (Vaswani et al., 2017) have been widelyadopted across multiple domains and tasks (Radfordet al., 2018;
Dong et al., 2018; Devlin et al., 2019;
Dosovitskiy et al., 2021; Ramesh et al., 2021; Jaegleet al., 2021).
The key component of Transformerlayer is a self-attention.
Self-attention allows to update each sequence element representation with information from all other elements in the sequence.
As a result, rich contextual representation for everyelement is generated at the end of encoding.
Thisway, global sequence-level and local information arestored in a single representation.
However, this mixing of two types of information in a single representation has limitations.
Distributed storage of globalfeatures across all sequence elements results in globalfeatures "blurring" and makes it harder to access them.
Another well-known deficiency of Transformers is poor scaling of self-attention with input sequencelength that hurts its applications to long inputs (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., 2020; Choromanski et al.,2020).
Our work introduces a memory-augmented segment-level recurrent Transformer named Recurrent
Memory Transformer (RMT).
RMT uses a memory mechanism based on special memory tokens (Burtsev et al., 2020) added to the input sequence.
Memory tokens provide additional reserved capacity tothe model that could be used to process information which is not directly representing any element inthe input sequence.
To process long sequences, we split them into segments and pass memory statesfrom a previous to a current segment.
This memory passing makes the model recurrent and removesthe input sequence length limitations.
RMT model can theoretically work with infinite lengths but, inpractice, 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 Transformermodel because modifications are made only to the input and output sequences of the model.
We tested RMT on the tasks that require a global information about the whole input sequence to besolved.
We use copy, reverse, and associative retrieval tasks in the setting where the input sequenceis split into segments.
RMT and Transformer-XL perfectly solve these tasks, but exceeding somevalue of sequence length, RMT starts to outperform Transformer-XL.
Also, we experimentally showthat the proposed Recurrent Memory Transformer requires less memory size to perform closely to
Transformer-XL on language modeling tasks.
RMT code and experiments are available 1 .
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-XL for sequenceprocessing tasks and on par with Transformer-XL on language modeling but requires lessmemory.
• We analysed how the Transformer model learns to use memory.
Specific interpretable memory read-write patterns of attention are shown.
Related work
In our study we add a memory to general purpose attention based neural architecture.
Memory isa 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 Backpropagation
Through Time learning algorithm (Werbos, 1990) and Long-Short Term Memory (LSTM) (Hochreiterand Schmidhuber, 1997) neural architecture.
Today memory-augmented neural networks (MANNs)usually rely on some kind of recurrent external-memory which is separate from the model’s parameters.
Neural Turing Machines (NTMs) (Graves et al., 2014) and Memory Networks (Westonet al., 2014) are equipped with a storage for vector representations that can be accessed with anattention mechanism.
Memory Networks (Weston et al., 2014; Sukhbaatar et al., 2015) were designedto enable reasoning by sequential attention over to the content of a 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 thesemodels are differentiable and can be trained via backpropagation through time (BPTT).
Parallel lineof 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 moreadvanced addressing mechanisms such as address-content separation and multi-step addressing wereproposed in (Gulcehre et al., 2016, 2017; Meng and Rumshisky, 2018).
The Global Context Layermodel (Meng and Rumshisky, 2018) uses the idea of address-content separation to solve the difficultyof training content-based addressing in the canonical NTM.
Recent rise of Transformer models also resulted in introduction of a number of new memory architectures.
Transformer-XL (Dai et al., 2019) introduces a segment-level recurrence at the level ofhidden representations.
These representations of a sequence are computed and stored in the cacheto be reused as an extended context for the next segment.
Compressive Transformer (Rae et al
2019) adds the second layer of memory to Transformer-XL.
This memory compresses and storesinformation from the cache.
∞ -former (Martins et al., 2021) utilizes continuous-space attention andrepresents input sequence as a continuous signal to make long-term memory unbounded.
Memory
Layers (Lample et al., 2019) model has a product key memory layer instead of a feed-forward layerwithin Transformer block to increase model capacity.
In many variations of Transformer different sorts of global representations are added.
Among themare Star-Transformer (Guo et al., 2019), Longformer (Beltagy et al., 2020), GMAT (Gupta and Berant,2020), Extended Transformer Construction (ETC) (Ainslie et al., 2020) and Big Bird (Zaheer et al.,2020).
All these architectures re-design self-attention mechanism to reduce it computational complexity 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 sequencewith special memory tokens.
Perceiver IO (Jaegle et al., 2021) maps an entire arbitrary input to thefixed number of latent representations.
Transformer layers do further processing over latent memoryrepresentations only.
Segment-level recurrence in Transformers is actively explored in a number of studies.
TransformerXL, Compressive Transformer keep previous states and re-use them in subsequent segments.
ErnieDoc (Ding et al., 2021) improves processing by using same-layer recurrence instead of attending toprevious layer outputs of a precedent segment.
Memformer (Wu et al., 2020) introduces a dedicatedmemory module to keep previous hidden states in summarized representations.
Memformer uses twospecial layers added to the Transformer model.
Memory cross-attention layer reads from memoryand memory slot attention layer updates it.
MART (Lei et al., 2020) has a similar approach as
Memformer but uses memory update rules analogous to LSTM (Hochreiter and Schmidhuber, 1997)and GRU (Cho et al., 2014).
FeedBack Transformer (Fan et al., 2020) goes further with full, and notsegment-level, recurrence.
FeedBack Memory merges past hidden representations from all layersinto a single vector and makes it accessible to the computations at any layer.
The disadvantage of fullrecurrence is that it is less parallelizable.
FeedBack Memory requires every sequence element to beprocessed sequentially.
In segment-level recurrent models, all elements of a segment are processed by
Transformer layers in parallel.
Only segments are processed sequentially.
Staircase Transformer (Juet al., 2021) combines segment-level recurrence and depth recurrence.
Staircase models use theoutput for previous segments and pass them as input for the next segment.
Our Recurrent Memory
Transformer is based on special memory tokens similar to Memory Transformer, segment-levelrecurrence as in Transformer-XL, and depth-recurrent mechanism for memory processing similar to
Staircase.
Recurrent Memory Transformer
[Figure]
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. Recurrence with memory has no limitations on effective context length, whereas Transformer-XL can use only finite context with cached states. All RMT memory/recurrence operations are made on the input and output level of the Transformer model.
3.1 Background
: Transformer-XL
Transformer-XL (Dai et al., 2019) extends Transformer model with state re-use cache mechanismfor segment-level recurrence and relative position encoding.
Input sequence is split on segmentsprocessed sequentially.
Hidden states computed for the previous segment M n are cached for eachtransformer layer n .
The input of the layer n consists of the last m states from the cached memoryand output of previous Transformer layer for the current segment τ :
[Equation]
here, SG stands for stop-gradient, ◦ denotes concatenation.
Cached states allow to increase effectivecontext size of Transformer model and save on compute operations.
Then, ˜ H n − 1 τ goes to Transformer layer to produce layer n outputs for segment τ :
[Equation]
In Transformer-XL, self-attention layers are modified to use relative position encodings to improvegeneralization to longer attention lengths.
The overall architecture is shown in the Figure 2.
Memory augmented Transformers such as GMAT, ETC, Memory Transformer (Gupta and Berant,2020; Ainslie et al., 2020; Burtsev et al., 2020) proposed to use special global tokens as storagefor 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 memorytokens at the start of the sequence to collect information from the subsequent tokens.
On the otherhand, if memory tokens are placed at the end of the sequence then preceding tokens unable toaccess their representations.
To solve this problem we add a recurrence to the sequence processing.
Representations of memory tokens placed at the end of the segment are used as an input memoryrepresentations at the start as well as at the end of the next segment.
In the Recurrent Memory Transformer input is augmented with special [mem] tokens, processed ina standard way along with the sequence of tokens.
Each memory token is a real-valued vector.
mmemory tokens are added at the beginning of the segment tokens representations H 0 τ and the same mtokens are added at the end:
[Equation]
here N is a number of Transformer layers.
The starting group of memory tokens functions as a read memory that allows sequence tokens toattend to memory states produced at the previous segment.
The ending group works as a writememory that can attend to all current segment tokens and update representation stored in the memory.
As result, H writeτ contains updated memory tokens for the segment τ .
Segments of the input sequence are processed sequentially.
To enable recurrent connection betweensegments, we pass outputs of the memory tokens from the current segment to the input of the nextsegment:
[Equation]
Both memory and recurrence in the RMT are based only on global memory tokens.
It allows tokeep the backbone Transformer unchanged and make RMT memory augmentation compatible withany model from the Transformer family.
Memory tokens operate only on the input and output of the model.
In this study we implement RMT on top of the original Transformer-XL code.
Botharchitectures are shown in Figure 2.
Recurrence in the RMT is different compared to the Transformer-XL because the former stores onlym memory vectors per segment.
On the other hand, the Transformer-XL stores m × N vectors persegment.
Also, in the RMT model memory representations from the previous segment are processedby Transformer layers together with the current segment tokens.
This makes memory part of RMTeffectively deeper in a number of applied Transformer layers τ × N .
Additionally, we allow allmemory tokens in the read/write block to access all other tokens in the same block.
The causalattention mask is applied only to tokens of the input sequence. The RMT attention mask is shown in Figure 6 (d).
We train the RMT with Backpropagation Through Time (BPTT).
During backward pass, unlike in
Transformer-XL, memory gradients are not stopped between segments.
The number of previoussegments to backpropagate is a hyperparameter of a training procedure.
We vary BPTT unroll in ourexperiments from 0 to 4 previous segments.
Increasing this parameter is computationally expensiveand requires a lot of GPU RAM.
However, such techniques as gradient checkpointing could be usedto alleviate this problem.
Experiments
We designed our experiments to evaluate the ability of Recurrent Memory Transformers to preservelong-term dependencies across multiple input segments.
The first set of experiments includes copy,reverse, associative retrieval and quadratic equations tasks.
The second addresses language modelingtask for word-level on WikiText-103 (Merity et al., 2017) and for character-level on enwik8 (Mahoney,2006).
We compare Recurrent Memory Transformer with Transformer and Transformer-XL models.
4.1 Algorithmic Tasks
Firstly, we evaluate RMT on algorithmic tasks that require information about the whole input sequenceto be solved successfully.
In a recurrent setting, the model has to keep information about all previoussegments to make predictions. decoder-only models.
Dotted lines show segment borders.
In the copy task, an input sequence should be replicated twice after a special start-to-generate token.
Inthe reverse task, an input sequence should be generated in a reverse order.
Input for the associativeretrieval task consists of N key-value pairs.
Then onekey is randomly selected, and the task is to producean appropriate value for the selected key.
Anothertask is to solve quadratic equations.
One exampleconsists of an equation, its solution with discriminant,and an answer.
The task is to generate a solution andanswer, while only answer quality is evaluated.
For all tasks, input and output sequences are splitinto segments and processed by models sequentially.
Datasets for algorithmic tasks were randomlypre-generated, and the same data was used in all experiments.
Transformer-XL and RMT are decoder-only Transformer models.
We do not compute loss overthe input sequence before the start-to-generate token.
The loss is computed over target sequencesegments only. This procedure is illustrated in Figure 3.
4.2 Language Modeling
We use two standard benchmarks for language modeling: WikiText-103 and enwik8.
[-WikiText-
-]{+WikiText
+}
(Merity et al., 2017) is used for word-level language modeling and contains 103M words from English Wikipedia articles.
Enwik8 (Mahoney, 2006) is used for character-level and consists of 10 8first bytes of XML text dump of the English Wikipedia.
We compare Recurrent Memory Transformer with decoder-only Transformer and Transformer-XLas baselines.
Model size and training parameters are selected to match Transformer-XL 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.
Weincreased 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 languagemodeling tasks harder for RMT because information has to be stored in the same amount of memoryfor more time steps.
Our RMT implementation is based on Transformer-XL repository 2 .
We also use WikiText-103,enwik8 data and processing from this repository.
Language modeling experiments follow the samemodel and training hyperparameters as Transformer-XL. WikiText-103 experiments use 16-layer Transformers, enwik8 – 12 layer Transformers.
We refer to Transformer-XL with memory size equalto zero as a Baseline. With this experimental setup we were able to reproduce results for the Transformer-XL model closeto the original paper (see Appendix A.4 for enwik8 and Table 2 for WikiText-103).
Results
Baseline, Tr-XL, RMT perform perfectly in the single segment setting on copy and reverse tasks(Figure 4).
In this case models do not need recurrence because the whole sequence is available.
Whenthe number of segments is larger than one, non-recurrent baseline struggles to solve tasks, but bothmemory models demonstrate ability to retain required information from the previous segments inmemory.
As the number of segments increases, Recurrent Memory Transformer starts to outperform Transformer-XL with memory sizes less than the number of all previous tokens.
With the number ofsegments up to 6 mean accuracy of Transformer-XL drops by up to 0.2 points, and with 9 segmentsscore plunges to an accuracy score of 0.2, close to the baseline without memory.
Associative retrieval results are similar with the number of segments up to 4.
RMT manages tosolve the task with Transformer-XL closely behind.
However, in the setting with 5 segments, RMTperformance slightly decreases and Transformer-XL average accuracy rises higher.
On the quadratic equations task (Table 1) we have checked that it is possible to solve the task withthe Transformer baseline and no segmentation used.
The baseline in this case defines upper bound for this task.
Enabling recurrence with multiple segments RMT solves the task perfectly, while Transformer-XL finds the task challenging.
Input sequence length is 180 tokens and consists of quadratic equation, solution, and answer.
The input sequence is split into a number of segments with an answer as the last segment.
Accuracy is 1.0 if the full answer is predicted correctly.
[Table]
Results of experiments on word-level languagemodeling on WikiText-103 are shown in Table 2.
In the first section with a segment length of 150, Tr-XL and RMT outperform the baseline and Memory Transformer (MT) by a large margin.
It shows the significance of increased effectivecontext length by Tr-XL cache or RMT memoryfor language modeling.
MT is slightly better than the Transformer baseline.
This is due to the fact that MT adds specialmemory tokens only to the beginning of an input sequence.
Autoregressive MT has no wayto write to memory because of the causal attention mask and, therefore, is unable to passinformation between segments.
Thus, in the autoregressive setting MT could be seen as equivalent to the prefix/promt-tuning (Li and Liang,2021; Lester et al., 2021).
[Table]
RMT improves over MT memory mechanismwith read/write blocks.
The best RMT modelswith memory size 10 and 25 show similar performance as Transformer-XL with a memory sizeequal to 75.
RMT learns to use smaller memorymore effectively than Transformer-XL.
Additionally, the smaller memory size of RMT leadsto reducing required GPU memory for runningthe model.
Decreasing the size of segments to 50, we forcemodels to work with longer recurrent dependencies as the number of recurrent steps increases.
RMT with memory consisting of a single vector shows similar results to Transformer-XLwith memory size 10.
It is worth noting that Transformer-XL memory consists of hidden representations from all layers (in this case, it is10 × 16 vectors) when RMT memory is onlymemory_size vectors.
Transformer-XL with memory size 50 and RMT with memory size 5 showssimilar perplexity values.
On enwik8 RMT models with memory size 5 and Transformer-XL with memory size 40 showsimilar results.
Confirming that RMT learns to use smaller amounts of memory representation moreeffectively.
All experiments results on enwik8 dataset are in Appendix A.4 and Table 3.
We analyze how number of segments, sequence length, length of training context, and memory sizeaffect models’ performance on different tasks in Figure 5.
As we split sequence into more segmentsit becomes more crucial to be able to pass information between segments.
For the copy task, wesplit 360 tokens sequence into multiple segments.
In Figure 5a we observe that Tr-XL performancestarts to degrade and eventually falls to the baseline model performance as the number of segmentsincreases.
In contrast, RMT continues to solve the task perfectly. In the more extreme setting, whenwe keep memory size fixed, but increase the total length of sequence Tr-XL fails shortly, while RMTstarts to degrade on 720 tokens sequence length (Figure 5b). Recurrent Memory Transformer learns to make predictions depending on #BPTT_unrolls previoussegments +1 current segment.
Transformer-XL does not use BPTT and relies only on memory_sizecached states and current segment making in total: memory_size + segment_length tokens. In [Figure] [Figure]
[Figure] [Figure] [Figure] Figure 5c, we compare RMT and Tr-XL according to the described value of visible context at trainingtime.
RMT with a single memory vector could be trained to achieve lower perplexity as Transformer-XLwith memory size 10.
It means that RMT can learn to store information from previous observationsmore compactly.
Another observation is that RMT with memory sizes 10 and 25 performs worse butclosely to Transformer-XL even when Transformer-XL has access to more non-compressed states(50, 100, 200) from previous segments.
Furthermore, we observed instabilities and out-of-memoryissues during RMT training for a larger number of BPTT unrolls and memory sizes.
Recurrent Memory Transformer does not benefit from increasing memory size from 5 to 50, butresults of Transformer-XL better scale with memory size (Figure 5d).
RMT models with memorysize 5 have close results to Transformer-XL with cache 50, confirming that RMT learns to store morecompact representations.
Dynamic of RMT perplexity suggests that there is some optimal memorysize for RMT to solve the task, and further increase does not add much.
Training RMT with one
BPTT unroll drastically improves its results showing the importance of BPTT training (Figure 5d).
[Figure]
To get an understanding of memory operations, learned by models during training algorithmic tasks,it is useful to look at attention maps.
Figure 6 shows some heatmaps from attention layers of modelstrained on copy and reverse.
The darkness of each pixel shows how much the element from thecorresponding row "attends" to its column.
In each RMT attention map sequence tokens are preceded by read memory, located at the top leftcorner, and followed by write memory at the bottom right.
Lines at the central part of (a), top imageshows classic attention of token sequence to itself, but the bottom line represents the operation ofwriting of sequence tokens to memory in straight order.
When completing reverse, the (a), bottomimage model learns to write the sequence in the reversed order, which is in line with common sense.
When it comes to reproducing the target sequence, model accesses memory, Figure 6 (b) and writesto the output sequence.
Another operation (c) is rewriting from read memory to write memory.
It iscommonly used by RMT in settings with larger number of segments to keep information about recentsegments longer.
Transformer-XL mechanism of accessing memory (d) does not allow straightforward writing to memory without changing sequence token representations.
Sequential reading from cache is representedby straight lines on Transformer-XL attention maps.
Using token representations as storage harms model performance in tasks with larger number ofsegments.
On reverse with 4 segments Transformer-XL with limited memory size 6, Figure 8 (b)attempts to mix representations of tokens and read multiple symbols from one cached state in nextsegments. This results in the average accuracy of 0.8 on the given task.
Despite having the samememory size, RMT manages to compress the whole segment in memory tokens Figure 8 (a) andachieve mean accuracy 1.0.
Visualizations from Figure 6 and Figure 8 provide evidence to support our hypotheses that Tr-XLhas to mix representations from previous and current segments in the same hidden states to passinformation between segments.
Also, visualizations show how memory tokens in RMT help mitigatesuch kind of mixing.
RMT ability of sequence compression to memory is illustrated in Figure 7.
For copy with 6segments RMT compresses and then reads the sequence of 12 tokens with just 6 memory tokens.
For Transformer-XL decreasing memory size harms the accuracy score significantly with number ofsegments larger than 2.
Conclusions
In this paper we introduced Recurrent Memory Transformer a simple recurrent memory augmentationof Transformer model.
RMT is implemented by extension of an input sequence with special globalmemory tokens and segment-level recurrence.
In our experiments we compared RMT with Transformer baseline and Transformer-XL whichis a well-known modification of Transformer for long sequences.
RMT almost perfectly solves
Copy, Reverse as well as quadratic equations tasks for sequences consisting of multiple segmentsoutperforming Transformer-XL.
It also demonstrates quality for associative retrieval task on parwith Transformer-XL.
As expected, baseline Transformer fails to solve these tasks for multi-segmentsettings.
RMT trained as language model significantly ahead of Transformer baseline and shows qualitymetrics similar to Transformer-XL but for up to 10 times smaller memory size.
Experimental resultsdemonstrate that for fixed memory size backpropagating gradients for more segments improvesperformance of RMT.
Analysis of attention maps suggests that better RMT performance can berelated to more effective storage of input representations in dedicated memory tokens compared tomixing representations storage in Transformer-XL.
Overall, results of the study show that dedicatedmemory storage and recurrence provided by Recurrent Memory Transformer make it a promisingarchitecture for applications that require learning of long-term dependencies and general purposein-memory processing, such as algorithmic tasks and reasoning.
Furthermore, we believe that RMTcould open the way for adding memory and recurrence to other models in the Transformer family.
Checklist
[Table]
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 instabilitiesand GPU RAM issues in Section 5.(c)
Did you discuss any potential negative societal impacts of your work?
[No] Theproposed 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]
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]
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 or as a URL)?
[Yes] We include
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. (b) Did you specify all the training details (e g., data splits, hyperparameters, how they were
chosen)?
[Yes] See Section 4.3, Appendix A, and provided supplementary material.
(c) Did you report error bars (e 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., typeof GPUs, internal cluster, or cloud provider)?
[Yes] We used different GPUs dependingon the task: 1080Ti, V100, A100.
We provide this information in Appendix A for eachtask.
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We refer to the
original Tr-XL code and Tr-XL paper.
We use it for establish baselines and setting our
methods.
See Section 4.
(b) Did you mention the license of the assets?
[No] Tr-XL licence is Apache 2.0 andavailable at its github repo.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
Our code is in the supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[No] We used publicly available Tr-XL code (Apache 2.0) and datasets.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[No] We use either synthetic data or datasets collected
from the Wikipedia (Wikitext-103, enwik8).
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A](c)
Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
A Training details and additional resultsA.1 Algorithmic tasks
Datasets were randomly generated by uniformly sampling tokens from dictionary into task sequencesand generating targets accordingly to the tasks.
After generation, datasets are fixed for all experiments.
Copy and reverse use sequences of sizes 24, 40, 120, 240, and 360, making total copy/reverseinput length 48/72, 80/120, 240/360, 480/720, 720/1080.
The associative retrieval task consists of 4key-value pairs and one randomly selected key; the answer consists of one value.
Train, validationand test sizes of copy 24, reverse 24 and associative retrieval datasets are 100000, 5000 and 10000.
Transformer-XL had the same cache size on training and validation to match RMT.
For training all models on copy and reverse, we used constant learning rate 1e-4 with reduction onplateau with decay factor of 0.5.
Copy and reverse were solved by models with 4 layers and 4 heads,associative retrieval models had 6 layers and 4 heads.
Models with the same context size and memorysize were trained for the same number of steps and the same training parameters.
Experiments with sequence length 24 were conducted on a single Nvidia GTX 1080 Ti GPU from 1hour to 2-3 days.
Copy and reverse on longer sequence lengths were done on more powerful Tesla
V100 using 1-3 devices with training time varying from 1 hour to 3-4 days.
A.2 Associative retrieval
We used code for the task dataset generation from (Ba et al., 2016) 3 .
A.3 Quadratic equations
This dataset consists of equations with integer coefficients with step-by-step solutions using adiscriminant.
Process of equation generation is started from uniformly sampling real roots x 1 , x 2from -100 to 100.
The answer of an equation is represented as x 1 , x 2 .
Next, we find the equation asmultiplication of two parentheses ( x − x 1 )( x − x 2 ) = 0 , which is expanded to x 2 − ( x 1 + x 2 ) x
+x 1 x 2 = 0 .
Next, we multiply all coefficients by a random natural number α from 1 to 10.
The finalequation form is αx 2 − α ( x 1 + x 2 ) x
+ αx 1 x 2 = 0 .
A dataset sample is made of these stages inreversed order.
We also provide a string with the discriminant calculation to help find the equationroots.
20 percent of equations in the dataset do not have real roots.
Example equation string:-4*xˆ2+392*x-2208=0 ,solution string:xˆ2-98*x+552=0;D=98ˆ2-4*1*552=7396=86ˆ2;x=(98-86)/2=6;x=(98+86)/2=92 ,and answer:6,
Each solution step is tokenized on char level and padded to the length of 30 tokens.
The total lengthof each training sample is 180, the dataset has 100000 training, 10000 validation and 20000 testsamples.
For this task we used models with 6 layers, 6 heads and segment sizes 180 and 30.
Trained wasperformed with the same schedule as copy and reverse on a single GTX 1080 ti for 1-2 days.
Memorysize for RMT and Transformer-XL was chosen equal to the segment length.
A.4 Enwik8
We verified our experimental setup by reproducing Transformer-XL results on enwik8 dataset(Table 3).
We used 12-layer Baseline (Transformer), Transformer-XL, RMT in all enwik8 experiments.
All results on enwik8 dataset are in Table 3.
We used 2 NVIDIA A100 80Gb GPUs, training timevaried from 10 to 30 hours depending on sequence length, memory size, and number of BPTT unrolls.
A.5 WikiText-103
We used 16-layer models in all experiments on WikiText-103 dataset.
Training hyperparameters wereused from (Dai et al., 2019) and authors PyTorch scripts 4 .
All results on WikiText-103 dataset are in Table 4.
In most of the WikiText-103 experiments, we used 2 NVIDIA A100 80Gb GPUs, trainingtime varied from 10 to 30 hours depending on sequence length, memory size, and number of BPTTunrolls.
[Table] [Table]