PaperlayerView exact revision on arXiv

Source-linked AI summary

Recurrent Memory Transformer

Aydar Bulatov, Yuri Kuratov, Mikhail S. Burtsev

arXiv:2207.06881v2cs.CLcs.LG

TL;DR

Transformers struggle to keep global and local information distinct and to process long inputs efficiently. RMT adds token-based memory and segment-level recurrence without changing the Transformer itself. It matches Transformer-XL on language modeling with up to 10 times less memory and outperforms it on several multi-segment tasks.

  • Problem

    Transformers mix global and local information in element-wise representations, while quadratic self-attention complexity limits input sequence length.

  • Method

    RMT extends a Transformer with special memory tokens and segment-level recurrence, passing memory states between segments without modifying the Transformer model.

  • Results

    RMT outperforms Transformer-XL on several multi-segment sequence-processing tasks and achieves similar language-modeling quality with up to 10 times smaller memory.

  • Takeaways & Limitations

    Dedicated memory storage and recurrence make RMT promising for long-term dependencies and general-purpose in-memory processing, including algorithmic tasks and reasoning.

  • Takeaways & Limitations

    Increasing the number of previous segments used for backpropagation is computationally expensive and requires substantial GPU memory.

Abstract

from arXiv · show

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.

1 Introduction

Transformers mix global and local information within element-wise representations while self-attention scales poorly for long inputs. RMT addresses these constraints with token-based memory and segment-level recurrence, outperforming Transformer-XL on longer-sequence tasks and using less memory for comparable language-modeling quality.

  • Motivation: Self-attention updates each sequence element using information from all other elements, producing contextual representations that mix global and local information.The paper notes that distributed global features can blur and become harder to access.
  • Motivation: Quadratic self-attention complexity limits Transformer applications to long input sequences.
  • Approach: RMT adds special memory tokens, splits long inputs into segments, and passes memory states between segments without modifying the Transformer model.The memory tokens provide reserved capacity for information not directly representing input elements.
  • Experiments: RMT and Transformer-XL solve copy, reverse, and associative-retrieval tasks, with RMT outperforming Transformer-XL beyond some sequence length.RMT also requires less memory to perform closely to Transformer-XL on language modeling.
  • Experiments: The study evaluates RMT, vanilla Transformer, and Transformer-XL on copy, reverse, associative retrieval, and language modeling tasks.
  • Further analysis: Combining Transformer-XL cache with RMT improves language-modeling performance, while attention analysis reveals interpretable memory read-write patterns.

2 Related work

Related work has explored external memory, global representations, and recurrence to extend Transformer capabilities. RMT belongs to this line of research while using token-based memory and segment-level recurrence.

  • Memory-augmented networks: Memory-augmented neural networks commonly use recurrent external memory separate from model parameters.Neural Turing Machines and Memory Networks access stored vector representations with attention.
  • Transformer memory: Transformer-XL introduces segment-level recurrence by caching hidden representations for reuse as extended context in the next segment.
  • Transformer memory: Compressive Transformer adds a second memory layer that compresses and stores information from the Transformer-XL cache.
  • Efficient attention: Several Transformer variants redesign self-attention or add global representations to reduce computational complexity and cover the input.Examples include Star-Transformer, Longformer, GMAT, ETC, and Big Bird.
  • Recurrent architectures: Other recurrent approaches retain, summarize, or reuse previous hidden states through mechanisms such as same-layer recurrence and dedicated memory modules.The cited examples include Ernie-Doc and Memformer.

3 Recurrent Memory Transformer

RMT adds global memory tokens and segment-level recurrence to a Transformer, passing memory states between sequentially processed segments. This preserves the Transformer backbone while enabling memory processing across longer contexts, with BPTT training through memory.

  • 3 Recurrent Memory Transformer: RMT adds special read and write memory tokens to the input, using updated write representations as memory for the next segment.Read tokens access previous-segment memory, while write tokens attend to current-segment tokens and update the stored representation.
  • 3 Recurrent Memory Transformer: Segments are processed sequentially, and memory-token outputs are passed from the current segment into the next to create recurrence.
  • 3 Recurrent Memory Transformer: The memory mechanism operates only on global tokens, leaving the backbone Transformer unchanged and making the augmentation compatible with Transformer-family models.
  • 3 Recurrent Memory Transformer: RMT stores m memory vectors per segment, whereas Transformer-XL stores m × N vectors, while processing RMT memory through Transformer layers across segments.RMT memory is effectively processed through τ × N applied Transformer layers.
  • 3 Recurrent Memory Transformer: RMT trains with BPTT without stopping memory gradients between segments, but increasing the unroll is computationally expensive and requires substantial GPU RAM.Experiments vary BPTT unroll from 0 to 4 previous segments; gradient checkpointing is suggested as a mitigation.

4 Experiments

The experiments test RMT on algorithmic tasks, language modeling, and long-text classification, comparing it with Transformer and Transformer-XL baselines. They vary recurrent context settings and also add recurrent memory to pretrained Transformer-family models.

  • 4 Experiments: Experiments evaluate long-term dependency preservation on copy, reverse, associative retrieval, and quadratic-equation algorithmic tasks.Inputs and outputs are split into segments and processed sequentially.
  • 4 Experiments: RMT is compared with decoder-only Transformer and Transformer-XL baselines on WikiText-103 word-level and enwik8 character-level language modeling.The setup matches Transformer-XL model size and training parameters, with context lengths of 150 tokens and 512 characters, respectively.
  • 4 Experiments: Language-modeling experiments increase recurrent steps by shortening segments, making the tasks harder while testing longer-term dependencies.Segments are reduced to 50 tokens for WikiText-103 and 128 characters for enwik8.
  • 4 Experiments: RMT is also evaluated on Hyperpartisan news classification by augmenting pretrained BERT-base, RoBERTa-base, DeBERTa-base, and T5-base models.Each model receives 500 input tokens and recurrent memory of size 10 before fine-tuning.

5 Results

Across algorithmic and language-modeling evaluations, RMT retains information across segments more effectively than Transformer-XL, while using smaller memory sizes in comparable language-modeling settings. Recurrence, memory size, and training-context length all affect performance, and combining RMT memory with a Transformer-XL cache performs best on WikiText-103.

  • Segmented algorithmic tasks: Both memory models solve copy and reverse perfectly with one segment, whereas the non-recurrent baseline struggles once sequences span multiple segments.The memory models retain required information from previous segments.
  • Segmented algorithmic tasks: RMT begins outperforming Transformer-XL on Copy and Reverse as the number of segments increases, despite using memory smaller than all preceding tokens.Transformer-XL accuracy drops by up to 0.2 points by six segments and approaches the no-memory baseline at nine segments.
  • Segmented algorithmic tasks: RMT solves the 360-token Copy task perfectly through nine segments, while Transformer-XL performance degrades toward the baseline as segmentation increases.The comparison uses a fixed sequence length of 360 tokens.
  • Segmented algorithmic tasks: With multiple segments, RMT solves the Quadratic Equations task perfectly, while Transformer-XL finds it challenging.The task places the answer in the final segment of a 180-token sequence.
  • Language modeling: On WikiText-103, RMT models with memory sizes 10 and 25 achieve performance similar to Transformer-XL with memory size 75, while requiring less memory.Combining Transformer-XL cache with recurrent memory gives the best performance.
  • Language modeling: Deeper BPTT unrolling improves RMT scores, but larger memory sizes and deeper unrolls introduce training instabilities and out-of-memory issues.Larger visible training context lowers perplexity for both RMT and Transformer-XL.

6 Conclusions

RMT extends Transformers with token-based memory and segment recurrence, achieving strong multi-segment task performance and language-model quality with substantially smaller memory. The results also suggest dedicated memory storage as a useful mechanism for long-term dependencies and transfer to pretrained models.

  • RMT almost perfectly solves Copy, Reverse, and quadratic-equation tasks across multiple segments, outperforming Transformer-XL, while matching it on associative retrieval.
  • Attention-map analysis suggests dedicated memory tokens store input representations more effectively than Transformer-XL’s mixed representation storage.
  • Dedicated memory storage and recurrence make RMT promising for long-term dependencies, in-memory processing, algorithmic tasks, and reasoning.

Checklist

The checklist reports that the paper documents its limitations, code, data, training details, error bars, compute resources, and asset licensing. It uses synthetic or publicly available datasets and reports no specific societal impacts or human-subject procedures.

  • The authors report training instabilities and GPU RAM issues as limitations discussed in Section 5.
  • Code, training scripts, raw experimental data, training details, and error bars are reported as available in supplementary materials or repositories.
  • The experiments use publicly available Tr-XL code and datasets, with licensing and asset attribution addressed in the checklist.
  • The data are synthetic or collected from Wikipedia, and the checklist reports no personally identifiable information or offensive-content discussion requirement.
  • Human-subject instructions, risks, and compensation are marked not applicable, while the paper reports no specific negative societal impacts.

A.1 Algorithmic tasks

The algorithmic-task experiments use randomly generated, character-tokenized segmented sequences to evaluate memory across copy, reverse, associative retrieval, and related tasks. RMT solves longer segmented tasks with limited memory, as illustrated by the reported figures and setup.

  • Datasets are randomly generated by uniformly sampling dictionary tokens, with fixed datasets used across experiments.
  • Copy and reverse evaluate sequence sizes 24, 40, 120, 240, and 360, while associative retrieval uses four key-value pairs and one selected key.
  • RMT solves copy and reverse almost perfectly with limited memory, unlike Transformer-XL, on test-set per-character accuracy evaluations.
  • Quadratic-equation samples reverse the stages of equation generation, use character-level padded solution steps, and include 20 percent equations without real roots.

A.5 WikiText-103

The WikiText-103 experiments compare RMT and Transformer-XL using test-set perplexity, while the enwik8 setup reports test-set bits-per-character for experimental verification.

  • WikiText-103 experiments use 16-layer models and mostly train for 200k batches on two NVIDIA A100 80Gb GPUs.
  • Table 4 reports test-set bits-per-character on enwik8 and indicates similar scores to the original Transformer-XL paper with segment length 512.
  • Table 5 reports test-set perplexity for RMT and Transformer-XL on WikiText-103.

B Operations with Memory

RMT compresses an entire segment into memory tokens, whereas Transformer-XL mixes token representations and reads multiple symbols from a cached state.

  • RMT compresses the whole segment into memory tokens.The figure compares compression and decompression for a sequence of length 12 and memory of size 6.
  • Transformer-XL mixes token representations and reads multiple symbols from one cached state.
Loading 2207.06881v2…