Scaling Transformer to 1M tokens and beyond with RMT

• URL: http://arxiv.org/abs/2304.11062v2
• Date: Tue, 6 Feb 2024 10:16:54 GMT
• Title: Scaling Transformer to 1M tokens and beyond with RMT
• Authors: Aydar Bulatov, Yuri Kuratov, Yermek Kapushev, Mikhail S. Burtsev
• Abstract summary: 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.
• Score: 5.60052250541419
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Taipan: Efficient and Expressive State Space Language Models with Selective Attention [100.16383527459429]
  Long-context language modeling is a significant challenge in Natural Language Processing (NLP) Recent State Space Models (SSMs) such as Mamba offer alternatives with constant memory usage, but they underperform in tasks requiring extensive in-context retrieval. We introduce Taipan, a novel hybrid architecture that combines Mamba-2 with Selective Attention Layers (SALs) Our experiments demonstrate Taipan's superior performance across various scales and tasks, offering a promising solution for efficient long-context language modeling.
  arXiv  Detail & Related papers  (2024-10-24T09:25:37Z)
• UIO-LLMs: Unbiased Incremental Optimization for Long-Context LLMs [111.12010207132204]
  UIO-LLMs is an incremental optimization approach for memory-enhanced transformers under long-context settings. We refine the training process using the Truncated Backpropagation Through Time (TBPTT) algorithm. UIO-LLMs successfully handle long context, such as extending the context window of Llama2-7b-chat from 4K to 100K tokens with minimal 2% additional parameters.
  arXiv  Detail & Related papers  (2024-06-26T08:44:36Z)
• Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers [4.674454841332859]
  Transformer-based models have emerged as one of the most widely used architectures for natural language processing. These huge models are memory hungry and incur significant inference latency even on cutting edge AI-accelerators. We propose LeanAttention, a scalable technique of computing self-attention for the token-generation phase.
  arXiv  Detail & Related papers  (2024-05-17T00:52:39Z)
• Ring Attention with Blockwise Transformers for Near-Infinite Context [88.61687950039662]
  We present a novel approach, Ring Attention with Blockwise Transformers (Ring Attention), which leverages blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices. Our approach enables training and inference of sequences that are up to device count times longer than those achievable by prior memory-efficient Transformers.
  arXiv  Detail & Related papers  (2023-10-03T08:44:50Z)
• Chunk, Align, Select: A Simple Long-sequence Processing Method for Transformers [24.109312575970456]
  We propose a simple framework to enable the offthe-shelf pre-trained transformers to process much longer sequences. Our method divides each long-sequence input into a batch of chunks, then aligns the interchunk information during the encoding steps. We learn an effective hidden selection policy, which regards the decoders of transformers as environments.
  arXiv  Detail & Related papers  (2023-08-25T05:52:05Z)
• Blockwise Parallel Transformer for Large Context Models [70.97386897478238]
  Blockwise Parallel Transformer (BPT) is a blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods.
  arXiv  Detail & Related papers  (2023-05-30T19:25:51Z)
• Confident Adaptive Language Modeling [95.45272377648773]
  CALM is a framework for dynamically allocating different amounts of compute per input and generation timestep. We demonstrate the efficacy of our framework in reducing compute -- potential speedup of up to $times 3$ -- while provably maintaining high performance.
  arXiv  Detail & Related papers  (2022-07-14T17:00:19Z)
• Recurrent Memory Transformer [0.3529736140137003]
  We study a memory-augmented segment-level recurrent Transformer (Recurrent Memory Transformer) We implement a memory mechanism with no changes to Transformer model by adding special memory tokens to the input or output sequence. Our model performs on par with the Transformer-XL on language modeling for smaller memory sizes and outperforms it for tasks that require longer sequence processing.
  arXiv  Detail & Related papers  (2022-07-14T13:00:22Z)
• Linearizing Transformer with Key-Value Memory Bank [54.83663647680612]
  We propose MemSizer, an approach to project the source sequence into lower dimension representation. MemSizer not only achieves the same linear time complexity but also enjoys efficient recurrent-style autoregressive generation. We demonstrate that MemSizer provides an improved tradeoff between efficiency and accuracy over the vanilla transformer.
  arXiv  Detail & Related papers  (2022-03-23T18:10:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.