Memory Stream: Enhancing Information Flow in Recurrent Memory Transformers for Efficient Long-Context Training
摘要
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.