Search papers, labs, and topics across Lattice.
This paper introduces PI-Mem, a novel memory mechanism that enhances long-context reasoning in large language models by processing all chunks in parallel and iteratively refining a shared memory. By addressing the limitations of sequential updates and inter-chunk dependencies, PI-Mem achieves significant improvements in both accuracy and inference speed on the HotpotQA benchmark, outperforming traditional recurrent-memory approaches. The results show an increase in performance by up to 7.81 absolute points and speedups of 6.1脳 and 2.1脳, effectively breaking the accuracy-efficiency trade-off in handling contexts of up to 3.6 million tokens.
PI-Mem achieves unprecedented long-context reasoning capabilities, outperforming traditional methods while accelerating inference by over 6 times.
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.
