💬 NLP

PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory

The paper introduces PI-Mem, a parallel-iterative memory mechanism that processes all context chunks simultaneously and iteratively refines a shared memory with reinforcement learning-based adaptive termination, enabling large language models to achieve superior accuracy and significant inference speedups on long-context reasoning tasks up to 3.6 million tokens.

Original authors: Dawei Liu, Haixu Song, Shuang Cheng, Shijie Wang, Haozheng Hou, Kaifeng Liu, Ermo Hua, Zhonghang Yuan, Zhijie Zhong, Yuchen Fan, Biqing Qi, Bowen Zhou

Published 2026-08-05

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dawei Liu, Haixu Song, Shuang Cheng, Shijie Wang, Haozheng Hou, Kaifeng Liu, Ermo Hua, Zhonghang Yuan, Zhijie Zhong, Yuchen Fan, Biqing Qi, Bowen Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive, 3.6-million-page mystery novel. You have a brilliant detective, a Large Language Model (LLM), who is incredibly smart but has a very short-term memory. If you hand the detective the whole book at once, their brain gets overwhelmed, and they start forgetting the clues from page 10 by the time they reach page 1,000,000. This is the "long-context" problem in artificial intelligence: how do we keep a computer's attention on the beginning of a story while it's reading the end?

To solve this, scientists have tried two main tricks. The first is like trying to stretch a rubber band; they try to force the model to remember everything at once, but the rubber band often snaps or gets too slow to use. The second, more popular trick, is the "Recurrent Memory" approach. Imagine a detective who reads the book one chapter at a time. After reading Chapter 1, they write a tiny summary note in a notebook, throw away the chapter, and move to Chapter 2. They read Chapter 2, update their note, throw away the chapter, and so on. The problem? If the detective reads a confusing chapter later on, they might accidentally scribble over their important note from Chapter 1, erasing a crucial clue forever. It's a slow, serial process where the detective has to wait for one note to be finished before starting the next, making the whole investigation drag on.

This is where a new team of researchers steps in with a fresh idea called PI-Mem (Parallel-Iterative Memory). Instead of the detective reading one chapter, updating the note, and moving on, PI-Mem gives the detective a superpower: the ability to read all the chapters at the same time, but with a twist.

Here is how PI-Mem works in the real world of this paper. Imagine the detective has a "Shared Global Memory" board. Instead of reading the book sequentially, the detective looks at every single chapter simultaneously, using the current state of the memory board as a guide. For each chapter, they ask: "Does this page have a new clue that isn't already on my board?" If the answer is yes, they pull that specific clue out. If the answer is no, they ignore the rest of the noise.

Once they've scanned all the chapters in parallel, they take only the new clues they found and merge them onto the memory board, creating a cleaner, more complete picture. Then, they repeat the process. They look at all the chapters again, but this time, they use the updated board to guide them. Maybe a clue they missed in Chapter 5 suddenly makes sense because they just found a matching piece in Chapter 200. They keep doing these "rounds" of scanning until no new clues are found or they hit a limit. Finally, they solve the mystery using only the consolidated clues on the board, not the entire 3.6-million-page book.

The researchers tested this method on two different AI models (Qwen3.5 and Qwen2.5) using a benchmark called HotpotQA, which involves answering tricky questions that require finding information across huge documents. They pushed the context length to a staggering 3.6 million tokens (roughly the size of a massive library).

The results were striking. Compared to the old "read-one-by-one" method (Recurrent Memory), PI-Mem didn't just get better answers; it got them much faster. On the 3.6-million-token test, PI-Mem improved the accuracy by 6.25 to 7.81 percentage points over the previous best method. Even more impressive, it was 6.1 times faster on the larger model and 2.1 times faster on the smaller one. The old method was so slow because the detective had to wait for every single step to finish before moving to the next. PI-Mem, by reading everything in parallel, turned a slow, winding road into a high-speed highway.

The team also used a special training technique called Reinforcement Learning to teach the AI when to stop. They gave the AI a "bonus" for finishing the investigation quickly if it had enough evidence, discouraging it from doing unnecessary extra rounds of reading. This ensured the system didn't waste time re-reading chapters it had already solved.

In short, the paper suggests that by changing how the AI "thinks" about long documents—switching from a slow, sequential note-taker to a parallel, iterative investigator—we can solve complex problems in massive amounts of text without losing the clues from the beginning. It proves that you don't need to force the AI to remember everything at once; you just need to give it a better way to organize the pieces of the puzzle as it finds them.

Technical Summary: PI-Mem

Problem Statement

Long-context reasoning remains a critical bottleneck for Large Language Models (LLMs). While recent recurrent-memory approaches attempt to handle inputs exceeding native context windows, they face two inherent challenges:

Evidence Overwrite: Sequential chunk-wise updates often overwrite early critical evidence with later, potentially irrelevant content before the relevance of the early evidence is fully established.
Latency and Parallelism Limits: The strict serial inter-chunk dependencies in recurrent workflows prevent parallel processing, causing inference latency to scale linearly with context length.

Existing alternatives, such as positional-encoding interpolation or sparse attention, often incur high computational costs or rely on difficult-to-construct synthetic corpora.

Methodology: PI-Mem

The authors propose PI-Mem (Parallel-Iterative Memory), a mechanism designed to process long contexts by gathering evidence in parallel and iteratively refining a shared memory state over a bounded number of turns.

Inference Workflow

The workflow operates in cycles (turns) defined by a Read-Select-Merge process:

Parallel Read: Unlike recurrent models that process chunks sequentially, PI-Mem reads all document chunks ( $C = \{c_i\}$ ) in parallel. Each chunk is processed conditioned on the same global memory state ( $m^{(k-1)}$ ) and the question ( $q$ ). This prevents the "forgetting" of early evidence due to sequential overwrites.
Select: For each chunk, the model generates an observation ( $o_i$ ) and explicitly outputs a <check> signal (yes or no). Only observations marked yes (indicating new or complementary evidence) are retained.
Merge: The selected observations are consolidated into a compact, updated global memory ( $m^{(k)}$ ). This step compresses and integrates evidence from different chunks, removing redundancy while preserving necessary details.
Iterative Refinement & Exit: The cycle repeats with the updated memory. The process terminates adaptively when no chunk yields new evidence (all <check> signals are no) or when a maximum turn limit ( $K$ ) is reached. The final answer is generated using only the question and the consolidated memory.

Training via Reinforcement Learning

PI-Mem is optimized end-to-end using Group Relative Policy Optimization (GRPO). The training objective combines:

Accuracy Reward: Based on the correctness of the final answer.
Turn-Efficiency Reward: An auxiliary reward ( $r_{turn}$ ) that penalizes redundant turns. This encourages the model to exit the loop as soon as sufficient evidence is accumulated, balancing performance with inference latency.

The reward is broadcast to all model calls (read, merge, and final answer) within a trajectory to optimize the entire workflow.

Key Contributions

Parallel-Iterative Mechanism: A novel architecture that decouples chunk processing from sequential dependencies, allowing parallel evidence extraction conditioned on a shared memory.
End-to-End RL Optimization: The introduction of a turn-efficiency reward to train the policy to adaptively exit the refinement loop, reducing redundant computation.
Empirical Validation: Comprehensive evaluation demonstrating that PI-Mem breaks the traditional accuracy-efficiency trade-off in long-context reasoning.

Experimental Results

The authors evaluated PI-Mem using Qwen3.5-35B-A3B and Qwen2.5-7B on the HotpotQA (HQA) benchmark within the RULER suite, scaling context lengths up to 3.6 million tokens.

Accuracy:
- On Qwen3.5-35B-A3B, PI-Mem improved the average score by +6.25 absolute points over the recurrent-memory baseline (MemAgent), achieving 81.25 vs. 76.56. At the specific 3.6M token setting, PI-Mem outperformed MemAgent by +7.81 absolute points.
- On Qwen2.5-7B, PI-Mem improved the average score by +6.56 points (84.06 vs. 77.50) and matched or outperformed MemAgent across all reported lengths, with especially large gains in the longer-context regime.
- PI-Mem also achieved superior performance on Out-of-Distribution (OOD) tasks and the LongBench v2 benchmark.
Efficiency (Latency):
- PI-Mem achieved significant inference speedups compared to both direct inference (Vanilla) and recurrent baselines.
- At 3.6M tokens, PI-Mem was 6.1× faster than MemAgent and 16.9× faster than Vanilla on Qwen3.5-35B-A3B.
- On Qwen2.5-7B, speedups were 2.1× over MemAgent and 14.4× over Vanilla.
Evidence Retention:
- In Multi-values Needle-in-a-Haystack (MV-NIAH) tasks, PI-Mem retained nearly 100% of ground-truth values in the final memory, whereas the recurrent baseline lost a substantial portion due to overwriting.

Significance

The paper claims that PI-Mem provides a scalable approach to complex multi-hop question answering over extremely long documents without relying solely on expanding native context windows. By processing chunks in parallel and iteratively refining a shared memory, the method effectively preserves evidence that recurrent models tend to overwrite. Furthermore, the integration of reinforcement learning allows the system to adaptively manage the trade-off between reasoning depth and computational cost, offering a promising path for ultra-long-context reasoning that maintains both high accuracy and low latency.