Collaborative Multi-Agent Optimization for Personalized Memory System

1. Problem Statement

Personalized Large Language Models (LLMs) rely on memory systems to overcome context window limitations and retain long-term user conversational history. Current memory systems typically employ multiple specialized agents (e.g., for constructing fine-grained memories, coarse-grained profiles, and retrieving information).

The Core Issue: Existing methods optimize these agents independently (via prompt engineering or separate reinforcement learning). This approach treats the system as a "collection of local experts" rather than a collaborative team.

• Local-Global Misalignment: Optimizing an agent for its local task (e.g., maximizing information coverage) without considering downstream agents can lead to storing redundant or noisy data that hinders retrieval. Conversely, retrieval agents trained without knowledge of construction strategies may over-filter critical details.
• Performance Gap: As illustrated in the paper's Figure 1, independently optimized agents often yield lower global system performance (query-answer accuracy) compared to jointly optimized systems, even if individual local tasks appear converged.
• Challenges: Joint optimization faces two main hurdles:
  1. Heterogeneity & Asynchrony: Agents operate with different configurations and execution frequencies (construction is periodic; retrieval is query-driven).
  2. Credit Assignment: It is difficult to determine how much each agent contributes to the final global outcome, leading to ambiguous reward signals if global rewards are distributed equally.

2. Methodology: CoMAM Framework

The authors propose CoMAM (Collaborative Reinforcement Learning Framework for Multi-Agent Memory Systems), which jointly optimizes heterogeneous agents through three key mechanisms:

A. Task-Specific Agent Definition

The system utilizes three specialized agents:

1. Extraction Agent: Constructs fine-grained memories ($M_f$) from raw history (e.g., key events, facts).
2. Profile Agent: Abstracts coarse-grained memories ($M_c$) from $M_f$ (e.g., user preferences, behavioral patterns).
3. Retrieval Agent: Retrieves relevant memories from the set $M = \{M_f, M_c\}$ to answer user queries and generate responses.

B. Execution Trajectory Regularization (MDP Formulation)

To handle heterogeneity and asynchrony, the authors model the entire pipeline as a sequential Markov Decision Process (MDP):

• State Space ($S$): Captures the system state at each step (Raw History $\to$ Fine Memory $\to$ Coarse Memory $\to$ Retrieved Info $\to$ Final Response).
• Action Space ($A$): The specific output of each agent (extraction, abstraction, retrieval/generation).
• Transitions ($T$): Embeds inter-agent dependencies; the output of one agent becomes the input state for the next.
• Rewards ($R$):
  • Local Rewards: Task-specific metrics (e.g., information coverage for extraction, abstraction rationality for profiling, retrieval precision).
  • Global Reward: Query-answer accuracy ($r_{ans}$).

This formulation allows for end-to-end optimization where agents learn to adapt their local strategies to benefit the global pipeline.

C. Adaptive Credit Assignment Mechanism

To solve the credit assignment problem and align local improvements with global performance, CoMAM introduces a dynamic weighting strategy:

1. Ranking Consistency: For a batch of trajectories, the system calculates the Normalized Discounted Cumulative Gain (NDCG) between an agent's local reward ranking and the global reward ranking.
2. Adaptive Weights: Agents with higher consistency (i.e., those whose local improvements correlate strongly with global success) receive higher weights ($w_n$).
3. Integrated Reward: The final reward for an agent is a weighted sum of its local task reward and the global system reward:
   $$r_{final, n} = r_{local, n} + w_n \cdot r_{global}$$
   This ensures agents are rewarded not just for doing their specific job well, but for doing it in a way that helps the whole system.

The optimization algorithm used is Group Relative Policy Optimization (GRPO), which eliminates the need for a separate critic model.

3. Key Contributions

1. Joint Optimization Framework: CoMAM is the first framework to explicitly model multi-agent memory systems as a sequential MDP to facilitate end-to-end collaborative training, moving beyond independent agent optimization.
2. Adaptive Credit Assignment: The paper proposes a novel mechanism using NDCG-based ranking consistency to dynamically allocate global rewards, effectively solving the "local vs. global" alignment problem without manual reward shaping.
3. Handling Heterogeneity: The MDP regularization allows heterogeneous agents (with different tasks and execution frequencies) to be trained simultaneously in a unified pipeline.
4. Comprehensive Evaluation: Extensive experiments on the PersonaMem benchmark demonstrate the framework's superiority over state-of-the-art baselines.

4. Experimental Results

The authors evaluated CoMAM on the PersonaMem benchmark (180+ long-term interaction histories, 6,000+ queries) across three context lengths (32K, 128K, 1M tokens) using Qwen and Llama backbones.

• Overall Performance (RQ1): CoMAM outperformed all baselines (including prompt-based systems like CAM and RL-based systems like Memory-R1).
  • On the 128K setting with Llama, CoMAM achieved 70% accuracy, a 16.7% improvement over the previous best.
  • It consistently outperformed independent optimization baselines across all context lengths.
• Ablation Studies:
  • Agent Necessity (RQ2): Removing any single agent (Extraction, Profile, or Retrieval) caused significant performance drops, confirming the complementary nature of the multi-agent design.
  • MDP Regularization (RQ3): "MDP RL" (joint optimization) significantly outperformed "Independent RL" and "Single Policy" approaches, validating the need for modeling inter-agent dependencies.
  • Credit Assignment (RQ4): The adaptive mechanism ("Ours") outperformed fixed-weight strategies (equal allocation) and strategies using only local or only global rewards, proving the efficacy of dynamic weight allocation.
• Efficiency (RQ5): While joint training requires parallel updates, the total convergence steps are fewer than the sum of independent training steps, indicating higher training efficiency. Inference latency remains comparable to independent systems.

5. Significance

• Paradigm Shift: The paper challenges the prevailing "independent optimization" paradigm in multi-agent LLM systems, demonstrating that collaborative training is essential for achieving global optimality in complex pipelines.
• Scalability: The framework is designed to handle long-context scenarios (up to 1M tokens), addressing a critical bottleneck in personalized LLM applications.
• Generalizability: The concepts of MDP regularization for asynchronous agents and NDCG-based adaptive credit assignment can be applied to other multi-agent systems beyond memory management, such as complex reasoning chains or tool-use agents.
• Practical Impact: By improving the accuracy of personalized responses in long-term conversations, CoMAM directly enhances the user experience in AI assistants, customer support bots, and personal digital companions.