PI-Mem：打破串行记忆覆盖，3.6M超长推理提速6.1倍且涨点7.8
PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory
在大模型迈向百万级上下文处理能力的进程中，推理效率与长程线索检索之间的矛盾始终未能得到根本解决。面对超长书籍、大型代码仓库或是跨越多篇文献的复杂多跳问答，盲目扩大模型的原生上下文窗口不仅面临极高的显存开销与注意力计算开支，还会陷入“大海捞针”般的线索遗忘问题。
ArXiv URL：https://arxiv.org/abs/2608.03048v1
为了绕开原生窗口的物理限制，以 MemAgent 为代表的“循环记忆”（Recurrent Memory）工作流应运而生。这类方法将数百万 Token 的长文档切分为数个分块（Chunks），维护一个固定长度的文本记忆，并按顺序逐块读取、覆盖式更新。然而，这种看似线性的精简设计却隐藏着致命缺陷：一是严格的串行依赖，处理数以百计的分块必须串行排队，GPU 硬件的并行吞吐能力被完全锁死；二是前后信息的严重不对称，早期分块中潜藏的蛛丝马迹在尚未被确认价值前，极易被后序涌入的大量无关内容反复洗刷，最终导致记忆被冲淡或覆盖。
由哈工大、上海人工智能实验室、清华大学、上海交通大学等团队联合提出的 PI-Mem（Parallel-Iterative Memory），彻底打破了这一“保精度就必须牺牲吞吐、求轻量就必然遗忘”的权衡困局。PI-Mem 抛弃了传统的时间步串行推进思路，转而采用以“全局记忆为锚点、全分块并行读取、跨轮次迭代提炼”的新范式。在扩展到 360 万（3.6M）Token 上下文的 HotpotQA 极长推理评测中，PI-Mem 相比经典的循环记忆基线取得了最高 +7.81 分的绝对准确率跃升，并换来了高达 6.1 倍的端到端推理加速。这一成果表明，依靠更具智慧的工作流拓扑与强化学习对齐，处理数百万量级的上下文完全可以做到又快又准。
串行覆盖与吞吐瓶颈：长文本记忆为何需要范式转换？
长文本理解的难点从来不是“读完所有字”，而是“串联起分布在不同角落的证据碎片”。在传统的循环记忆机制中，系统以流水线的方式推进：模型先读取分块 1，生成更新后的记忆；接着带着新记忆去读取分块 2，再次覆写更新。
这种流式覆盖更新在面临真实复杂的多跳（Multi-hop）推理时极为脆弱。许多跨文档问答中，证据 A 出现在文本前 5% 的位置，证据 B 却落在后 90% 的位置，单独看证据 A 甚至无法判断它是否与用户的问题直接相关。在传统的循环记忆机制下，模型在读到证据 A 时只能给予极低的权重甚至忽略；而在流经中间数百万 Token 的干扰段落时，原本微弱的早期证据早已在数十次重写中被覆盖殆尽。等到系统终于读到证据 B 时，能够与它产生“化学反应”的前置信息早已荡然无存。
更为严重的是工程落地层面的吞吐灾难。由于后一个分块的推理必须严格等待前一个分块完成记忆覆写，整个推理过程形成了一条极其漫长、无法被打断的串行关键路径（Serial Critical Path）。现代 GPU 针对高并发批处理（Batching）设计，但在串行处理超长序列时，硬件利用率极度受限，推理总延迟随着上下文长度呈线性甚至更糟的阶梯式递增。
直接进行原生全注意力推理会导致显存和 KV Cache 崩溃，而传统的循环记忆又陷入延迟高与易覆盖的泥潭。PI-Mem 的设计初衷，正是要把这种长程文本的依赖关系，从传统的“空间切片、时间排队”转变为“空间并发、多轮提炼”。
并行迭代的机制重塑：Read-Select-Merge 闭环
PI-Mem 的核心逻辑是，通过将“跨分块关联”的解决机制从单次串行扫描转移到有限轮次的宏观迭代中，实现单轮内部的完全并行化。它把长输入切分为等长分块后，不再逐个顺序穿梭，而是将每一轮记忆演进解构为三个标准化步骤：并行读取（Read）、局部筛选（Select）与紧凑融合（Merge）。
在每一个更新轮次 $k$ 中，系统首先维护着一个在所有分块间共享的全局记忆状态 $m^{(k-1)}$（初始为空）。随后的三步操作形成一个高内聚的闭环：
并行读取（Parallel Read）：模型将用户提问 $q$ 与当前共享的全局记忆 $m^{(k-1)}$ 广播到每一个分块 $c_i$ 上。所有分块的读取操作（ReadCall）彼此独立，可以在 GPU 集群中以高度并发的批处理方式瞬间触发。由于每个分块在读取时都参考了相同的全局记忆上下文，模型在评估分块局部内容时拥有完全一致的全局视角，消除了先后顺序造成的线索偏倚，也绕开了等待前置分块处理的延迟屏障。
严格筛选（Select）：并发读取后，各分块会产出局部的观测提炼。PI-Mem 并不是一股脑将所有分块的产出全部塞进下一轮，而是设计了自省校验机制。模型在生成观测的同时输出一个结构化判断标记（例如通过输出特定标签确认该分块是否包含增量或互补信息）。只有通过验证的有效证据才会被收入当前轮次的候选集 $O^{(k)}$ 中，这构成了抵御无关上下文污染的第一道防火墙。
紧凑融合（Merge）：获取各并行分块涌现出的有效证据后，系统调用合并算子（MergeCall），将提问 $q$、上一轮记忆 $m^{(k-1)}$ 以及本轮收集的全部有效观察集合 $O^{(k)}$ 统一交付给模型进行蒸馏。模型的工作不是机械拼接字符，而是将分散在各个分块中的新证据与既有记忆进行去重、因果关联与信息压缩，最终凝练成一段保持紧凑的全新全局记忆 $m^{(k)}$。这种设计确保了哪怕文本被切分为几百个分块，最终沉淀出的记忆体积始终处于严格受控的边界之内，避免了中间文本爆炸。
单轮并行的“读取-筛选-融合”解决的是并发和局部噪声抑制问题，而多跳推理的深层破解则依赖于跨轮次迭代（Iterative Refinement）。在第一轮中，初始记忆为空，模型只能在各分块中检索出与问题直接匹配的浅层事实；当这些事实被沉淀到 $m^{(1)}$ 并开启第二轮时，所有分块再次进行全量并行读取。此时，模型拥有了前一轮已发掘出的关键线索，原本孤立不可见的深层关联在第二轮读取中迅速显形，真正实现了跨越数百万 Token 的跨块信息交织。
轨迹级强化学习：引入轮次效率奖励打破冗余
将这样一套由多次大模型调用组成的复杂工作流落到实处，若只依赖预训练模型的自身指令跟随能力，往往会出现执行动作走样：要么模型缺乏筛选纪律，将大量噪声写入记忆；要么陷入无限迭代的死循环，为了追求所谓“更全的证据”反复消耗计算资源。
为了给 PI-Mem 注入自我节制与精准筛选的能力，作者团队采用了基于 GRPO（Group Relative Policy Optimization）的轨迹级（Trajectory-level）强化学习进行端到端优化。
在强化学习的建模中，一个完整回答的过程构成了一条长轨迹 $\tau_i$。该轨迹不仅包含最终给出答案的回答样本，还完整串联了中间所有分块的读取样本（Read calls）与记忆融合样本（Merge calls）。为了全面校准这一复合工作流，奖励函数的设计包含了双重目标：不仅要保证最终答案的正确性，还要引导模型尽可能以最少的迭代轮次完成任务。最终的轨迹级复合奖励为：
\[R_i = r_{\mathrm{acc}}(a_i) + \lambda_{\mathrm{turn}} r_{\mathrm{turn}}(\tau_i)\]其中，$r_{\mathrm{acc}}(a_i)$ 为对最终答案正确性赋予的奖励；而 $r_{\mathrm{turn}}(\tau_i) = \frac{K - k_i}{K - 1}$ 则是关键的轮次效率奖励（Turn-Efficiency Reward），其中 $K$ 是系统允许的最大轮次上限，$k_i$ 是当前轨迹实际消耗的轮次。这个奖励项赋予提前终结迭代的动作以显式的数学偏好——只要当前积累的全局记忆已经足以支撑回答问题，模型在读取阶段就会自发判定没有新证据需要提取，从而触发提前退出（Adaptive Exit）。
在策略更新层面，算法遵循无须标准差归一化的组相对优势计算方法，将整条轨迹计算得到的标量优势值（Advantage）直接广播赋权给该轨迹内的每一个离散生成调用。这种端到端的奖励分配，让局部读取的“真假判断”、融合阶段的“去重提炼”与最终阶段的“答案推理”统一在同一个长程信用分配机制下，引导模型在兼顾高召回率的同时，学会极具经济性的退出策略。
3.6M 上下文极端实测：效率与准确率的双重飞跃
为了检验 PI-Mem 应对极端文本规模的能力，研究团队基于开源的高水准基座模型进行了大规模测试，测试模型涵盖了采用混合注意力与 MoE 架构的 Qwen3.5-35B-A3B，以及采用稠密全注意力机制的 Qwen2.5-7B。评测的核心场景集中在极具挑战的多跳问答基准 HotpotQA（HQA）上，并将上下文长度持续推高至令人咋舌的 360 万（3.6M）Token。
在 HotpotQA 的全长评测中，随着上下文长度从数万 Token 攀升到 3.6M，各类方法的表现呈现出鲜明的梯队分化：
直接采用原生模型进行全上下文推理（Vanilla）在短文本时尚有一战之力，但面对 3.6M 的极限长度，其巨大的显存占用与注意力计算开销使得推理极为缓慢，同时在超长干扰段落下面临严重的注意力分散；而以 MemAgent 为代表的传统循环记忆基准，虽然显存受控，但因其串行覆写的固有缺陷，得分随长度增加出现明显衰减。
相比之下，PI-Mem 展现出了极为顽强的长程推理鲁棒性。在 3.6M Token 的超长上下文场景下：
-
使用 Qwen3.5-35B-A3B 时，PI-Mem 相比 MemAgent 实现了 +6.25 分的绝对性能优势；而在端到端推理耗时上，PI-Mem 的速度达到了 MemAgent 的 6.1 倍，相较于直接全文本推理（Vanilla）更是快了 16.9 倍。
-
使用 Qwen2.5-7B 时，PI-Mem 的领先优势进一步扩大，相比 MemAgent 取得了 +7.81 分的绝对提升，推理延迟缩减为 MemAgent 的近一半（2.1 倍加速），相比原生直接推理提速 14.4 倍。
为了更透彻地回答“PI-Mem 是否真正解决了串行记忆覆盖问题”，研究人员在 RULER 的多值大海捞针（Multi-values Needle-in-a-Haystack）任务上对最终记忆的证据覆盖率（Evidence Coverage）进行了微观切片分析。该任务要求从百万级长文本的不同分布区间中完整找出关联至同一键值的 4 个独立目标数值。测试结果显示出巨大的反差：MemAgent 在历经数百个分块的连续覆盖后，遗失了绝大部分已读证据；而 PI-Mem 在多轮结束后的紧凑全局记忆中，近乎 100% 完整保全了所有分散的目标键值。这一证据直接证实了并行读取与集中融合机制在杜绝历史信息被随机冲刷方面的绝对优势。
机制拆解：各组件如何协同抗击长程噪声？
在严谨的消融实验中，PI-Mem 各个核心构件的价值得到了深入拆解。
首先是关于筛选（Select）与融合（Merge）组件的必要性。当在系统中移除 Select 模块、强行将所有分块产出的观察无差别推向 Merge 阶段时，模型在超长距离下的准确率显著下挫。这是由于无筛选的记忆极易受到长文本中海量无关背景知识的侵蚀，大量噪声进入候选池引发了“上下文污染”（Context Pollution）。反过来，如果保留 Select 但移除 Merge 模块，改为将筛选出的有效片段简单机械拼接，记忆体积会随着分块增加呈现线性发散，推理延迟随之剧增，碎片化的未消化文本也会阻碍模型在最终阶段的逻辑推理。这两者的结合证明，长文本记忆系统既需要严苛的局部准入控制，也需要结构化的集中代谢机制。
其次是轮次效率奖励（Turn-Efficiency Reward）对系统动态行为的深远塑造。消融数据表明，在未加入轮次奖励的基线强化学习训练中，模型学到的退出策略极为迟缓且不稳定，经常在已经搜集到足够线索的情况下继续无休止地运行至最大轮次，甚至因过度提炼而引入新的幻觉。引入轮次奖励后，模型的训练曲线表现出了高度自律的收敛特性：多余的无效轮次在训练早期被快速压制，模型学会在获取核心证据后立即自发执行早期退出。更值得注意的是，减少轮次不仅没有损伤下游任务的求解，其后期的最终准确率反而显著优于无节制多轮运行的版本。这印证了论文中的判断：对于长文本记忆而言，过度迭代往往伴随着潜在的信息畸变，适时收手不仅能节约算力，更是保障证据纯度的最优解。
最后，强化学习所带来的泛化红利同样显著。研究人员利用基于 HotpotQA 合成的数据训练策略模型，并将其直接放置在 RULER 的各类域外分布（OOD）任务中进行测试。实验显示，在 8K 短文本时，经过 RL 训练的模型与未训练版本差异不大；但随着上下文长度扩展至 100 万 Token，经过强化学习校准的 PI-Mem 展现出了极为明显的优势拉大态势。这说明端到端 RL 赋予模型的并非针对特定问题的死记硬背，而是一套通用的长程记忆检索调度“元能力”：如何根据记忆空缺有针对性地在各分块中抓取线索，以及如何将零散证据有效折叠整合。
从硬扛窗口到算法架构协同
长期以来，工业界与学术界在应对超长上下文时存在一种路径依赖：要么把注意力机制硬生生改造成各类稀疏近似甚至线性复杂度结构，要么依靠工程优化把原生上下文窗口强行堆叠到几百万乃至上千万。然而，由于 Transformer 在注意力分布上的弥散效应，简单的窗口物理延展往往无法换来逻辑推理能力的等比例增长，其高昂的预填充（Prefill）与 KV 显存开支更是阻碍实时应用落地的一道鸿沟。
PI-Mem 的价值在于提供了一条截然不同的解题思路。它并未试图在底层注意力矩阵上动大手术，而是从工作流架构设计与端到端策略学习的维度切入，将超长上下文推理视作一种具备内在交互机制的系统级协作。
通过将线性的单向依赖彻底重构成空间上的“全域并发读取”与时间上的“有限轮次提炼”，PI-Mem 在底层充分榨取了现代 AI 芯片的并行批处理红利，在上层又以数学形式构筑了抗信息遗忘的记忆融合闭环。随着大模型在代码库级 Agent、万卷长书全景洞察等复杂推理场景下的持续渗透，这种将计算图解耦、赋予模型反思性信息代谢机制的并行记忆工作流，无疑为攻克超长上下文落地瓶颈提供了极具启发意义与工程可行性的新蓝图。