让大模型处理360万字长文本推理,速度更快且更准。
PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory

- 并行迭代更新共享记忆,避免早期信息被覆盖
- 在360万词长文本上准确率提升超7个百分点,推理提速达6.1倍
- 适合需要长文档多跳推理的科研与工程场景
长上下文推理仍是大语言模型的关键瓶颈。现有循环记忆方法存在两个固有缺陷:逐块顺序更新会因后续无关内容覆盖早期关键证据,串行块间依赖限制并行性并导致延迟随上下文长度增加。为此,我们提出PI-Mem(并行迭代记忆)机制,可并行处理所有文本块,并在有限轮次内迭代优化共享记忆。每轮中,PI-Mem基于当前记忆并行读取所有块,从中选择新或互补证据,并融合为紧凑共享记忆以供下一轮使用。为减少冗余轮次,我们引入强化学习优化流程,加入辅助的“轮次效率”奖励,使模型能自适应退出。我们在HotpotQA基准上,使用Qwen3.5-35B-A3B和Qwen2.5-7B模型,在最长达360万词的上下文下评估,结果表明其相比循环记忆基线分别提升+6.25和+7.81绝对准确率,推理速度分别提升6.1×和2.1×。这些结果证明PI-Mem打破了长上下文推理中的准确率-效率权衡,为极长文档上的复杂多跳问答提供了可扩展方案。
原文摘要 · Abstract (English)
Long-context reasoning remains a critical bottleneck for large language models, as recent recurrent-memory approaches face two inherent challenges: sequential chunk-wise updates can overwrite early critical evidence with later irrelevant content, and serial inter-chunk dependencies limit parallelism and cause latency to increase with context length. To address these issues, we propose PI-Mem (Parallel-Iterative Memory), a mechanism that processes all chunks in parallel and iteratively refines a shared memory over a bounded number of turns. In each turn, PI-Mem reads all chunks in parallel conditioned on the current memory, selects new or complementary evidence from each chunk, and merges the selected evidence into a compact shared memory for the next turn. To discourage redundant turns, we optimize the workflow through reinforcement learning with an auxiliary turn-efficiency reward, enabling the model to adaptively exit once sufficient evidence has been accumulated. We evaluate PI-Mem with Qwen3.5-35B-A3B and Qwen2.5-7B on the HotpotQA benchmark across context lengths up to 3.6 million tokens and find that it outperforms the recurrent-memory baseline by +6.25 and +7.81 absolute points while achieving 6.1$\times$ and 2.1$\times$ inference speedups, respectively. These results demonstrate that PI-Mem breaks the accuracy--efficiency trade-off in long-context reasoning and provides a scalable approach to complex multi-hop question answering over extremely long documents.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。