Tree-based Credit Assignment for Multi-Agent Memory System

TL;DR

This paper introduces TreeMem, a tree-structured credit assignment method for multi-agent memory systems in LLMs, enabling effective agent-specific learning from final rewards without task-specific annotations.

Contribution

It proposes a novel tree-based credit assignment approach that derives agent-specific signals from final rewards, improving multi-agent memory system training without requiring costly annotations.

Findings

TreeMem outperforms strong baselines on long-horizon benchmarks.

The method effectively assigns credit to individual agents without task-specific annotations.

Experimental results validate the approach's ability to enhance memory system performance.

Abstract

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.