Joint Optimization of Multi-agent Memory System

TL;DR

This paper introduces CoMAM, a joint optimization framework for multi-agent memory systems in LLMs, using end-to-end reinforcement learning to improve collaboration and performance.

Contribution

It proposes a novel joint optimization method with adaptive credit assignment for multi-agent memory systems, addressing limitations of independent agent training.

Findings

CoMAM outperforms existing memory systems in experiments.

Joint optimization improves agent collaboration and overall system performance.

Adaptive credit assignment enhances targeted feedback for agents.

Abstract

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.