Memory agents improve answers by learning from process rewards
All-in-One Memory: Process-Rewarded Reinforcement Learning for Memory Agents
Artificial Intelligence
Summary
Long conversations and documents are hard for AI to remember and use well because their memory processes are often handled separately. The authors propose a new method that treats remembering, searching, writing, and stopping as one connected process with shared rewards. This method helps AI keep better track of important information and decide what to remember or retrieve. It was tested on tasks requiring understanding long contexts, showing it improves answer accuracy compared to older methods.
What this means in practice
- •For conversational ai developers: Build chatbots that remember and use information from long, multiple sessions more accurately by optimizing memory actions as a unified process.
- •For document processing teams: Enhance automated systems that handle lengthy documents by improving memory retrieval and writing decisions within a shared reward framework.
- •For customer support platform engineers: Create systems that track and utilize long customer interaction histories more effectively to provide better and contextually accurate responses.$Commercial implications: Enables development of support tools that recall and process long-term conversation context, improving service quality and automation.
Authors
Bingbing Wang, Jingjie Lin, Weiming Qiao, Jing Li, Ruifeng Xu
Abstract
Large language models increasingly rely on external memory to answer queries grounded in long-horizon histories such as multi-session dialogues and lengthy documents. However, existing memory methods often optimize retrieval and memory construction separately, without a shared reward space for evaluating search, writing, and stopping decisions. To address this, we propose All-in-One Memory (AiO-Mem), which organizes these operations as transitions over a shared evidence state and directly evaluates the semantic quality of each action in a unified reward space. AiO-Mem compiles heterogeneous histories into source-traceable multi-view events, providing a shared substrate for retrieval and memory writing. A unified memory policy actor then rolls out complete interaction trajectories over this shared evidence state. A separate trajectory-conditioned process reward model is trained through evidence-grounded supervised fine-tuning and trajectory-level DPO to assign process scores to the collected trajectories. These scores are aggregated into an action-balanced trajectory reward and used by Process-Aligned DAPO (PA-DAPO) to optimize the memory policy. Across long-context memory benchmarks, AiO-Mem consistently improves answer accuracy over both static memory systems and RL-trained memory agents.