Jaejin Cho
2026
Reasoning with Memory: Adaptive Information Management for Retrieval-Augmented Generation
Hieu Man | Ro-ee Tal | Abhishek Kumar | Jaejin Cho | Benjamin Hsu
Findings of the Association for Computational Linguistics: ACL 2026
Hieu Man | Ro-ee Tal | Abhishek Kumar | Jaejin Cho | Benjamin Hsu
Findings of the Association for Computational Linguistics: ACL 2026
Multi-hop reasoning remains a fundamental challenge for Retrieval-Augmented Generation (RAG) systems. Recent approaches—from adaptive retrieval to agentic pipelines—struggle to maintain coherent intermediate reasoning states as chains grow longer. We introduce State-Aware RAG, a framework that addresses this limitation through an explicit working memory that serves as a dynamic cognitive workspace for reasoning. Our modular architecture features a lightweight, trainable extractor that learns to actively filter, consolidate, and update this working memory via a novel Path-Outcome Dual Reward paradigm, which balances local coherence with global strategy. The retriever and generator remain frozen, enabling plug-and-play flexibility. Experiments on eight QA benchmarks demonstrate state-of-the-art results, on average achieving +8.6% over the best memory-augmented baseline and +9.3% over the best RL-enhanced baseline. Our architecture generalizes seamlessly to stronger generators and retrievers without retraining, establishing dynamic memory management as a critical yet underexplored dimension for advancing RAG systems.