InfMem: Learning System-2 Memory Control for Long-Context Agent
InfMem enhances long-context reasoning accuracy by 10.17 points using the PreThink-Retrieve-Write protocol.
Key Findings
Methodology
InfMem employs a PreThink-Retrieve-Write protocol to actively monitor evidence sufficiency, perform targeted retrieval, and apply evidence-aware joint compression to update bounded memory. It combines supervised fine-tuning and reinforcement learning to align retrieval, writing, and stopping decisions with end-task correctness.
Key Results
- On Qwen3-1.7B, InfMem improved accuracy by 10.17 points while reducing inference time by 3.9x.
- On Qwen3-4B, accuracy increased by 11.84 points, significantly enhancing inference efficiency.
- On Qwen2.5-7B, accuracy improved by 8.23 points, with early stopping reducing redundant computation.
Significance
InfMem provides significant improvements in accuracy and efficiency for long-context reasoning, addressing the issue of losing low-salience bridging evidence in existing methods. This method holds broad potential applications in academia and industry, especially in scenarios requiring ultra-long document processing.
Technical Contribution
InfMem introduces a System-2-style control loop, distinct from traditional streaming methods, offering new theoretical guarantees and engineering possibilities. Its innovative training strategy combines supervised fine-tuning and reinforcement learning, optimizing memory control decisions.
Novelty
InfMem is the first to apply System-2 control to long-context reasoning, significantly enhancing complex problem-solving capabilities through active retrieval and compression mechanisms, contrasting sharply with existing streaming frameworks.
Limitations
- In extremely complex multi-hop reasoning tasks, memory insufficiency may still occur.
- Requires substantial computational resources for training and optimization.
Future Work
Future research directions include optimizing memory compression algorithms, exploring more efficient training strategies, and validating in more practical application scenarios.
AI Executive Summary
In the realm of ultra-long document reasoning tasks, existing streaming methods often fail to preserve low-salience bridging evidence necessary for multi-hop reasoning. InfMem introduces a System-2-style control mechanism through the PreThink-Retrieve-Write protocol, actively monitoring evidence sufficiency, performing targeted retrieval, and applying evidence-aware joint compression to update bounded memory. This innovative approach excels across multiple ultra-long QA benchmarks, significantly improving accuracy and inference efficiency.
The core technologies of InfMem include a training strategy that combines supervised fine-tuning and reinforcement learning to ensure retrieval, writing, and stopping decisions align with end-task correctness. Experimental results show that InfMem consistently outperforms existing MemAgent models on Qwen series, particularly in complex multi-hop reasoning tasks.
Despite the significant advancements achieved by InfMem in long-context reasoning, memory insufficiency may still pose a challenge in extremely complex tasks. Future research will focus on optimizing memory compression algorithms and exploring more efficient training strategies to further enhance reasoning capabilities and efficiency.
Deep Analysis
Background
In recent years, with the remarkable performance of large language models across various tasks, long-context reasoning has become an important research area. Traditional streaming methods process documents in segments, maintaining a fixed-size memory, but often lose critical bridging evidence in multi-hop reasoning tasks.
Core Problem
The core problem of long-context reasoning is synthesizing sparse evidence scattered across distant segments under strict memory constraints. This problem is challenging due to the need to ensure both reasoning accuracy and computational efficiency.
Innovation
InfMem's core innovation lies in introducing a System-2-style control mechanism through the PreThink-Retrieve-Write protocol, actively monitoring evidence sufficiency, performing targeted retrieval, and applying evidence-aware joint compression to update bounded memory.
Methodology
- �� PreThink: Assess whether current evidence is sufficient.
- �� Retrieve: Perform targeted retrieval from a global document index.
- �� Write: Jointly compress retrieved passages, the current chunk, and existing memory under a fixed budget.
Experiments
Experiments utilized Qwen series models to evaluate InfMem's performance on ultra-long QA benchmarks. Comparisons with MemAgent validated InfMem's advantages in accuracy and inference efficiency.
Results
Results showed InfMem consistently outperformed MemAgent across benchmarks, particularly in complex multi-hop reasoning tasks. Adaptive early stopping significantly reduced inference time.
Applications
InfMem can be applied in scenarios requiring ultra-long document processing, such as legal contract analysis and large codebase understanding. Its efficient reasoning capability and accuracy hold broad potential in industry.
Limitations & Outlook
Despite significant advancements, memory insufficiency may still pose a challenge in extremely complex tasks. Future research will focus on optimizing memory compression algorithms and exploring more efficient training strategies.
Plain Language Accessible to non-experts
Imagine you are in a huge library searching for answers. Traditional methods are like searching each shelf for a book and remembering its content, but it's easy to forget important details. InfMem is like a smart librarian who first thinks about which books are needed for the question, quickly finds them, and notes down the important content. This way, when you need the answer, it's already prepared without searching the entire library again.
ELI14 Explained like you're 14
Imagine you're playing a super long adventure game where you need to find clues hidden in different levels. Regular methods are like just moving forward and trying to remember each level's details, but it's easy to forget important clues. InfMem is like a smart helper who first thinks about which clues are needed for the question, quickly finds them, and notes down the important content. So when you need the answer, it's already prepared without going back to find the clues again.
Glossary
PreThink
The process of evaluating whether current evidence is sufficient to answer the question.
Used in InfMem to decide whether retrieval is needed.
Retrieve
The process of targeted retrieval from a global document index to gather necessary evidence.
Used in InfMem to acquire supplementary evidence.
Write
The process of jointly compressing retrieved passages, the current chunk, and existing memory under a fixed budget.
Used in InfMem to update memory.
System-2
A mode of thinking that requires deliberate reasoning and logic.
Describes the control mechanism in InfMem.
MemAgent
A traditional streaming method using fixed-size memory for reasoning.
Benchmark method compared with InfMem.
Open Questions Unanswered questions from this research
- 1 How to further improve memory efficiency in extremely complex multi-hop reasoning tasks?
- 2 How to optimize memory compression algorithms without increasing computational costs?
Applications
Immediate Applications
Legal Contract Analysis
InfMem can be used to analyze ultra-long legal contracts, quickly extracting key clauses and improving legal review efficiency.
Long-term Vision
Large Codebase Understanding
InfMem can be used to understand large codebases, helping developers quickly locate key code segments and improve development efficiency.
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.