Scaling Transformer to 1M tokens and beyond with RMT

Cool. The idea of augmenting a transformer with an RNN isn't new, but I'll take a look.

See also the threads here, mentioning other methods for scaling up context length up to the millions of tokens:

https://news.ycombinator.com/item?id=35502187

Another roadblock to AGI starting to fall. WOW.

This paper presents a novel approach to incorporate memory into a transformer. It does not demonstrate that this approach works in a useful manner. While the approach is interesting, I’m skeptical that the RNN has enough capacity to encode the memory in it's output. I would have liked to see more detail on the synthetic benchmark they used. The memory component may be learning the benchmark rather than a generalized feature.

Abstract

This technical report presents the application of a recurrent memory to extend the context length of BERT, one of the most effective Transformer-based models in natural language processing. By leveraging the Recurrent Memory Transformer architecture, we have successfully increased the model's effective context length to an unprecedented two million tokens, while maintaining high memory retrieval accuracy. Our method allows for the storage and processing of both local and global information and enables information flow between segments of the input sequence through the use of recurrence. Our experiments demonstrate the effectiveness of our approach, which holds 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

The Transformer model (Vaswani et al., 2017) has been widely adopted and used in various research areas and industrial applications. The most important issue of the model is quadratic complexity of attention operation, that makes large models increasingly difficult to apply to longer inputs.

This report we show that by using simple token-based memory mechanism introduced in (Bulatov et al., 2022) can be combined with pretrained transformer models like BERT (Devlin et al., 2019) with full attention and full precision operations can be applied to sequences longer than 1 million tokens using a single Nvidia GTX 1080Ti GPU.

Contributions

1. We enhance BERT by incorporating token-based memory storage and segment-level recurrence with recurrent memory (RMT).

2. We demonstrate that the memory-augmented BERT can be trained to tackle tasks on sequences with lengths up to seven times its originally designed input length (512 tokens).

3. We discovered the trained RMT’s capacity to successfully extrapolate to tasks of varying lengths, including those exceeding 1 million tokens with linear scaling of computations required.

4. Through attention pattern analysis, we found the operations RMT employs with memory, enabling its success in handling exceptionally long sequences.

Discussion

The problem of long inputs in Transformers has been extensively researched since the popularization of this architecture. In this work, we demonstrate that applying Transformers to long texts does not necessarily require large amounts of memory. By employing a recurrent approach and memory, the quadratic complexity can be reduced to linear. Furthermore, models trained on sufficiently large inputs can extrapolate their abilities to texts orders of magnitude longer.

Synthetic tasks explored in this study serve as the first milestone for enabling RMT to generalize to tasks with unseen properties, including language modelling. In our future work, we aim to tailor the recurrent memory approach to the most commonly used Transformers to improve their effective context size.