Technical Summary: Tree-based Credit Assignment for Multi-Agent Memory System (TreeMem)

Problem Statement

Memory systems are increasingly utilized to extend Large Language Models (LLMs) beyond their context windows for long-horizon tasks, typically organized as multi-agent pipelines involving Builder, Summarizer, and Retrieval agents. While Reinforcement Learning (RL) has been explored to optimize these agents, existing paradigms face significant limitations:

1. Uniform Final Rewards: Applying the final downstream task reward (e.g., QA accuracy) uniformly to all agents provides coarse and ambiguous supervision. It fails to distinguish the specific contributions of heterogeneous agents, limiting their ability to specialize.
2. Task-Specific Rewards: Designing intermediate rewards for specific subtasks (e.g., evidence coverage for Builders, abstraction quality for Summarizers) requires costly manual annotations and is prone to unreliability. Furthermore, handcrafted rewards can induce "reward hacking," where agents optimize the metric without improving downstream performance.

There is a lack of methods that can derive fine-grained, agent-specific credit signals directly from the final downstream reward without relying on expensive task-specific annotations.

Methodology: TreeMem

The authors propose TreeMem, a tree-based RL training framework that derives agent-specific credit from the final reward through a tree-structured rollout and Monte Carlo averaging.

1. Tree-Structured Rollout

Instead of a linear pipeline, TreeMem expands the multi-agent workflow into a tree structure during training:

• Builder Stage: Given a history $H$, the Builder agent samples $G$ candidate memory actions ($a_1^i$).
• Summarizer Stage: For each builder action, the Summarizer samples $J$ branches of summarized memory ($a_2^{i,j}$).
• Retrieval Stage: For each builder-summarizer pair, the Retrieval agent samples $K$ final responses ($a_3^{i,j,k}$).
  This results in $G \times J \times K$ complete trajectories, forming a tree of candidate paths from raw history to final answers.

2. Credit Assignment via Monte Carlo Averaging

TreeMem assigns credit to intermediate actions by propagating the final reward upward through the tree:

• Retrieval Agent: Receives the direct final reward $R$ of its generated answer.
• Summarizer Agent: Its credit is estimated by averaging the final rewards over the $K$ response branches dependent on its specific summary.
• Builder Agent: Its credit is estimated by averaging the final rewards over all $J \times K$ branches dependent on its specific memory construction, with an added length penalty to discourage verbatim copying of long histories.

This process converts a sparse, final-only signal into granular, agent-specific optimization signals.

3. Policy Optimization

The framework constructs a training group by randomly selecting one complete trajectory from each of the $G$ subtrees rooted at the builder actions. These selected trajectories, paired with their computed agent-specific credits, are used to update the policies of the Builder, Summarizer, and Retrieval agents simultaneously. The optimization follows a group-relative principle (inspired by GRPO), normalizing credits within the group to compute advantages. This allows heterogeneous agents to specialize effectively while remaining aligned with the global task objective. Notably, the tree expansion is used only during training; at inference, the system operates with the optimized policies directly, without executing tree branches.

Key Contributions

• Novel Framework: Introduction of TreeMem, the first framework to derive agent-specific credit for multi-agent memory systems directly from final rewards without task-specific annotations.
• Tree-Structured Estimation: A mechanism to model the multi-agent pipeline as a tree, using Monte Carlo averaging to isolate the impact of intermediate agent actions on the final outcome.
• Specialization without Annotation: Demonstrates that fine-grained credit assignment can be achieved without the costly manual labeling required by previous task-specific reward methods.
• Simultaneous Optimization: A strategy to update heterogeneous agent policies jointly using subtree-estimated credits, promoting effective role specialization.

Experimental Results

TreeMem was evaluated on long-horizon benchmarks including PersonaMem (32K, 128K, 1M token contexts), LongMemEval, and LOCOMO, using Qwen and Llama backbone models.

• Performance: TreeMem consistently outperformed strong baselines, including context-based methods, prompt-driven memory systems, and existing RL-optimized methods (e.g., Mem1, Memory-R1, CoMAM). On PersonaMem, it achieved an average improvement of +7.5% over the best RL baseline.
• Ablation Studies:
  • Reward Signals: TreeMem outperformed variants using only final rewards or manually designed task-specific rewards, validating that tree-based credit assignment provides superior training signals.
  • Agent Impact: Removing the tree-based signal for any single agent (Builder or Summarizer) resulted in performance degradation, confirming that all agents benefit from agent-specific credit.
• Efficiency: While the tree rollout introduces computational overhead during training, the method accelerates convergence (requiring fewer steps) and improves token efficiency (producing more compact memories). At inference, it incurs no extra overhead compared to standard baselines.
• Hyperparameter Sensitivity: The system demonstrated robust performance across a wide range of group sizes ($G$) and branch counts ($J, K$), requiring no exhaustive grid search.

Significance and Claims

The paper claims that TreeMem addresses a critical bottleneck in multi-agent memory optimization: the inability to provide specific training signals to heterogeneous agents without expensive annotations. By transforming coarse final rewards into fine-grained, agent-specific signals through tree-structured rollouts, TreeMem enables memory agents to specialize effectively in their respective roles (building, summarizing, retrieving).

The authors assert that this approach validates the importance of fine-grained credit assignment in complex multi-agent LLM systems. The results suggest that TreeMem can significantly enhance the performance of long-horizon memory systems, making them more reliable and efficient for tasks requiring dynamic information storage and retrieval over extended interactions. The work positions itself as a practical solution that balances system-level alignment with the need for specialized agent behavior, avoiding the pitfalls of both uniform supervision and unreliable manual reward engineering.