Recurrent Memory Transformer

By Aydar Bulatov, Yuri Kuratov, Mikhail S. Burtsev. 2022

205 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

Landmark Attention: Random-Access Infinite Context Length for Transformers (2023, methodology): “Various methods have been proposed to introduce memorization capabilities to Transformers through recurrence [5, 40].”
RWKV: Reinventing RNNs for the Transformer Era (2023, methodology): “In parallel, RNN-style (Hochreiter and Schmidhuber, 1997; Chung et al., 2014) recursive components have also been modiﬁed to increase context length, such as the Recurrent Memory Transformer (Bulatov et al., 2022, 2023) and Linear Recurrent Units (Orvieto et al., 2023).”
LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding (2023, background): “…2020; Beltagy et al., 2020; Zaheer et al., 2020; Wang et al., 2020; Fedus et al., 2022; Ding et al., 2023), recurrent and memory modules (Dai et al., 2019; Rae et al., 2020; Wu et al., 2022; Martins et al., 2022; Bulatov et al., 2022; Orvieto et al., 2023; Liang et al., 2023; Zhou et al., 2023).”
Extending Context Window of Large Language Models via Positional Interpolation (2023, background): “Several works add memory capabilities to Transformers through recurrence, which increase the models’ capability of handling very long sequences (Bulatov et al., 2022; Wu et al., 2020; Dai et al., 2019; Wu et al., 2022; Martins et al., 2021; Mu et al., 2023).”
Adapting Language Models to Compress Contexts (2023, methodology): “We overcome these challenges by adapting the Recurrent Memory Transformer (RMT) (Bulatov et al., 2022), which can be combined with existing language models.”
Think before you speak: Training Language Models With Pause Tokens (2023, background): “…only as an input has found its use in various forms, most commonly as <cls> (Chang et al., 2023; Liu, 2019; Devlin et al., 2019), <sep> or <mask> in BERT (Devlin et al., 2019) and in a line of work on adding memory to transformers (Burtsev et al., 2020; Bulatov et al., 2022; Darcet et al., 2023).”
Vision Transformers Need Registers (2023, unknown): “In follow-up work, Bulatov et al. (2022) address complex copy-repeat-reverse tasks.”