🤖 AI

MemReread: Enhancing Agentic Long-Context Reasoning via Memory-Guided Rereading

MemReread is a memory-guided framework that enhances agentic long-context reasoning by triggering question decomposition and rereading passes when initial memory is insufficient, thereby recovering discarded evidence and supporting non-linear reasoning while maintaining linear time complexity.

Original authors: Baibei Ji, Xiaoyang Weng, Juntao Li, Zecheng Tang, Yihang Lou, Min Zhang

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Baibei Ji, Xiaoyang Weng, Juntao Li, Zecheng Tang, Yihang Lou, Min Zhang

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

The Big Problem: Reading a Novel in a Hurricane

Imagine you are trying to solve a mystery, but the clues are hidden inside a 100,000-page book. You can only hold a few pages in your hands at a time.

The Old Way (Standard AI): The AI tries to read the whole book at once. But the book is so huge that the AI's brain gets "distracted" (a problem called attention dilution). It forgets the clues from page 50 by the time it reaches page 90,000.
The "Streaming" Way (Previous Agents): To fix this, the AI reads the book page-by-page (or chunk-by-chunk). As it reads, it writes a summary in a small notebook (Memory). Once a page is read, it's thrown away.
- The Flaw: Sometimes, a clue on page 10 is useless until you read page 50. But by the time the AI gets to page 50, it has already thrown away page 10 and forgotten the clue. The AI can't go back.
The "Retrieval" Way (Other Agents): These agents try to fix the "throwing away" problem by keeping a giant library of every page they've ever read. When they need a clue, they ask a librarian (Search Engine) to find it.
- The Flaw: The librarian is often confused. They might bring back the wrong pages because the AI asked a vague question, or they might bring back pages that were already thrown away. This "noise" confuses the AI, making it worse at solving the mystery.

The Solution: MemReread (The "Double-Check" Detective)

The authors propose MemReread. Instead of just reading once or frantically searching a library, the AI acts like a careful detective who knows when to stop and re-read.

Here is how it works, step-by-step:

1. The First Pass (The Scan)

The AI reads the long document from start to finish, just like the "Streaming" agents. It keeps a small, running summary in its memory.

Analogy: You are walking through a forest looking for a specific bird. You take notes on what you see.

2. The "Wait, I Missed Something" Moment

Once the AI finishes reading the whole book, it looks at its notes and asks: "Do I have enough information to answer the question?"

If Yes: It gives the answer.
If No: It realizes it missed a crucial link (like the clue on page 10 that connects to page 50).

3. The "Reread" (The Targeted Re-Scan)

Instead of searching a messy library, the AI breaks the big question into a smaller, specific question (a "sub-question"). It then goes back to the beginning of the book and reads it again, but this time, it is looking only for the answer to that specific sub-question.

Analogy: You realize you missed the bird's nest. Instead of searching the whole forest randomly, you say, "Okay, I need to find the nest specifically." You walk through the forest a second time, eyes wide open, looking only for the nest.

4. The Update (The Smart Notebook)

Once the AI finds the missing piece during the second read, it updates its main notebook with the new fact. It then asks itself again: "Do I have everything now?" If not, it breaks down the next missing piece and reads again.

The Secret Sauce: The "Smart Coach" (Reinforcement Learning)

The paper also introduces a special training method (called Rereading-Adaptive GRPO) that teaches the AI when to stop.

The Problem: If the AI reads the book 10 times for every question, it will be very accurate but incredibly slow and expensive.
The Fix: The "Coach" rewards the AI for solving the problem with the fewest number of re-reads possible.
- If the question is easy, the AI learns to answer after one read.
- If the question is hard (like a complex puzzle), the AI learns that it needs to read twice or three times to get it right.
- Analogy: Think of it like a student taking a test. If they get the answer right on the first try, they get a gold star. If they have to check their work twice, they still get a star, but a smaller one. If they check it five times, they get no star. This teaches the AI to be efficient.

Why This is Better

No More "Lost in the Middle": Because the AI can go back and re-read specific parts, it doesn't lose the "indirect clues" that were thrown away earlier.
No More "Bad Librarian": It doesn't rely on a search engine that might bring back junk. It just re-reads the text itself, so it never gets confused by irrelevant noise.
Efficient: It keeps the memory small (linear complexity) and only spends extra time when the problem actually requires it.

The Results

The paper tested this on huge documents (up to 1 million words).

MemReread consistently beat the other methods (MemAgent and ReMemR1).
It was especially good at "Out-of-Distribution" tasks (questions it hadn't seen before), proving it actually learned how to reason, not just memorized answers.
It solved complex multi-step puzzles that other AI models failed because they couldn't connect the dots between the beginning and the end of the text.

Summary

MemReread is an AI that reads a long document, realizes it missed a clue, and then intelligently decides to re-read the document specifically to find that missing piece. It uses a smart training system to know exactly how many times it needs to re-read to get the answer right without wasting time. It's like a detective who knows when to scan the crime scene once, and when to go back for a second, more careful look.

Technical Summary: MemReread

Problem Statement

Large Language Models (LLMs) struggle with long-context reasoning tasks, particularly when context lengths approach or exceed their training windows. This performance degradation is primarily attributed to attention dilution, which impairs the identification of critical facts and disrupts coherent reasoning chains. While training to extend context windows is often prohibitively expensive due to the quadratic complexity of standard attention mechanisms, existing alternative approaches face significant limitations:

Retrieval-Augmented Memory Agents: These frameworks attempt to recover discarded information by querying historical memory or context chunks. However, they suffer from permanent latent evidence loss. If a fact is discarded early because its relevance is only recognized later (requiring a "bridge" fact from a future chunk), retrieval cannot recover it because the latent evidence was never archived. Additionally, these agents generate invalid queries when they cannot distinguish between information that has been overwritten versus information that simply resides in unread context, leading to memory pollution and interference.
Pure Memory Agents: These maintain a dynamically updated memory while processing chunks linearly. While they avoid retrieval noise, they are constrained by inherent sequentiality, preventing them from revisiting overlooked facts or performing non-linear reasoning effectively.

Methodology: MemReread

To address these limitations, the authors propose MemReread, a memory-guided agent framework that enhances streaming reading through an adaptive rereading mechanism. Unlike retrieval-based approaches, MemReread circumvents intermediate retrieval entirely.

Core Workflow

MemReread operates on a "memorize-while-reading" paradigm but introduces a feedback loop for insufficient information:

Read & Answer: The agent processes context chunks sequentially, maintaining a bounded memory buffer. Upon completing a pass, it attempts to answer the question based solely on the accumulated memory.
Decompose: If the final memory lacks sufficient evidence, the agent decomposes the original question into a specific sub-question. This sub-question is designed to isolate the highest-priority missing information (e.g., identifying an entity alias before tracking its value).
Reread: Guided by the generated sub-question, the agent performs a targeted rereading of the context chunks. This process is linear and streaming, focusing solely on acquiring the specific missing facts.
Integrate: Once the sub-question is answered, the resulting (sub-question, sub-answer) pair is integrated into the root memory. The intermediate sub-memory is discarded to maintain bounded complexity.
Iterate: This cycle repeats until the root memory contains all necessary information to answer the original question.

Training Strategy: Rereading-Adaptive GRPO

To optimize the trade-off between reasoning accuracy and computational cost, the authors introduce Rereading-Adaptive Group Relative Policy Optimization (ReA-GRPO).

Process Supervision: Utilizes state rewards from previous works (ReMemR1) to provide dense feedback on memory updates.
Outcome Supervision: Introduces a novel advantage calculation that dynamically adjusts based on the number of rereading passes ( $p$ $p$ ):
- For correct trajectories, the model is incentivized to minimize rereading passes (brevity).
- For incorrect trajectories, the model is encouraged to perform more rereading passes to recover missing facts.
- This mechanism allows the agent to learn when to stop reading and when to invest additional compute for difficult samples.

Key Contributions

Memory-Guided Rereading: A novel framework that recovers indirect facts discarded during initial reading by triggering question decomposition and targeted rereading, avoiding the pitfalls of retrieval-based recall (e.g., invalid queries and permanent evidence loss).
Linear Complexity: The architecture maintains linear time complexity ( $O(p_c \cdot n)$ , where $p_c$ is the max rereading passes and $n$ is context length) and constant space complexity (bounded memory), making it scalable for ultra-long contexts without the quadratic cost of attention or the storage burden of chunk-level retrieval.
Adaptive RL Training: The development of ReA-GRPO, which enables agents to dynamically determine the number of rereading passes based on task complexity, balancing performance gains with inference overhead.
Diagnostic Analysis: The authors provide a diagnostic dataset ("Global Reasoning") and empirical analysis demonstrating that retrieval-based agents suffer from specific failure modes (permanent latent evidence loss and interference) that rereading effectively mitigates.

Experimental Results

Extensive experiments on HotpotQA (in-distribution) and 2WikiMultiHopQA (out-of-distribution) across context lengths ranging from 8K to 1M tokens demonstrate:

Superior Performance: MemReread consistently outperforms baseline frameworks (MemAgent and ReMemR1) across various model scales (1.7B to 4B). Notably, the 4B MemReread model achieves up to 12.1% higher accuracy than ReMemR1 on out-of-distribution tasks.
Robustness to Length: The method maintains high performance even at 1M token contexts, where baselines often degrade significantly.
Efficiency: While inference time is higher than single-pass streaming (due to rereading), it remains linear. Crucially, the space overhead is constant and comparable to MemAgent, unlike retrieval-based agents whose memory usage scales linearly with context length.
Ablation Studies: Results confirm that increasing the maximum rereading limit improves performance, and the ReA-GRPO training strategy effectively reduces unnecessary rereading on simple tasks while maintaining high accuracy on complex ones.

Significance and Claims

The paper claims that MemReread represents a significant step forward in agentic long-context reasoning. By decoupling reading (information acquisition) from reasoning (identifying gaps and triggering rereading), the framework allows models to recover critical indirect facts that are typically lost in streaming paradigms. The authors emphasize that this approach achieves non-linear reasoning capabilities while preserving the inherent logical flow of document comprehension and maintaining linear computational complexity.

The work suggests that for ultra-long contexts, the ability to "look back" via structured rereading is more effective than attempting to retrieve from a fragmented memory or relying on a single pass of attention. The proposed ReA-GRPO strategy further demonstrates that reinforcement learning can effectively teach agents to self-regulate their computational effort, optimizing the performance-efficiency trade-off. The authors conclude that this framework offers a practical and scalable solution for long-context tasks, advancing the capabilities of LLM agents in scenarios where context length exceeds standard training windows.