Joint Optimization of Multi-agent Memory System
About
Memory systems are critical for LLMs, mitigating context window limitations and supporting long-horizon user-LLM interactions. Such systems typically comprise multiple agents responsible for memory construction and retrieval. Existing approaches often optimize each agent independently under a shared global objective (e.g., downstream QA accuracy), treating other agents as a static environment. However, this design has two key limitations: (1) independent optimization ignores inter-agent dependencies and lacks agents' co-adaptation, and (2) relying solely on sparse global rewards provides limited guidance for optimizing specialized agents and causes ambiguous credit assignment. These may ultimately limit agents' effective collaboration in the memory system. To address these limitations, we propose CoMAM, a joint optimization framework that promotes collaboration among agents via end-to-end reinforcement learning and an adaptive credit assignment mechanism. Specifically, we model the multi-agent pipeline as a Markov decision process (MDP) to expose inter-agent dependencies during end-to-end training. Agents are then jointly optimized using a combination of their local task reward and an adaptively weighted global reward, enabling agents to co-adapt while receiving targeted feedback for their respective roles. Experiments show that CoMAM consistently outperforms leading memory systems, validating the effectiveness of the joint optimization framework.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Query Answering | PersonaMem 128K context length | Query-Answering Accuracy70 | 60 | |
| Query Answering | PersonaMem 32K context length | Query-Answering Accuracy64 | 60 | |
| Query Answering | PersonaMem 1M context length | Query-Answering Accuracy69 | 38 | |
| Multiple-choice Query Answering | PersonaMem (Average) | Accuracy67 | 22 |