Reasoning with Memory: Adaptive Information Management for Retrieval-Augmented Generation

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 1 , 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. * Work conducted during internship at Amazon. † Corresponding author. 1 https://github.com/amazon-science/state-aware-rag information is treated statically-simply appended to context-causing irrelevant, redundant, or conflicting information to accumulate and compound across reasoning steps. Recent advances focus almost exclusively on information acquisition-improving what enters the system, through query reformulation (