Collaborative Multi-Agent Optimization for Personalized Memory System

Abstract: Memory systems are crucial to personalized LLMs by mitigating the context window limitation in capturing long-term user-LLM conversations. Typically, such systems leverage multiple agents to handle multi-granular memory construction and personalized memory retrieval tasks. To optimize the system, existing methods focus on specializing agents on their local tasks independently via prompt engineering or fine-tuning. However, they overlook cross-agent collaboration, where independent optimization on local agents hardly guarantees the global system performance. To address this issue, we propose a Collaborative Reinforcement Learning Framework for Multi-Agent Memory Systems (CoMAM), jointly optimizing local agents to facilitate collaboration. Specifically, we regularize agents' execution as a sequential Markov decision process (MDP) to embed inter-agent dependencies into the state transition, yielding both local task rewards (e.g., information coverage for memory construction) and global rewards (i.e., query-answer accuracy). Then, we quantify each agent's contribution via group-level ranking consistency between local and global rewards, treating them as adaptive weights to assign global credit and integrate local-global rewards. Each agent is optimized by these integrated rewards, aligning local improvements with the global performance. Experiments show CoMAM outperforms leading memory systems, validating the efficacy of our proposed collaborative reinforcement learning for joint optimization.

• The paper introduces CoMAM, a joint reinforcement-learning framework that coordinates extraction, profiling, and retrieval agents through sequential MDP regularization and adaptive credit assignment.
• CoMAM achieves 0.62–0.70 accuracy across Qwen and Llama models at 32K–1M context lengths, outperforming independent RL baselines by 8.5–16.7% relatively.
• Adaptive NDCG-based reward weighting improves global answer accuracy and agent-specific rewards while reducing training convergence steps through parallel multi-agent optimization.

Problem statement and motivation

Memory systems allow LLMs to personalize responses by retaining and selectively retrieving long-term user-LLM conversation histories that exceed the context window. Contemporary systems decompose this task into multi-granular memory construction (e.g., fine-grained facts and coarse-grained user profiles) and retrieval-augmented answering, each handled by specialized agents. The prevailing optimization paradigm, however, treats these agents independently: prompt engineering specializes each agent for its local task, and RL methods such as Mem1 and Memory-R1 fine-tune agents separately, often against a single query-answer accuracy signal. The central claim of this paper is that such independent optimization is structurally inadequate: individually well-optimized agents can produce strategies that conflict when composed, yielding "locally optimal but globally suboptimal" systems. The authors demonstrate this empirically with a teaser experiment showing that agents whose local rewards have converged still underperform a jointly optimized configuration across all history lengths. This motivates two technical challenges the paper addresses: (1) heterogeneity and asynchrony — construction and retrieval agents have different configurations, discrete parameter spaces, and asynchronous execution schedules; and (2) local-global alignment — naively distributing the global reward equally across agents introduces credit assignment ambiguity, since agent contributions to the global outcome differ substantially.

The CoMAM framework

CoMAM (Collaborative reinforcement learning framework for Multi-Agent Memory systems) instantiates a three-agent pipeline: an Extraction Agent producing fine-grained memories $\mathcal{M}_f$ from raw history $\mathcal{H}$; a Profile Agent abstracting coarse-grained memories $\mathcal{M}_c$ (preferences, behavioral patterns) from $\mathcal{M}_f$; and a Retrieval Agent that retrieves top-$K$ memories and generates the response. Each agent carries a task-specific local reward: coverage- and precision-weighted overlap with ground-truth evidence for extraction (weighted $\alpha = 0.8$ toward coverage) and retrieval (weighted $\beta = 0.2$ toward precision), and an LLM-as-judge score $\mathcal{V}$ for profile abstraction rationality. The global reward is binary query-answer accuracy $r_{\text{ans}}$.

Two mechanisms constitute the method's core contributions.

Sequential MDP regularization

To train heterogeneous, asynchronously executing agents jointly, CoMAM casts the full pipeline as a three-step sequential MDP. States chain agent outputs into downstream inputs: $s_0 = \mathcal{H}$, $\mathcal{H}$0, $\mathcal{H}$1, $\mathcal{H}$2, with each action $\mathcal{H}$3 produced by a dedicated policy. This formulation embeds inter-agent dependencies into state transitions, allowing simultaneous gradient updates of all policies while each retains task specialization. The ablation results substantiate the design choice: heterogeneous policies consistently outperform a single shared policy, and MDP-based joint RL outperforms independent RL under both policy paradigms (e.g., 0.64/0.70/0.66 versus 0.57/0.61/0.59 accuracy at 32K/128K/1M for heterogeneous policies).

Adaptive credit assignment

Rather than distributing the global reward equally, CoMAM quantifies each agent's contribution via group-level ranking consistency. For each input, $\mathcal{H}$4 MDP trajectories are sampled; the NDCG between the ranking induced by agent $\mathcal{H}$5's local rewards and the ranking induced by global rewards yields a consistency score $\mathcal{H}$6, normalized via softmax into credit weights $\mathcal{H}$7. Each agent is then optimized with GRPO on the integrated reward $\mathcal{H}$8. The intuition is that agents whose local performance correlates strongly with global outcomes receive larger shares of the global signal, aligning local improvement with system performance while preserving local task gradients.

Experimental results

Evaluation uses the PersonaMem benchmark (over 180 long-term histories, roughly 6,000 multiple-choice queries across seven personalization-oriented question types) at 32K, 128K, and 1M token context lengths, with Qwen2.5 (3B/7B-Instruct) and Llama-3.x (3.2-3B/3.1-8B-Instruct) backbones. Results are averaged over at least three seeds.

Method	Qwen 32K	Qwen 128K	Qwen 1M	Llama 32K	Llama 128K	Llama 1M
Base (no memory)	0.41	0.39	0.38	0.35	0.31	0.32
RAG	0.48	0.45	0.41	0.43	0.39	0.36
CAM (prompt-based)	0.53	0.50	0.45	0.48	0.45	0.43
Mem1 (RL)	0.59	0.57	0.58	0.56	0.58	0.61
Memory-R1 (RL, independent)	0.58	0.60	0.60	0.57	0.61	0.60
CoMAM	0.64	0.70	0.66	0.62	0.68	0.69

CoMAM exceeds the strongest independent-optimization baseline (Memory-R1) by 8.5–16.7% relative across settings, and the gains hold across both backbone families and all seven question types, including fine-grained factual recall and preference-aligned recommendation. The consistent gap between RL-based and prompt-based memory systems (e.g., Memory-R1's 0.57–0.61 versus CAM's 0.45–0.53) independently confirms the value of RL for memory agents; the further gap to CoMAM isolates the contribution of joint optimization.

The credit assignment ablation is particularly informative: local-only rewards reach 0.57–0.65, global-only rewards 0.60–0.64, and equal local-global integration 0.62–0.65, while adaptive assignment achieves 0.64–0.70. Notably, adaptive credit assignment also improves each agent's local task reward relative to local-only training, indicating that global alignment does not come at the cost of task specialization. Sensitivity analysis shows the adaptive weights dominate any fixed weighting scheme. On efficiency, joint training reduces total convergence steps relative to the sum over independently trained agents (e.g., 68 versus 160 steps at 32K) because agents train in parallel; inference latency is unaffected by the training paradigm.

Limitations and open questions

The paper concedes three limitations. First, the agent set is deliberately compact; capabilities such as memory modification for consistency maintenance and redundancy deletion are deferred, leaving open how credit assignment scales to larger agent teams with more entangled dependencies. Second, the MDP regularization is static: training assumes fixed invocation ordering and frequencies, whereas deployed systems exhibit dynamic invocation patterns and real-time memory updates; the train-deployment distribution shift this introduces is not quantified. Third, local-local alignment between agents is modeled only implicitly through the shared global objective; whether explicit pairwise alignment terms would further improve collaboration remains an open question. A further implicit assumption worth noting is the reliance on ground-truth evidence $\mathcal{H}$9 for local rewards, extracted with a proprietary model and human-validated — the method's applicability where such evidence annotations are unavailable is not addressed.

Conclusion

CoMAM reframes multi-agent memory system optimization as a joint RL problem, combining sequential MDP regularization of heterogeneous agent execution with NDCG-based adaptive credit assignment to integrate local task rewards and global answer accuracy. The empirical evidence — consistent improvements over independent-optimization RL baselines across two backbone families and three context lengths, together with ablations isolating the contributions of MDP regularization and adaptive weighting — supports the paper's thesis that cross-agent collaboration must be optimized directly rather than assumed to emerge from local specialization. The main open questions concern scaling the credit assignment to richer agent sets, closing the train-deployment gap introduced by the static MDP, and relaxing the dependence on annotated evidence for local rewards.