Maglev: Sliding Recurrent Memory
Maglev combines sliding-window attention with fixed-size recurrent memory via a parallel-trained predictor, boosting long-sequence modeling by 1-2% accuracy.
Key Findings
Methodology
Maglev employs a dual-model setup: a causal prefiller Q, utilizing full attention to generate target memories m′t, and a decoder P, which uses only sliding-window attention with recurrent K/V injection to produce current memories mt. During training, Q generates all target memories in parallel, supervised by a consistency loss aligning mt with m′t. P learns to predict tokens and produce memories, enabling inference solely with P. Parameter sharing between Q and P reduces model size while maintaining performance. This design balances long-range information capture with training efficiency, leveraging the parallelizable nature of the prefiller and the fixed memory size of the decoder.
Key Results
- On 43.52B tokens, Maglev reduces validation BPB from 0.7413 to 0.7251 and improves downstream accuracy from 54.1% to 56.4%, outperforming sliding-window and latent recurrent baselines. Parameter sharing maintains most gains with fewer parameters.
- Ablation studies show that the consistency loss weight λ influences performance, with moderate λ balancing target accuracy and model flexibility.
- The model excels in tasks requiring long context understanding, demonstrating effective long-term memory retention without increasing inference complexity.
Significance
Maglev addresses the core challenge of training nonlinear recurrent Transformers efficiently for long sequences. By decoupling target memory generation from inference, it enables parallel training while maintaining fixed-size, persistent memory during inference. This innovation opens new avenues for scalable long-context language models, impacting applications such as document understanding, dialogue systems, and knowledge bases, where long-range dependencies are critical. The approach bridges the gap between efficiency and expressive power, offering a practical solution for real-world long-sequence tasks.
Technical Contribution
The paper introduces a novel architecture combining a parallel-trained causal prefiller with a sliding-window decoder that injects shifted memories via K/V pathways. The consistency loss aligns decoder memories with target memories, enabling effective nonlinear recurrence without sequential unrolling. Parameter sharing reduces model size, and the design maintains bounded inference costs. This framework extends Transformer capabilities, providing a scalable, efficient method for long-sequence modeling with fixed memory constraints.
Novelty
This work is the first to utilize a parallel causal prefiller to generate target memories for training a fixed-memory recurrent Transformer. Unlike prior linear attention or state-space models, Maglev introduces a nonlinear, parameter-efficient recurrence mechanism that is trained via consistency supervision, enabling effective long-range modeling without sacrificing training parallelism or inference efficiency.
Limitations
- The success heavily depends on the quality of the target memories generated by Q; poor predictions can impair overall performance.
- Parameter sharing may limit the model’s capacity to handle extremely complex tasks or very long sequences.
- The current design may struggle with tasks requiring highly dynamic or context-dependent memory updates, necessitating further adaptation.
Future Work
Future research will explore optimizing the capacity and structure of the prefiller Q, including pretraining strategies and architectural variations. Investigations into alternative memory injection methods, such as residual streams or cross-attention, could further enhance memory expressivity. Extending the framework to multimodal data and multi-task learning, as well as scaling to larger models and datasets, are promising directions for advancing long-sequence modeling capabilities.
AI Executive Summary
Long-sequence modeling remains a fundamental challenge in natural language processing. Traditional Transformer models, while powerful, face computational bottlenecks as sequence length grows, due to the quadratic complexity of full attention. Sliding-window attention reduces this cost but sacrifices information outside the window, limiting long-term memory. Recurrent models like LSTMs offer persistent memory but struggle with parallel training and capacity. Recent efforts, including linear attention and state-space models, improve efficiency but often lack nonlinear expressivity or flexibility.
Maglev introduces a novel architecture that combines the strengths of these approaches. It features a parallel-trained causal prefiller Q, which generates target memories m′t using full attention, and a decoder P that employs sliding-window attention with recurrent K/V injection. During training, Q produces all target memories in parallel, supervised by a consistency loss that aligns decoder memories mt with target memories m′t. At inference, Q is discarded, and P runs recurrently, maintaining fixed-size memory windows. This design ensures efficient long-range modeling without increasing inference complexity.
Experimental results on 43.52 billion tokens demonstrate that Maglev outperforms baseline sliding-window and latent recurrent Transformers, reducing validation BPB from 0.7413 to 0.7251 and increasing downstream task accuracy from 54.1% to 56.4%. The parameter sharing strategy further reduces model size while preserving performance. These findings confirm that the combination of parallel supervision and fixed-size recurrent memory can significantly enhance long-context language modeling.
Overall, Maglev offers a scalable, efficient solution for long-sequence tasks, bridging the gap between expressivity and training efficiency. Its innovative framework paves the way for future research into nonlinear recurrence, memory injection techniques, and broader applications in NLP and beyond. Despite current limitations, such as dependence on the quality of target memory generation, the approach sets a promising direction for next-generation long-range models.
Deep Analysis
Background
长序列建模一直是深度学习中的核心难题。早期模型如LSTM和GRU通过递归机制压缩历史信息，但受限于线性状态更新，难以捕获复杂的长距离依赖。Transformer引入全注意力机制，极大提升表达能力，但在超长文本中计算成本呈指数增长。为解决这一问题，出现多种变体：Longformer、Reformer等采用局部注意力或压缩记忆，Transformer-XL引入片段复用机制，增强长距离依赖能力。尽管如此，这些方法在保持效率的同时，仍面临信息遗失或建模能力不足的挑战。线性注意力和状态空间模型如S4、Hyena通过线性变换实现高效递归，极大降低复杂度，但缺乏非线性表达能力。近年来，结合递归与注意力的混合模型逐渐兴起，试图在保持长距离信息的同时提升训练效率。
Core Problem
现有模型在长序列建模中面临两大瓶颈：一是递归模型难以高效训练，因其序列依赖性强，训练过程需逐步展开；二是全注意力机制在超长文本中计算成本过高，限制了模型的扩展性。滑动窗口注意力虽能降低复杂度，但信息传递受限，难以实现持久记忆，影响推理性能。线性模型虽高效，但在非线性表达方面不足，难以应对复杂推理任务。如何在保证训练效率的同时，增强模型的长距离记忆和非线性表达能力，成为研究的核心难题。
Innovation
Maglev的核心创新在于引入偏预测器Q，利用全注意力在训练时生成目标记忆，提供并行监督，打破递归训练的序列依赖。解码器P采用滑动窗口注意力结合递归K/V注入，保持推理时的固定记忆容量。通过一致性损失对齐两者，确保模型在推理中依赖自身记忆。参数共享设计减少模型复杂度，提升训练效率。这一机制融合了长距离信息捕获与高效训练的优势，为非线性递归Transformer提供了新思路。
Methodology
- �� Q模型采用全注意力机制，输入完整序列，输出目标记忆m′t，训练时利用交叉熵和一致性损失优化。
- �� P模型使用滑动窗口注意力，结合递归K/V注入，将偏预测器生成的目标记忆融入每个时间步的注意力计算中。
- �� 训练过程中，Q在前一轮生成所有目标记忆，P在下一轮学习预测下一词和重建记忆，二者通过损失函数对齐。
- �� 在推理阶段，Q被剔除，P仅用自身记忆递归预测，保持固定大小的记忆窗口。
- �� 参数共享策略在不同模型变体中应用，减小模型存储需求，提升训练效率。
Experiments
采用43.52B Token的训练规模，基于NanoChat架构（20层，1280维，10头，最大序列长度2048），对比滑动窗口、全注意力和潜在递归Transformer。评估指标包括BPB、困惑度、任务准确率。实验还包括不同λ值的损失权重、参数共享与否的对比，以及不同任务（如文本理解、推理）的性能表现。通过消融验证偏预测器容量和一致性损失的影响，确保模型在长文本任务中的优越性。
Results
Maglev在43.52B训练Token后，BPB由0.7413降至0.7251，平均下游任务准确率从54.1%提升到56.4%。参数共享模型几乎保持性能，验证偏预测器目标的有效性。消融显示，适当调整λ值能进一步优化模型表现，表明偏预测器容量和解码器能力的平衡关键。模型在长文本推理和知识存储任务中表现优异，验证了偏预测器生成目标记忆的实用性。
Applications
Maglev适用于长文本理解、对话系统、知识库问答等场景，能在保持训练效率的同时，提供持久的长距离记忆能力。其固定大小的记忆窗口确保推理成本可控，适合大规模部署。未来可结合预训练模型进行微调，扩展到多模态和多任务环境，推动实际应用落地。
Limitations & Outlook
模型对偏预测器的依赖较大，偏预测器性能不足会影响整体效果。参数共享可能限制模型容量，难以应对极端长序列。偏预测器生成目标的准确性在复杂任务中仍有提升空间。未来需优化偏预测器结构，提升其泛化能力，减少训练成本。
Plain Language Accessible to non-experts
想象你在厨房里做饭。传统的厨师会记住所有食材和步骤，用心烹饪，但每次都要从头开始，效率低。现在，有个智能助手会提前准备好所有食材的清单，告诉你每一步该做什么，还能记住之前的菜谱。你只需要跟着助手的指示做菜，自己不用记太多。Maglev就像这个助手，它用一个“偏预测器”提前准备好长长的记忆清单，帮助“厨师”——模型——在做长篇文章或对话时，记住很多信息。最后，助手可以被“丢掉”，模型自己继续工作，就像厨师凭借助手提前准备的食材，独立完成大餐。这种方式既快又省力，还能做出复杂的菜肴（长文本）。
ELI14 Explained like you're 14
想象你在玩一个超级长的游戏，每次你都要记住之前发生的事情，但游戏太长，记忆变得很难。以前的游戏角色（模型）要一边玩一边记，特别慢。现在，有个聪明的朋友（偏预测器）会提前告诉你所有重要的线索，你只需要跟着线索走就行了。这个朋友帮你准备好信息，你不用每次都自己回忆。等到游戏结束，你可以把这个朋友“送走”，自己用之前学到的线索继续玩。这样既快又能记住很多事情，玩得更顺畅。Maglev就像这个聪明的朋友，让游戏变得更简单、更快，也能记住更长的故事！
Abstract
We introduce \ours{}, a recurrent Transformer architecture with fixed-size memory that generalizes sliding-window attention while remaining parallelizable during training. \ours{} consists of two coupled models: a prefiller $Q$, which leverages full attention\footnote{In practice, we use interleaved full and sliding-window attention for $Q$, as this yields stronger performance. The essential requirement is that $Q$ be more expressive than $P$, with access to the full history.} to produce memory targets $m'_t$, and a decoder $P$, which uses only sliding-window attention and recurrent K/V injection to produce decoder memories $m_t$ for next-token prediction. We train \ours{} with a memory consistency loss that aligns $m_t$ with $m'_t$, allowing inference to use $P$ alone. Empirically, \ours{} improves validation loss and downstream pretraining benchmarks over sliding-window and latent recurrent transformer baselines. Moreover, sharing parameters between $P$ and $Q$ reduces parameter memory while preserving most of the gains.