Maglev: sliding recurrent memory

TL;DR Maglev 用预填充器 Q 与解码器 P 耦合并施加记忆一致性损失，让固定大小记忆的递归 Transformer 在训练可并行的同时，推理仅用 P 即超越滑动窗口和潜在递归基线。

问题

问题背景

Transformer 在长上下文建模中需要权衡记忆容量、推理成本与训练并行性。

现有方法局限

• 滑动窗口注意力 将注意力限制在固定窗口内，虽然 KV 缓存和计算代价可控，但直接丢弃窗口外信息，缺乏跨长距离的持久记忆。
• 经典循环模型（如 LSTM）使用非线性 token-wise 状态更新，但隐藏状态维度较小，容量有限，且串行传播限制序列并行。
• Memory Transformer 只在 segment/block 边界更新记忆，粒度粗，无法在每个 token 级别注入或读取信息。
• 线性注意力 / 状态空间模型 虽然实现 token 级并行循环，但采用结构化线性或仿射更新，表达力受限于固定形式。

为什么这个问题难 / 重要

核心挑战在于 非线性 token-wise 记忆更新 与 并行训练 存在内在冲突。全注意力可以提供强表达力但推理 KV 缓存随上下文增长；纯循环状态固定但难以并行且容量有限。业界对长上下文 LLM 的需求持续上升，希望在固定内存预算下保留远距离检索与复制能力，同时预训练不能牺牲并行效率。因此需要一种既能 token 级更新、又可在训练时并行监督的架构。

行业类比

类似实时代码助手或长会议纪要系统：内存固定，不能缓存全部历史，但必须持续记住用户早期关键意图并逐 token 更新当前决策。

核心洞察

• Maglev 的核心洞察是将递归记忆的训练与推理解耦：通过一个全注意力 prefiller 提供理想的内存表示，再用一致性损失蒸馏到仅用滑动窗口注意力的 decoder，从而在保持并行训练的同时实现 token-wise 递归更新。这与 Mamba、RWKV 等线性递归模型不同，后者为了并行化而强制使用结构化线性更新，限制了记忆的非线性表达能力；也与分段记忆 Transformer 不同，Maglev 的递归是 token 级别的，不依赖块边界，更细粒度。
• Maglev 将固定大小内存视为滑动窗口注意力的泛化：decoder 在每个位置通过递归注入过去 K/V 对，使模型能访问窗口外的压缩历史，而训练时由 prefiller 指导如何有效压缩。这不同于 Longformer 等仅扩大窗口的方法，也不同于 Transformer-XL 的分段缓存，Maglev 的内存是固定大小、可学习的，并且更新公式通过注意力机制实现，能捕捉非线性依赖，同时参数共享进一步降低了成本。

方法

输入与双模型结构

Maglev 接收 token 序列作为输入，并维护固定大小的循环记忆。架构由两个耦合模型组成：

• Prefiller Q：表达力更强，使用交错的全注意力与滑窗注意力，可访问完整历史，为每个位置生成记忆目标 m'_t。
• Decoder P：仅使用滑窗注意力和循环 K/V 注入，根据局部窗口与上一时刻记忆 m_{t-1} 更新并输出当前记忆 m_t，用于下一 token 预测。

训练与推理流程

训练时同时运行 Q 和 P：Q 借助全局信息产生高质量记忆目标，P 则通过 token 级循环更新产生自身记忆，二者通过记忆一致性损失对齐。该损失强迫 P 的 m_t 逼近 m'_t，因此在推理时可单独使用 P。可选地，P 与 Q 共享参数，减少额外参数量而保留大部分收益。推理阶段只运行 P，计算成本与滑窗注意力相当，同时保留固定大小的循环记忆，无需访问完整历史。

与同类方法的差异

与传统滑窗注意力或分段记忆 Transformer 不同，Maglev 通过 token 级循环记忆和预填充器并行监督，在不牺牲训练并行性的前提下获得持久记忆能力。

实验

实验设计

论文提出 Maglev 训练框架，包含两个耦合模型：prefiller Q（使用全注意力或交错全注意力/滑动窗口，生成高质量记忆目标 m'_t）和 decoder P（仅用滑动窗口注意力与循环 K/V 注入，生成用于下一 token 预测的记忆 m_t）。训练通过 memory consistency loss 对齐两者，推理阶段仅使用 P，降低推理成本。

对比基线包括：

• 纯 sliding-window attention 模型
• latent recurrent transformer 模型

评估指标：验证损失与下游预训练基准（具体数据集未在摘要中披露）。

关键发现

• 验证损失与下游基准：Maglev 相比滑动窗口与潜在循环 Transformer 基线取得提升，但摘要未提供具体数值。
• 参数共享：在 P 与 Q 之间共享参数可显著降低参数内存，同时保留大部分性能增益。
• 训练并行化：Maglev 在训练时保持并行性，避免了逐 token 顺序传播的瓶颈。

基线对比解读

• 相较 sliding-window attention：Maglev 引入固定大小的循环记忆，弥补滑动窗口丢弃远距离信息的缺陷，且不牺牲训练的并行性。
• 相较 latent recurrent Transformer：Maglev 采用 token-wise 记忆更新与并行监督，避免了段级更新造成的人工粒度或线性/仿射更新的表达力限制。

工程启示：该架构实现了训练（并行、高表达力）与推理（仅用 P，低 KV-cache 与注意力成本）的解耦，适合长上下文场景下的高效部署。

行业影响

落地场景

长文档理解与 RAG：金融研报、法律合同、医疗病历等超长文本场景，Maglev 的固定大小记忆可替代全注意力或粗暴截断，在保持全局信息的同时将推理时 KV-cache 控制为常数级，适合需要精准引用与跨段落推理的产品。

对话式客服与智能助理：电商平台或企业服务中，客服需维护跨多轮会话的长期上下文。Maglev 推理时仅用滑动窗口 decoder，不随对话长度累积显存，可支持更长会话与更低延迟。

商业价值

• 降本：推理阶段 KV-cache 不再线性增长，直接降低长上下文服务的 GPU 显存与计算成本；训练可并行，缩短模型迭代周期。
• 体验提升：通过 prefiller 蒸馏出的全局记忆，减少长文本问答中的信息遗忘与幻觉，提升关键任务准确率。

与现有产品/工作流接口

Maglev 可作为 Transformer 中 attention 层的替代模块或 hybrid 组件，集成到现有训练框架（如 PyTorch / Hugging Face）。训练时新增 memory consistency loss，无需修改优化器与数据管线。推理侧只需部署 decoder P，兼容 vLLM 等主流 serving 引擎的自定义 attention/kernel 接口。对于已有长上下文 LLM 产品，可通过轻量 finetune 将部分层替换为 Maglev，在保持模型主体结构不变的情况下降低推理成本。

局限

• **固定容量限制**：记忆大小固定，可能无法容纳足够的长程依赖，尤其是需要精确复制或检索远处 token 的任务（如代码补全、长文档问答）。与全注意力相比，在需要全局信息的场景下仍可能存在性能差距。论文未深入分析不同记忆容量对下游任务的影响，也未讨论动态扩展或分层记忆机制，这限制了其在超长上下文任务中的适用性。
• **训练成本与超参数**：训练需同时运行 prefiller Q 和 decoder P，其中 Q 使用全注意力或交错注意力，导致训练计算量与内存占用显著高于纯滑动窗口模型。记忆一致性损失引入额外超参数（如损失权重、记忆维度），可能增加调参难度。此外，参数共享虽然减少模型大小，但可能限制 Q 与 P 的差异化表达能力，影响记忆目标的质量。
• **与高效序列模型的对比**：相较线性注意力/状态空间模型（如 Mamba、RWKV），Maglev 的推理仍需维护滑动窗口 KV 缓存，且循环状态更新为非线性操作，可能不利于极致的内存和计算优化。实验仅覆盖预训练基准，缺乏长上下文下游任务（如 LongBench、RULER）的评测，实际部署收益尚不明确，在真实应用中的优势有待进一步验证。