PaperlayerView exact revision on arXiv

Source-linked AI summary

Mem-T: Densifying Rewards for Long-Horizon Memory Agents

Yanwei Yue, Boci Peng, Xuanbo Fan, Jiaxin Guo, Qiankun Li, Yan Zhang

arXiv:2601.23014v2cs.LGcs.CL

TL;DR

Long-horizon memory agents receive sparse, delayed rewards, making end-to-end optimization of memory management difficult. Mem-T combines a hierarchical memory agent with MoT-GRPO, which densifies feedback through tree-based backpropagation and hindsight credit assignment. Experiments report strong performance and favorable efficiency, including up to 14.92% improvement and approximately 24.45% lower inference tokens per query than the cited baselines.

  • Problem

    Long-horizon memory agents receive sparse, delayed rewards, hindering joint optimization of memory construction and retrieval.

  • Method

    Mem-T uses a hierarchical memory agent, while MoT-GRPO converts terminal feedback into dense supervision through memory operation trees and hindsight credit assignment.

  • Results

    Mem-T achieves state-of-the-art performance and a favorable accuracy-efficiency Pareto frontier across in-domain and out-of-domain benchmarks.

  • Takeaways & Limitations

    The reported results position attribution-centric, fully learnable memory management as a route toward self-evolving agents capable of lifelong learning.

Abstract

from arXiv · show

Memory agents, which depart from predefined memory-processing pipelines by endogenously managing the processing, storage, and retrieval of memories, have garnered increasing attention for their autonomy and adaptability. However, existing training paradigms remain constrained: agents often traverse long-horizon sequences of memory operations before receiving sparse and delayed rewards, which hinders truly end-to-end optimization of memory management policies. To address this limitation, we introduce Mem-T, an autonomous memory agent that interfaces with a lightweight hierarchical memory database to perform dynamic updates and multi-turn retrieval over streaming inputs. To effectively train long-horizon memory management capabilities, we further propose MoT-GRPO, a tree-guided reinforcement learning framework that transforms sparse terminal feedback into dense, step-wise supervision via memory operation tree backpropagation and hindsight credit assignment, thereby enabling the joint optimization of memory construction and retrieval. Extensive experiments demonstrate that Mem-T is (1) high-performing, surpassing frameworks such as A-Mem and Mem0 by up to $14.92\%$, and (2) economical, operating on a favorable accuracy-efficiency Pareto frontier and reducing inference tokens per query by $\sim24.45\%$ relative to GAM without sacrificing performance.

1. Introduction

Existing memory agents face sparse, delayed rewards across long-horizon operation sequences, limiting process-level training. Mem-T introduces a hierarchical, trainable framework with MoT-GRPO to jointly optimize memory construction and retrieval using dense supervision.

  • Motivation: Existing memory agents may execute hundreds of operations across approximately 500 turns before receiving a binary reward from sporadic QA accuracy signals.Their terminal reward is indiscriminately propagated across operations without dense supervision or process-level attribution.
  • Motivation: The central challenge is jointly optimizing memory construction and retrieval with dense rewards and accurate process-level attribution.
  • Mem-T: Mem-T is a streamlined hierarchical memory agent that supports memory formation, evolution, and multi-turn retrieval over dynamic information streams.Its architecture integrates factual, experiential, and working memory while orchestrating the memory-operation lifecycle.
  • MoT-GRPO: MoT-GRPO constructs memory operation trees, back-propagates outcome rewards to intermediate retrieval nodes, and assigns utility to source memories through hindsight credit assignment.These mechanisms provide dense process-level signals for retrieval and supervise corresponding formation and evolution operations.
  • Results: Up to 14.92% F1 gains and approximately 24.45% lower inference tokens per query are reported relative to the stated baselines.The evaluation describes state-of-the-art performance alongside a superior accuracy-efficiency Pareto frontier.

2. Related Work

Memory-agent research has progressed from heuristic systems toward agentic architectures and reinforcement-learning-based memory management. Existing work spans multiple memory functions and time scales, while comparison categories distinguish whether capabilities are absent, heuristic-based, or trainable.

  • Memory Agent Architectures: Memory-agent architectures have evolved from heuristic systems such as MemoryBank and MemGPT toward agentic systems including Mem0, MemOS, and A-Mem.
  • Memory Agent Architectures: Prior systems span factual memory for declarative knowledge, experiential memory for continual self-improvement, and working memory for task execution context.
  • Comparison Dimensions: Table 1 compares memory systems across factual, experiential, and working memory and across formation, evolution, retrieval, and process attribution stages.Its legend distinguishes capabilities that are not included, included heuristically, or included and trainable.
  • Reinforcement Learning for Memory Agents: Reinforcement learning is used to give language models adaptive memory-management capabilities across short-term working, long-term factual, and experiential memory.

3. Method

Mem-T models memory management as a dual-track process: continuous construction updates a hierarchical store, while on-demand retrieval searches it over multiple turns. MoT-GRPO trains these long-horizon operations with tree-based rewards and hindsight credit assignment.

  • Mem-T Workflow: Mem-T maintains a hierarchical memory state over a continuous information stream and combines working, factual, and experiential memory modules.Working memory preserves within-episode coherence, while long-term modules store factual and experiential information across sessions.
  • Memory Construction: During continuous construction, π_form extracts facts, strategies, raw data, or working summaries from each input chunk.The formation action space contains CrtFact, CrtExp, CrtRaw, and UpdWork.
  • Memory Construction: The evolution policy integrates each candidate memory through ADD, UPDATE, DELETE, or IGNORE actions applied to the existing store.Relevant existing memories are considered when selecting the evolution action and determining additions or removals.
  • Memory Retrieval: For each query, π_retr performs sequential searches over memory modules, accumulates relevant memories, and terminates with Finish before generating an answer.The retrieval policy conditions later actions on the query, retrieved-memory history, and reasoning state.
  • MoT-GRPO for Retrieval: MoT-GRPO constructs branching retrieval trees, assigns dense node rewards from validity, evidence density, and expected performance, and estimates both local and global advantages.The method uses iterative branching rollouts and combines intra-tree and inter-tree comparisons to guide retrieval-policy optimization.
  • MoT-GRPO for Construction: Hindsight credit assignment transfers downstream retrieval advantages to upstream memory-construction actions, linking construction quality to evidence alignment and retrieved-memory utility.The credit coefficient uses evidence alignment and a retrieval-trace signal weighted by λ=0.1.

4. Experiments

Experiments evaluate Mem-T across in-domain and out-of-domain long-context benchmarks, its accuracy-efficiency trade-off, parameter sensitivity, and memory-processing behavior. Mem-T achieves strong performance and cost-effectiveness, while MoT-GRPO and selected retrieval settings improve results with diminishing returns beyond moderate complexity.

  • Evaluation Setup: Mem-T is evaluated on LoCoMo, LongMemEval, HotpotQA, and NarrativeQA using Qwen3-4B and Qwen3-8B backbones.LoCoMo uses a 1:1:8 train-validation-test split, while the other three benchmarks are treated as out-of-domain datasets.
  • Main Results: 14.92 F1 points (34.13% ↑) improve Mem-T over baselines on LoCoMo with Qwen3-4B.With Qwen3-8B, the corresponding improvement is 14.55 F1 points (33.08% ↑).
  • Accuracy-Efficiency Trade-off: 19.94% ∼24.45% lower inference overhead per query accompanies Mem-T’s 5.17 ∼28.48 F1 Score improvement over GAM.Mem-T lies on the Pareto front for both LoCoMo and HotpotQA.
  • Sensitivity Analysis: F1 increases from 53.45 to 58.65 as maximum retrieval steps rise from 2 to 6, but further increases to 10 yield gains below 0.5% while token use rises from ∼9k to ∼21k.The experiments therefore use six maximum retrieval steps.
  • Sensitivity Analysis: Increasing operation trees G from 1 to 3 raises LoCoMo F1 from 54.20 to 58.65 and HotpotQA F1 from 62.49 to 66.54.Increasing G to 5 provides only a marginal average improvement of 0.35 while increasing computational cost by approximately 67%; the selected setting is G=3.
  • Case Study: Mem-T more accurately extracts information, distinguishes memory updates from additions, and performs deeper retrieval than the Qwen3-4B baseline.The case study contrasts these capabilities with baseline failures in timestamp resolution, memory evolution, and ambiguous raw-query retrieval.

5. Conclusion

Mem-T combines a hierarchical memory framework with MoT-GRPO, which converts sparse terminal rewards into dense step-wise supervision for joint memory construction and retrieval optimization. Experiments report strong benchmark performance and a favorable accuracy-efficiency trade-off.

  • Conclusion: Mem-T combines hierarchical memory with MoT-GRPO, using memory operation trees to provide dense step-wise supervision for joint construction and retrieval optimization.The framework is presented as an attribution-centric reinforcement-learning paradigm for memory agents.
  • Conclusion: Mem-T achieves state-of-the-art performance across in-domain and out-of-domain benchmarks while maintaining superior Pareto efficiency between task accuracy and inference overhead.The conclusion characterizes the approach as a shift from heuristic-based storage toward fully learnable, attribution-centric memory systems.

A.1. Dataset Description

The evaluation spans long-context conversational, multi-hop reasoning, sustained interactive memory, and narrative comprehension benchmarks, with datasets designed to test retrieval, reasoning, and memory abilities across challenging contexts.

  • LoCoMo: LoCoMo evaluates long-range memory and reasoning over extended conversations spanning dozens of sessions and hundreds of turns.Its questions cover single-hop, multi-hop, open-domain, temporal, and adversarial reasoning categories; the experiments exclude adversarial questions.
  • HotpotQA: HotpotQA tests multi-hop reasoning by requiring aggregation of evidence across multiple supporting documents.The 56K-token synthetic version embeds golden paragraphs among distractors, turning the task into long-range retrieval and inference.
  • LongMemEval: LongMemEval assesses sustained memory in prolonged user–assistant interactions through information extraction, multi-session reasoning, temporal reasoning, knowledge updates, and abstention.It uses 500 manually curated questions embedded in scalable chat histories.
  • NarrativeQA: NarrativeQA measures comprehension and reasoning over long narratives such as books and movie scripts, with questions requiring synthesis across global document structure.The dataset contains tens of thousands of human-written question–answer pairs associated with over a thousand story documents.

A.2. Implementation Details

The implementation uses distributed inference and training infrastructure, extended contexts, mixed precision, and memory-efficient parallelization for memory retrieval and construction.

  • Retrieval training: Retrieval training uses Ray with vLLM, XFormers, a global batch size of 32, and a peak learning rate of 5 × 10^-6.A KL divergence coefficient of 0.001 is used to support training stability and prevent reward hacking.
  • Retrieval training: Retrieval supports prompts up to 40,960 tokens and observation histories up to 20,480 tokens.FSDP with parameter, gradient, and optimizer offloading is used for computational efficiency in bfloat16 precision.
  • Memory construction: Memory construction training runs on one 8-GPU node using LLaMA-Factory, DeepSpeed ZeRO-3, and Flash Attention 2.The maximum sequence length is truncated to 6,144 tokens.
  • Memory construction: Memory construction uses a global batch size of 32, 200 optimization steps, cosine scheduling, and a peak learning rate of 5 × 10^-6.The setup uses bfloat16 precision and reserves 10% of the dataset for validation.

B.1. Generalization Experiments Across Other LLMs

Mem-T generalizes to the Qwen3-8B backbone, outperforming existing baselines across most LoCoMo metrics, while training further improves its performance.

  • Generalization results: Mem-T significantly outperforms existing baselines across most LoCoMo metrics when applied to Qwen3-8B.The comparison uses F1 and BLEU-1 as evaluation metrics.
  • Generalization results: 50.08 Overall F1 is achieved by the training-free variant, surpassing previously trained models including Memory-R1 and MemAgent.This result demonstrates that the hierarchical memory system transfers effectively even without additional training.
  • Generalization results: 58.53 F1 is achieved when Mem-T is combined with the proposed training, with particularly strong performance on Single-Hop and Temporal reasoning tasks.The result supports transferability across different LLM backbones.

B.2. Sensitivity Analysis

Sensitivity analysis shows that expanding more memory-operation branches improves performance up to a moderate breadth, while further expansion produces diminishing returns.

  • Branch expansion: Figure 7 analyzes sensitivity to the number of nodes selected for branch expansion during MoT-GRPO training on LoCoMo and HotpotQA.The figure supports comparison of expansion breadth across the two datasets.
  • Branch expansion: 56.08→58.65 LoCoMo F1 and 64.89→66.35 HotpotQA F1 result from increasing branch-expansion nodes from 1 to 3.Increasing the expansion breadth beyond 3 nodes yields diminishing returns, with both datasets plateauing or slightly declining at 5 nodes.

B.3. Training Curves

MoT-GRPO shows improving retrieval rewards and rapidly converging construction loss. Retrieval rewards rise despite raw variance, while construction loss stabilizes at a low level after an initially sharp decline.

  • Memory retrieval training: Retrieval rewards increase from approximately 0.30 to over 0.50 by step 200 under MoT-GRPO.The smoothed reward curve rises consistently, although raw rewards remain substantially variable.
  • Memory construction training: Construction loss drops from 1.8 to below 0.4 within the first 50 of 180 steps.The sharp initial descent indicates rapid convergence during memory construction training.
  • Memory construction training: Construction loss later fluctuates marginally around 0.25, indicating stable memory synthesis and state-update training.The stabilized curve suggests the model captured underlying patterns for memory synthesis and state updates.
Loading 2601.23014v2…