InfMem: Learning System-2 Memory Control for Long-Context Agent

AI summary

Overview

Research area: Natural Language Processing, specifically long-context LLM reasoning, memory-augmented agents, and reinforcement learning for agentic control.

Technical level: Intermediate to Advanced. Readers should be comfortable with LLM inference pipelines, retrieval-augmented generation, streaming/chunked processing, and the basics of policy-gradient reinforcement learning (GRPO) and supervised fine-tuning.

Scope: The paper proposes a bounded-memory agent called InfMem that replaces the passive, fixed compression routine of existing streaming long-context agents with an explicit, learnable control loop (PreThink–Retrieve–Write) trained via a two-stage SFT→RL recipe, and evaluates it on QA benchmarks spanning 32k to 1M tokens.

What This Paper Is About

When an LLM must answer a question using a document far longer than its context window, a common workaround is to read the document chunk by chunk while maintaining a small fixed-size memory that gets overwritten at each step. The problem is that this routine compresses every chunk the same way, so subtle "bridging" facts needed for multi-hop reasoning tend to get silently discarded before the agent realizes they matter. InfMem's goal is to give the agent deliberate, learned control over three decisions: whether it already knows enough to answer, what to go back and search for inside the document, and how to merge new evidence into its limited memory.

Key Contributions

1. A PreThink–Retrieve–Write control loop. InfMem adds an explicit decision layer over streaming memory. At each step, a controller inspects the question and accumulated memory, decides whether to stop or search, synthesizes a retrieval query, and predicts how many document units to fetch. Retrieval is followed by evidence-aware joint compression into a bounded memory.
2. Separation of the reading view from the retrieval view. The agent reads the document as coarse sequential chunks, but searches a separately pre-built, globally indexed set of finer-grained units (paragraphs). This lets the agent jump backward or forward in the document to recover lost evidence without abandoning the linear streaming cost profile.
3. An adaptive early-stopping mechanism. Once the controller judges memory sufficient, the loop halts and answers immediately, preventing redundant overwrites from degrading consolidated evidence and cutting inference cost.
4. A practical SFT→RL training recipe with a composite reward. Stage 1 distills protocol-consistent trajectories from a Qwen3-32B teacher into smaller students. Stage 2 applies GRPO to align retrieval, writing, and stopping decisions with end-task correctness, using four reward signals: ground-truth answer correctness, early-stop shaping, functional-call well-formedness, and memory-update completeness.

Main Findings

• Consistent double-digit gains over the strong streaming baseline. InfMem improves average absolute accuracy over MemAgent by +10.17, +11.84, and +8.23 points on Qwen3-1.7B, Qwen3-4B, and Qwen2.5-7B respectively on synthesized ultra-long QA benchmarks.
• The advantage grows with context length and reasoning depth. At 1M tokens on MuSiQue with Qwen3-4B+RL, InfMem scores 56.86 versus MemAgent's 35.91; on 2WikiMultiHopQA at 1M tokens it reaches 66.39 versus 35.18. On single-hop SQuAD the gap narrows, which the authors attribute to fixed compression being adequate for direct fact extraction.
• Large inference efficiency wins. Adaptive early stopping reduces inference latency by 3.9× on average (up to 5.1×). At roughly 896k context, InfMem decodes 33.1K tokens per query versus 117.4K for MemAgent (0.28×), and its token cost grows about 5× from 112k to 896k where MemAgent's grows about 11×.
• Baselines degrade sharply past 128k tokens. YaRN and vanilla RAG collapse to single-digit accuracy at 1M tokens; iterative RAG (RESP) and bounded-memory baselines (Mem-α) mitigate but still trail MemAgent and InfMem on average.
• Gains transfer to standard benchmarks. On LongBench QA, InfMem+RL improves over the YaRN baseline by +8.93, +11.47, and +16.92 average points on the three backbones, showing the benefit is not limited to artificially long synthesized inputs.
• RL adds gains beyond SFT. SFT primarily teaches protocol mechanics; the RL stage, which operates without any teacher hints or question decompositions, is responsible for refining the control policy itself.

Methodology in Plain English

The authors frame the problem as a streaming agent reading a document in chunks while carrying a memory of fixed token budget. Their agent runs a three-part loop at every step:

• PreThink looks only at the question and the current memory. If the memory already looks sufficient, it emits a stop signal. Otherwise it emits a search action, a query written in natural language, and a number k saying how many document units to fetch. Optionally it can write a short rationale, which helps interpretability but does not affect execution.
• Retrieve uses that query against a pre-built index of finer-grained units taken from the same document (no external corpus), returning the top-k passages and concatenating them with separators.
• Write merges three sources under the token budget: the newly arrived chunk, the retrieved passages, and the previous memory. The instruction is to preserve verifiable facts and bridging links while discarding stylistic or speculative content, prioritizing new complementary evidence without dropping previously useful facts.

The training pipeline has two stages. In supervised fine-tuning, a Qwen3-32B teacher executes the full protocol on data built from HotpotQA, SQuAD, and MuSiQue, where each original question is embedded into a 32K-token document by inserting gold evidence once and padding with same-corpus distractors. Only correct trajectories survive filtering, yielding 29,717 single-turn dialogue instances; the student is trained with masked next-token prediction on agent-generated tokens only, so it never sees teacher decompositions or gold document IDs.

In the reinforcement learning stage, the same synthesis pipeline extends HotpotQA to roughly 28K tokens, and GRPO optimizes a weighted sum of four rewards: binary answer correctness, an exponentially decaying bonus for stopping soon after the memory first becomes sufficient, a check that every function call parses, and a check that every memory update is complete and within budget. The same trajectory-level reward is propagated back to all preceding steps to handle long-horizon credit assignment.

Why This Matters

Research impact. The paper argues that as context windows scale, the bottleneck is shifting from raw memory capacity to cognitive control — the ability to know what is already known. It provides a concrete instantiation of that idea, showing that an explicit, learned sufficiency-and-retrieval controller outperforms uniformly applied recurrent compression, and it demonstrates that early-stopping policy is something that can be shaped by reward rather than left as a heuristic.

Real-world applications:

• Legal and regulatory review, where answers depend on specific clauses scattered across hundreds of pages of contracts or filings and the reviewer must know when the evidence base is complete.
• Large-codebase question answering, where a fix or behavior depends on a handful of lines across many files and the agent must trace cross-file dependencies.
• Scientific and medical literature synthesis, where a conclusion requires chaining findings from separated sections or papers rather than extracting one sentence.
• Enterprise knowledge and customer-support agents, where long internal document stores must be searched with bounded compute and answers returned quickly.

Industry relevance. The efficiency results matter directly for serving costs. Halving or better the number of decoded tokens per query at extreme context lengths, combined with early termination, makes bounded-memory agents more viable for production deployments where latency and GPU budget are hard constraints, and the SFT→RL recipe is designed to be applied to relatively small backbones (1.7B to 7B) rather than frontier models.

Future Directions

• Reducing reliance on a heuristic sufficiency check. The early-stop reward currently uses the exact-match protocol to determine when memory first becomes sufficient, which the authors acknowledge can misjudge hard multi-hop cases; learning a dedicated sufficiency estimator is an open problem.
• Joint learning of the retrieval index. The retrieval units are pre-constructed and fixed; whether the agent could learn or adapt the granularity and contents of that index alongside the control policy is unaddressed.
• Extending beyond single-document, closed-corpus settings. The current retrieval view is built from the same document streamed by the agent; generalizing to open corpora or multiple documents raises new questions about index consistency and distractor handling.
• Bridging bounded memory with parametric consolidation. The related work notes a plasticity–stability dilemma in methods that absorb context into model weights; combining InfMem's explicit memory control with weight-based consolidation is a natural extension.
• Scaling and broader domain transfer. The paper reports a zero-shot transfer of the protocol to a non-Qwen backbone, but testing at larger model scales and in domains beyond QA-style benchmarks remains open.

Target Audience

Researchers and engineers working on long-context LLMs, memory-augmented agents, and agentic RL will benefit most, particularly those building streaming or bounded-memory inference systems for document QA. The paper is also relevant to practitioners in retrieval-augmented generation who want to understand how retrieval decisions, memory writes, and stopping criteria can be jointly optimized rather than hand-tuned. Readers unfamiliar with policy-gradient RL or agentic tool-calling may find the training section dense, though the architectural contributions and empirical results remain accessible on their own.

Authors’ abstract

Reasoning over ultra-long documents requires synthesizing sparse evidence scattered across distant segments under strict memory constraints. While streaming agents enable scalable processing, their passive memory update strategy often fails to preserve low-salience bridging evidence required for multi-hop reasoning. We propose InfMem, a control-centric agent that instantiates System-2-style control via a PreThink-Retrieve-Write protocol. InfMem actively monitors evidence sufficiency, performs targeted in-document retrieval, and applies evidence-aware joint compression to update a bounded memory. To ensure reliable control, we introduce a practical SFT-to-RL training recipe that aligns retrieval, writing, and stopping decisions with end-task correctness. On ultra-long QA benchmarks ranging from 32k to 1M tokens, InfMem consistently outperforms MemAgent across backbones. Specifically, InfMem improves average absolute accuracy by 10.17, 11.84, and 8.23 points on Qwen3-1.7B, Qwen3-4B, and Qwen2.5-7B, respectively, while reducing inference time by 3.9x on average and by up to 5.1x through adaptive early stopping. Code is available at https://github.com/UCMP13753/InfMem.