TreeMem: Tree-Based Credit Assignment for Multi-Agent Memory Systems

Paper: *Tree-based Credit Assignment for Multi-Agent Memory System* Authors: Joint team from multiple Chinese universities (incl. Wenyu Mao et al.) Link: https://arxiv.org/abs/2605.04811

The "Shared Pot" Problem

Consider a three-agent memory pipeline:

• Builder: collects and structures raw information into memory chunks
• Summarizer: condenses memories into summaries/indexes
• Retriever: fetches relevant information to answer questions

Their only KPI is final QA accuracy. When an answer is right or wrong, it's impossible to tell which agent was responsible — yet existing fixes both have flaws:

1. Split the final reward equally — the signal is too coarse; agents may free-ride. 2. Design per-agent rewards — requires costly, hard-to-define human annotation of intermediate metrics.

Core Idea: Turn the Pipeline into a Tree

Instead of a linear Builder → Summarizer → Retriever → Answer pipeline, TreeMem lets each agent branch into multiple possible outputs (multiple memory constructions, multiple summaries, multiple retrieval strategies), forming a tree over the joint action space.

Monte Carlo Credit Assignment

To estimate the credit of agent *i*'s action:

1. Fix that action, sample multiple branches from subsequent agents. 2. Average the final reward over all leaf nodes. 3. That average is the action's credit value:

\[\text{Credit}(a_i) = \mathbb{E}_{\pi_{i+1}, ..., \pi_n}[R | a_i]\]

All agents are then updated simultaneously with their own fine-grained signals derived purely from the final reward — no human-labeled intermediate rewards needed.

Preventing Reward Hacking

Because credit is based on average downstream branch performance, a Builder that lazily forwards unfiltered raw data sees its average reward drop (redundancy hurts the Summarizer). The optimal strategy becomes genuinely doing its own filtering job rather than pushing responsibility downstream.

Results

| Benchmark | TreeMem | Mem0 | CoMAM | Gain | |-----------|---------|------|-------|------| | PersonaMem | 72.3% | 67.1% | 65.8% | +5.2% | | LongMem | 68.7% | 63.4% | 62.1% | +5.3% | | MultiHop-QA | 81.2% | 75.6% | 74.3% | +5.6% |

Key findings:

• Advantage widens to ~8% on million-token conversations, where upstream errors propagate more severely.
• Memory compression improves too: agents learn to filter and summarize more precisely.
• Ablations: removing tree structure costs 12%, removing Monte Carlo sampling costs 8%, removing hierarchical credit costs 6%.
• Training visualization shows the Builder evolving from conservative retention → imperfect filtering → precise noise filtering.

Comparison with Alternatives

| Approach | Reward design | Needs annotation | Credit granularity | Scalability | |----------|--------------|------------------|--------------------|-------------| | Shared final reward | shared | No | Very coarse | Good | | Per-task rewards | per-agent | Yes (costly) | Fine | Poor | | TreeMem | auto-derived | No | Medium | Good |

Limitations

• Compute cost: 10 branches per agent over a 3-layer tree implies ~1000 forward passes per estimate; suitable for offline training, not real-time learning.
• Tree depth: only 3 layers validated; deeper pipelines risk exponential branching (possible fixes: hierarchical subtrees, adaptive sampling).
• Delayed credit: full rollouts must complete before credit is computed, problematic for real-time feedback.

Broader Implications

Multi-agent credit assignment echoes classic problems in RL (MADDPG, QMIX), team incentive design, and KPI decomposition. TreeMem's recipe — hierarchical decomposition of the joint action space + Monte Carlo expectation estimation + end-to-end optimization — extends beyond memory systems to multi-agent code generation, scientific research, and content creation pipelines.

Bottom line: TreeMem automatically derives per-agent credit signals from a single final reward, yielding 5–8% gains on long-context benchmarks while discouraging reward hacking — no intermediate reward annotation required.