PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Transformer to 1M tokens and beyond with RMT

Aydar Bulatov, Yuri Kuratov, Yermek Kapushev, Mikhail S. Burtsev

arXiv:2304.11062v2cs.CLcs.AIcs.LG

TL;DR

Transformer self-attention’s quadratic compute limits input length, motivating methods to extend pretrained models’ context. This paper adds recurrent token memory and curriculum learning to process long sequences; its experiments report handling up to two million tokens on specialized tasks with linear computational scaling.

Abstract

from arXiv · show

A major limitation for the broader scope of problems solvable by transformers is the quadratic scaling of computational complexity with input size. In this study, we investigate the recurrent memory augmentation of pre-trained transformer models to extend input context length while linearly scaling compute. Our approach demonstrates the capability to store information in memory for sequences of up to an unprecedented two million tokens while maintaining high retrieval accuracy. Experiments with language modeling tasks show perplexity improvement as the number of processed input segments increases. These results underscore the effectiveness of our method, which has significant potential to enhance long-term dependency handling in natural language understanding and generation tasks, as well as enable large-scale context processing for memory-intensive applications.

Introduction

Self-attention gives Transformers context-aware representations, but its quadratic compute limits input length. The paper adapts pretrained Transformers with recurrent token memory and evaluates their ability to process much longer sequences.

Introduction

  • RMT adds special memory tokens to pretrained Transformers and trains the model to control memory operations and sequence representations without changing the Transformer.The study combines the approach with pretrained BERT and GPT-2 using full attention and full-precision operations.Source: Introduction · p. 1
  • Segment-wise curriculum learning is proposed to fine-tune most pretrained Transformer models for potentially unlimited sequences.The authors also introduce memory acquisition and retention tasks scalable to sequences of millions of tokens.Source: Introduction · p. 1
  • RMT stores facts for sequences up to two million tokens, which the authors describe as a record for the longest sequence task processed by a deep neural network.The authors also report no technical limitations preventing further scaling.Source: Introduction · p. 1
  • RMT offers linear inference scaling and constant memory, reducing the computational burden of processing long sequences.The paper compares its computational complexity with other Transformer models.Source: Introduction · p. 1

Related Work

Prior work combines neural networks and Transformers with forms of memory or reduced attention to handle long inputs. RMT is presented as a recurrent alternative that can use a common supported interface without changing the backbone architecture.

Related Work

  • Neural memory systems include attention-accessed storage in NTMs and Memory Networks and recurrent memory writing in DNC and Sparse DNC.These models are differentiable and trainable via backpropagation through time.Source: Related Work · pp. 1–2
  • Transformer recurrence processes long inputs as segments, with methods variously reusing hidden states, compressing memory, or updating dedicated memory modules.FeedBack Transformer also implements recurrence beyond the segment level.Source: Related Work · p. 2, Figure 1 · p. 2
  • Many recurrent methods require architectural modifications, whereas RMT can be built on models with a common supported interface.This contrast concerns ease of application to different pretrained models.Source: Related Work · p. 2
  • Reduced-attention methods limit attention distance or use global representations, but their memory requirements grow with input size during training and inference.Reported Longformer, Big Bird, and Long T5 lengths are below 33,000 tokens; CoLT5 handles up to 64,000, while CoLT5 and LongNet encounter GPU memory limits with longer inputs.Source: Related Work · p. 2

Recurrent Memory Transformer

RMT wraps a Transformer with token memory passed between sequentially processed input segments while leaving the backbone unchanged. With fixed segment length, it reduces inference compute relative to full-attention models as sequence length grows.

Recurrent Memory Transformer

  • At 2,048,000 tokens, the figure compares RMT’s inference FLOPs with full-attention OPT models, with RMT using ×29 fewer FLOPs for OPT-175B and ×295 fewer FLOPs than OPT-135M.The RMT segment length is fixed at 512 tokens; its scaling is linear, while full-attention OPT models are closer to quadratic scaling.Source: Figure 2 · p. 3
  • RMT uses global memory tokens for memory and recurrence, leaving the Transformer backbone unchanged and making the augmentation compatible with Transformer-based models.The approach is adapted as a plug-and-play wrapper for popular Transformers.Source: Recurrent Memory Transformer · pp. 2–3
  • The FLOP estimates cover OPT model configurations and forward passes, with estimates modified to account for RMT recurrence.The configurations include vocabulary size, layer count, hidden size, intermediate hidden size, and attention heads.Source: Recurrent Memory Transformer · p. 3
  • RMT divides the input into segments, computes full attention within each segment, and passes memory outputs from one segment to the next.For encoder-only models such as BERT, memory is prepended once per segment.Source: Recurrent Memory Transformer · pp. 2–3
  • For fixed segment length, RMT scales linearly with sequence length and requires fewer FLOPs than non-recurrent models beyond one segment.In this study, one segment is 512 tokens, and the reported maximum reduction is ×295.Source: Recurrent Memory Transformer · p. 3

Memorization Tasks

The synthetic datasets test whether models can retain facts amid irrelevant text and use them to answer questions at the sequence end. Tasks progress from storing a fact at the start to finding randomly placed facts and reasoning over two facts.

Memorization Tasks

  • The synthetic tasks combine answer-relevant facts with unrelated natural-language background and place the question at the sequence end.Each task is formulated as multi-class classification, with each class representing an answer option.Source: Memorization Tasks · p. 3, Figure 3 · p. 4
  • In Memorize, a fact starts the sequence and irrelevant text grows between it and the final question, testing extended storage across segments.The task asks the model to use the opening fact to answer the question at the end.Source: Memorization Tasks · p. 3, Figure 3 · p. 4
  • Detect and Memorize randomizes the fact’s position, requiring the model to distinguish it from irrelevant text, store it, and later answer the final question.Source: Memorization Tasks · pp. 3–4, Figure 3 · p. 4
  • Reasoning places two facts randomly in the sequence and asks a final question that requires both facts to answer correctly.The task uses the Two Argument Relation bAbI task; facts are generated from bAbI and background text comes from QuALITY questions.Source: Memorization Tasks · pp. 3–4, Figure 3 · p. 4

Learning Memory Operations

RMT is trained to carry information through memory across segments, using a curriculum that increases task length. It generalizes to longer synthetic tasks, including sequences of millions of tokens, while learned attention patterns correspond to memory operations.

Learning Memory Operations

  • Curriculum training begins with a one-segment task and increases sequence length by one segment after convergence, improving training stability.The practical BERT segment holds 499 tokens after reserving three special tokens and ten memory placeholders from its 512-token input.Source: Learning Memory Operations · p. 4
  • Models trained on more than five segments generalize well to longer input lengths across memorization, detection-and-memorization, and reasoning tasks.Figure 4 evaluates checkpoints trained on one to seven segments with memory size 10.Source: Figure 4 · p. 4
  • Most models perform well on shorter tasks, but a model trained on a single reasoning segment struggles when evaluated on longer sequences.The authors suggest it may stop expecting the question in the first segment once task size exceeds one segment.Source: Learning Memory Operations · pp. 4–5
  • Attention patterns show RMT detecting and writing facts to memory, preserving memory across an uninformative segment, appending another fact, and reading memory to answer.The model learns these recurrent operations without specific memory read/write modules, although the operations were not explicitly motivated by the task loss.Source: Figure 6 · p. 5, Learning Memory Operations · p. 5
  • After training on 5 or more segments, RMT generalizes nearly perfectly to tasks twice as long and holds up across validation tasks of up to 4,096 segments or 2,043,904 tokens.Detect & Memorize is the easiest task at extreme lengths, while Reasoning is the most complex.Source: Learning Memory Operations · p. 5, Figure 5 · p. 5

Natural and Formal Language Modeling

RMT carries information between segments in natural-language modeling and formal theorem proving. It improves reported perplexity in both settings, while the reported gains in lemma detection and use for proof generation are less notable.

Natural language modeling

  • The language-modeling experiment trains GPT-2 with two memory tokens on ArXiv documents from The Pile, processing histories and target segments sequentially.Documents are split into non-overlapping fixed-length segments, each prepended with its history.Source: Natural and Formal Language Modeling · pp. 5–6
  • RMT trained for the same number of steps as baseline GPT-2 has substantially lower perplexity, and multi-segment training improves tolerance to longer histories.Memory models trained without history suffer on long contexts, but improve after multi-segment training.Source: Natural and Formal Language Modeling · p. 6, Figure 7 · p. 6
  • RMT carries information from previous segments, supporting equally good predictions across token positions, while exact prior context yields larger loss gains at higher inference cost.The comparison concerns GPT-2 with context size 0 versus RMT memory and exact-context sizes 128 and 768.Source: Natural and Formal Language Modeling · p. 6, Figure 8 · p. 6

Formal language modeling

  • In formal theorem proving, each sequence contains a theorem, randomly ordered relevant and irrelevant lemmas, and a human-written proof, split into fixed-size segments.The experiment uses GPTNeo with 1.3B parameters, 10 memory tokens, and a segment size of 2028.Source: Natural and Formal Language Modeling · p. 6
  • RMT improves perplexity over the memory-less GPTNeo baseline, which is trained on two segments using the same number of tokens as RMT.Figure 9 reports perplexity for all sequence tokens and for the last segment.Source: Natural and Formal Language Modeling · p. 6, Figure 9 · p. 7
  • Training with four or more segments does not improve longer-sequence predictions, and gains from detecting and using lemmas in proof generation are less notable.The authors suggest more careful loss construction and data preparation may improve the latter result.Source: Natural and Formal Language Modeling · pp. 6–7

Conclusion

The work extends pretrained encoder-only and decoder-only Transformers with token-based recurrent memory and segment-level recurrence. RMT handles sequences exceeding one million tokens with linear computational scaling, while models pretrained on shorter sequences can adapt to longer tasks.

Conclusion

  • The work expands pretrained encoder-only and decoder-only Transformers through recurrent memory augmentation.The conclusion presents this as an extension of conventional Transformer scalability.Source: Conclusion · p. 7
  • Curriculum learning adapts models pretrained on shorter sequences to tasks involving significantly longer sequences.The paper argues that a model’s original input length does not necessarily restrict its potential capabilities.Source: Conclusion · p. 7
  • RMT handles sequences exceeding 1 million tokens while computational requirements scale linearly as task length increases.The work describes integrating token-based memory storage with segment-level recurrence.Source: Conclusion · p. 7
  • The authors present the work as a direction for handling longer sequences and improving Transformer adaptability.Source: Conclusion · p. 7

Limitations and Discussion

The paper identifies curriculum design, training cost, and evaluation scope as important boundaries. It also notes that full attention might retain a performance edge with unlimited resources and general-purpose information to remember.

Limitations and Discussion

  • The curriculum procedure substantially affects RMT generalization, making careful curriculum implementation necessary.Source: Limitations and Discussion · p. 7
  • 2M-token scaling is demonstrated only on specialized tasks, and established NLP benchmarks at those lengths are unavailable.Source: Limitations and Discussion · p. 7
  • BPTT training still requires significant computation; larger models cost more to train, and some tasks may need more than seven unrolled segments.The experiments found a maximum unroll of seven segments sufficient to demonstrate generalization to much longer sequences.Source: Limitations and Discussion · p. 7
  • With unlimited resources and general-purpose information to remember, full-attention models might retain a performance edge over RMT.RMT operates on compressed memory states rather than exact hidden states from the past.Source: Limitations and Discussion · p. 7

Efficiency of recurrence

In the reported comparison, RMT is faster and more memory-efficient than the standard Transformer, including with BPTT overhead. During inference, its compute grows linearly while memory use remains constant.

Efficiency of recurrence

  • RMT is faster and more memory-efficient than the standard Transformer, while the baseline runs out of memory at 8,000 tokens.The comparison used one Nvidia A100 80GB GPU, FP32, and a Hugging Face GPT-2 implementation without FlashAttention or other optimizations.Source: Efficiency of recurrence · p. 11
  • RMT remains faster and more memory-efficient even with the additional overhead of backpropagation through time.The BPTT overhead can be turned off when resources are limited.Source: Efficiency of recurrence · p. 11
  • During inference, RMT’s computational requirements scale linearly while memory requirements remain constant.Only one segment is kept in GPU memory at a time during inference.Source: Efficiency of recurrence · p. 11

Synthetic task generation

The discussed experiments construct synthetic memory tasks, evaluate curriculum training, and test combining RMT with parameter-efficient methods. The passages also identify a distractor-distribution limitation and report improved language-modeling perplexity when recurrence is added to parameter-efficient methods.

Synthetic task generation

  • The synthetic tasks use bAbI facts and questions, with facts placed at the beginning or in a random segment; distractor text comes from a different distribution, making question identification trivial.The memorization and Detect & Memorize datasets use bAbI’s single-supporting-fact subset; the reasoning task uses its two-argument-relations subset.Source: Synthetic task generation · p. 11

Impact of curriculum

  • Curriculum enables task learning and length generalization, while mixing examples from earlier segment counts improves generalization across lengths.Without curriculum, training directly on the maximum of four segments fails to learn the task or extrapolate; length mixing also improves performance at smaller segment counts and on unseen lengths.Source: Impact of curriculum · p. 12, Figure 11 · p. 12, Figure 12 · p. 12

Usage of Parameter-Efficient Methods

  • Adding RMT recurrence to LoRA or Parallel Adapter significantly improves Pythia-70m perplexity on Arxiv language modeling compared with those parameter-efficient methods alone.Table 1 reports the comparison for the Arxiv dataset and Pythia-70m; the backbone architecture remains unchanged.Source: Usage of Parameter-Efficient Methods · p. 12, Table 1 · p. 13

Related papers

Loading 2304.11062v2…