Abstract <p>A fundamental limitation of Transformer-based models is their quadratic computational complexity with respect to input length, which limits their applicability to long-context tasks. Recurrent Memory Transformer (RMT) addresses this by introducing a memory mechanism that enables segment-wise recurrent processing. However, RMT relies on a multi-stage training curriculum that increases computational costs and complexity during fine-tuning. In this work, we propose the Recurrent Memory Transformer with a Memory Stream (RMT-MS), a novel architecture with layer-wise memory states and horizontal memory connections across segments. These mechanisms increase memory capacity and improve information flow, reducing the need for curriculum learning. We evaluate RMT-MS alongside RMT and ARMT on three long-context tasks: associative retrieval, BABILong QA1, and QA3. Our experiments show that RMT-MS achieves strong performance in single-stage training, matching curriculum-trained baselines on simpler tasks, and narrowing the gap on more complex ones. These results highlight the potential of RMT-MS for efficient long-context modeling without costly training schedules.</p>

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Memory Stream: Enhancing Information Flow in Recurrent Memory Transformers for Efficient Long-Context Training

  • M. Kairov,
  • A. Bulatov,
  • Yu. Kuratov

摘要

Abstract

A fundamental limitation of Transformer-based models is their quadratic computational complexity with respect to input length, which limits their applicability to long-context tasks. Recurrent Memory Transformer (RMT) addresses this by introducing a memory mechanism that enables segment-wise recurrent processing. However, RMT relies on a multi-stage training curriculum that increases computational costs and complexity during fine-tuning. In this work, we propose the Recurrent Memory Transformer with a Memory Stream (RMT-MS), a novel architecture with layer-wise memory states and horizontal memory connections across segments. These mechanisms increase memory capacity and improve information flow, reducing the need for curriculum learning. We evaluate RMT-MS alongside RMT and ARMT on three long-context tasks: associative retrieval, BABILong QA1, and QA3. Our experiments show that RMT-MS achieves strong performance in single-stage training, matching curriculum-trained baselines on simpler tasks, and narrowing the gap on more complex ones. These results highlight the potential of RMT-MS for efficient long-context modeling without costly training schedules.