Joint Optimization of Multi-agent Memory System

Wenyu Mao; Haoyang Liu; Haosong Tan; Yaorui Shi; Jiancan Wu; An Zhang; Xiang Wang

cs.MAMar 13, 2026

Joint Optimization of Multi-agent Memory System

Authors: Wenyu Mao, Haoyang Liu, Haosong Tan, Yaorui Shi, Jiancan Wu, An Zhang, Xiang Wang

Organizations: University of Science and Technology of China · Institute of Dataspace, Hefei Comprehensive National Science Center

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.

Figures & tables

Figure 1: ( Left ) Training curves of coarse/fine memory construction and retrieval agents, which are optimized independently. ( Middle ) Global system performance under independent and joint optimization across different context lengths. ( Right ) Conceptual comparison between existing methods that optimize agents independently with shared global rewards and our joint optimization with adaptive credit assignment.

Figure 2: The overview of our joint optimization framework for the multi-agent memory system, CoMAM, which promotes collaboration among agents via end-to-end reinforcement learning and an adaptive credit assignment mechanism.

Method	Qwen			Llama
Method	32K	128K	1M	32K	128K	32M
	No memory
Base	0.41	0.39	0.38	0.35	0.31	0.32
RAG [ 14 ]	0.48	0.45	0.41	0.43	0.39	0.36
	Prompt-based memory
CAM [ 15 ]	0.53	0.50	0.45	0.48	0.45	0.43

Table 1: Query-Answering Accuracy on the PersonaMem Benchmark across 32K, 128K, and 1M context lengths. Bold indicates the best performance, and underlining denotes the second best.

Method	$F_{1}\uparrow$	$B_{1}\uparrow$	$J_{1}\uparrow$
RAG	18.27	14.57	22.20
SeCom	19.03	16.57	10.00
A-Mem	41.55	36.58	54.80
Mem0	38.44	34.53	46.80
GAM	51.75	48.96	52.00
Ours	62.17	51.88	57.60

Table 2: Results on LongMemEval using Qwen2.5-7B.

Table 5

Figure 3: Sensitivity of CoMAM to the hyperparameters

\alpha

and

\beta

as introduced in Sec. 4.1 .

Figure 4: The left demonstrates CoMAM’s sensitivity to the credit assignment weight, and the right shows the credit assignment’s impact on each local agent’s performance on the 32K setting.

Table 8

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	32K	128K	1M
#Train	313	1,744	1,830
#Validate	100	200	200
#question types	7	7	7

Appendix

Table 7: The statistics of queries after the train-validation split across different dataset settings.

Dataset	32K	128K	1M
batch_size	128	128	128
group_size	8	8	8
max_model_length	32k	32k	32k
epochs	5	5	5
max_output_tokens	3k	3k	3k

Appendix

Table 8: Hyperparameters of CoMAM on PersonaMem benchmark across different context lengths.

Figure 5: The curves for local task-specific reward

r_{n}

and global accuracy reward

r_{3}

Figure 6: Detailed performance of different methods across seven question types on the PersonaMem benchmark across three history length settings ( i.e., 32K, 128K, and 1M). The number of question types is consistent with those listed in Section 5.1

Memory-Augmented Language Model Agents

Memory-augmented language model agents use persistent facts, interaction histories, and skills to support personalization and long-horizon tasks.

See all 443 papers

Tree-based Credit Assignment for Multi-Agent Memory System
May 6, 2026Marina Mao, Alexandr Liu, Pengbo Li +3Credit Assignment Long-Horizon LLM Agents
University of Science and Technology of China · 2LLM Department, Tencent · 3The Hong Kong University of Science and Technology +1
CoMem: Collective-Individual Memory Synergy for Evolutionary Multi-Agent Systems
Sep 14, 2026Chengxin Yu, Zhaoxin Fan, Faguo Wu +3Multi-Agent LLM Systems Multi-Agent Collaboration
Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing, School of Artificial Intelligence, Beihang University · Beijing academy of blockchain and edge computing · National University of Defense Technology +1
AMA: Adaptive Memory via Multi-Agent Collaboration
Jan 28, 2026Weiquan Huang, Zixuan Wang, Hehai Lin +6Long-Term Conversational Memory Multi-Agent LLMs
The Hong Kong University of Science and Technology (Guangzhou) · Shandong University · Nanyang Technological University +1

May 6, 2026J/K move · Enter open · S save

Tree-based Credit Assignment for Multi-Agent Memory System

Marina Mao, Alexandr Liu, Pengbo Li, Siheng Li+2

University of Science and Technology of China · 2LLM Department, Tencent · 3The Hong Kong University of Science and Technology+1

Credit Assignment Long-Horizon LLM Agents Multi-Agent Reinforcement Learning Memory-Augmented Language Model Agents

Open paper AskPDF arXiv

Memory systems are widely adopted to enhance LLMs for long-horizon tasks, and are commonly organized as multi-agent pipelines with memory building, summarizing, and retrieval agents. To empower this system, existing RL-based methods either apply final downstream task rewards (e.g., QA accuracy) for all agents uniformly, which are coarse and ambiguous, or design task-specific rewards for agents on different subtasks, which require costly annotations (e.g., key evidence) and are difficult to define reliably. To address these limitations, we propose Tree-based Credit Assignment for Multi-Agent Memory Systems (TreeMem), which derives agent-specific credit from the final reward without task-specific annotations. Specifically, TreeMem extends the multi-agent pipeline (builder--summarizer--retrieval) into a tree structure, where each agent's outputs are expanded into multiple subsequent branches. The contribution of each agent is estimated via Monte Carlo averaging over its subsequent branches, capturing how intermediate agent actions may influence the final reward. This converts the coarse final reward into agent-specific optimization signals. These signals are then used to update all agent policies simultaneously, helping heterogeneous agents specialize effectively. Experiments on long-horizon benchmarks show that TreeMem improves memory system performance over strong baselines, validating the effectiveness of tree-structured credit assignment for the multi-agent memory system.

Abstract

Memory-Augmented Language Model Agents

Explore similar work

Tree-based Credit Assignment for Multi-Agent Memory System

CoMem: Collective-Individual Memory Synergy for Evolutionary Multi-Agent Systems

AMA: Adaptive Memory via Multi-Agent Collaboration