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

TL;DR

InfMem introduces a System-2–style memory-control framework for ultra-long document QA under bounded compute. It implements a structured PreThink--Retrieve--Write loop with early stopping to actively manage memory and retrieve only necessary evidence, backed by a practical SFT-to-RL training pipeline and verifier-based rewards. Across up to 1M-token contexts, InfMem consistently surpasses MemAgent in accuracy while markedly reducing latency through adaptive stopping, with gains validated across Qwen backbones and transfer to LongBench QA. The work demonstrates that cognitive control of evidence—rather than sheer memory capacity—is a key bottleneck in long-context reasoning, suggesting a pathway to scalable, efficient reasoning in extreme-length regimes.

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 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.9\times$ on average (up to $5.1\times$) via adaptive early stopping.

Figures (13)

• Figure 1: The InfMem System-2 Framework. Unlike passive streaming agents, InfMem instantiates an active System-2 control loop (PreThink--Retrieve--Write) to manage bounded memory. (1) PreThink acts as a cognitive controller, monitoring memory sufficiency to decide whether to answer immediately (Early Stop) or seek more information. (2) Retrieve executes targeted global search, fetching sparse evidence $\mathbf{r}_t$ from the index $\{p_j\}$ to bridge logical gaps. (3) Write performs joint compression, synthesizing the retrieved evidence with the current stream $\mathbf{c}_t$ to update the memory $\mathbf{m}_t$. This loop enables the agent to actively maintain evidence fidelity under extreme context lengths.
• Figure 2: Long-context scaling of Qwen3-4B up to 1M tokens on synthesized long-context QA benchmarks. InfMem demonstrates remarkable robustness in long-context scaling, maintaining consistent accuracy on synthetic benchmarks up to 1M tokens without performance degradation
• Figure 3: Inference Efficiency versus QA Performance on 1M Context Scaling. Notably, InfMem exhibits exceptional proficiency in long-range multi-hop reasoning, preserving high-fidelity performance without the computational overhead typically associated with extreme sequence lengths.
• Figure 4: Prompt template for the Retrieval Planner, which decides whether to call retrievesearch again or stop, without answering the question.
• Figure 5: Prompt template for memory updating, integrating both retrieved and recurrent chunks to refine the memory state.