PaperlayerView exact revision on arXiv

Source-linked AI summary

MemBuilder: Reinforcing LLMs for Long-Term Memory Construction via Attributed Dense Rewards

Zhiyu Shen, Ziming Wu, Fuming Lai, Shaobing Lian, Yanghui Rao

arXiv:2601.05488v4cs.CL

TL;DR

Long-term dialogue systems need memory that tracks evolving historical states, while existing methods rely on static prompting or ineffective sparse-reward training. MemBuilder trains a 4B model to construct multi-dimensional memory with synthetic session-level rewards and contribution-aware gradient weighting. It reaches state-of-the-art results across three benchmarks, including 84.23% on LoCoMo, and generalizes to out-of-distribution dialogue settings.

  • Problem

    Existing long-term memory methods rely on static prompting or sparse rewards, while global rewards obscure which memory components contribute to downstream outcomes.

  • Method

    MemBuilder trains one lightweight 4B model to manage Core, Episodic, Semantic, and Procedural memory using synthetic session-level rewards and contribution-aware gradient weighting.

  • Results

    MemBuilder achieves state-of-the-art performance across LoCoMo, LongMemEval, and PerLTQA, including 84.23% on LoCoMo and strong out-of-distribution generalization.

  • Takeaways & Limitations

    A properly trained lightweight open-source model can replace expensive closed-source memory-construction APIs within the evaluated long-term dialogue benchmarks.

  • Takeaways & Limitations

    Evaluation uses Claude 4.5 Sonnet for question answering and Claude 4.5 Opus for synthetic question generation, whose generated QA pairs may contain inaccuracies or ambiguities.

Abstract

from arXiv · show

Maintaining consistency in long-term dialogues remains a fundamental challenge for LLMs, as standard retrieval mechanisms often fail to capture the temporal evolution of historical states. While memory-augmented frameworks offer a structured alternative, current systems rely on static prompting of closed-source models or suffer from ineffective training paradigms with sparse rewards. We introduce MemBuilder, a reinforcement learning framework that trains models to orchestrate multi-dimensional memory construction with attributed dense rewards. MemBuilder addresses two key challenges: (1) Sparse Trajectory-Level Rewards: we employ synthetic session-level question generation to provide dense intermediate rewards across extended trajectories; and (2) Multi-Dimensional Memory Attribution: we introduce contribution-aware gradient weighting that scales policy updates based on each component's downstream impact. Experimental results show that MemBuilder enables a 4B-parameter model to outperform state-of-the-art closed-source baselines, exhibiting strong generalization across long-term dialogue benchmarks.

1 Introduction

Long-term dialogue memory must represent evolving historical states, but existing approaches face sparse rewards and ambiguous credit assignment across memory components. MemBuilder addresses these issues with session-level dense rewards and contribution-aware updates for multi-dimensional memory construction.

  • RAG treats retrieved information as independent static chunks, limiting its ability to represent evolving information and superseded historical facts.
  • Memory-augmented frameworks decompose information into timestamped events and structured semantic units before storage, shifting work from inference-time processing to retrieval.
  • End-of-trajectory rewards provide insufficient credit assignment because models cannot identify which session’s memory operations produced the final outcome.
  • Global rewards fail to distinguish memory operations across components with different downstream impacts.
  • MemBuilder trains a lightweight 4B model over Core, Episodic, Semantic, and Procedural memory using synthetic session-level questions and contribution-aware gradient weighting.Session rewards provide immediate feedback after each session, while update magnitudes scale with each component’s retrieval usage.

2 Related Works

Related work spans benchmark-driven long-term dialogue evaluation, structured external-memory frameworks, and training-based memory management. Training approaches differ in memory representation and reward design, with implicit methods trading interpretability and controllability for efficiency and explicit methods exposing distinct reinforcement-learning limitations.

  • Long-term memory benchmarks evaluate LLM agents through multi-session question answering, temporal reasoning, and evolving user profiles.
  • Structured external-memory frameworks organize information around distinctions such as episodic, semantic, and procedural memory.Examples include hierarchy-based, personalized-memory, and multi-dimensional prompting implementations.
  • Latent-memory methods encode information into hidden states or latent tokens, while decoupled architectures use separate memory encoders and retrievers.These implicit representations are efficient but sacrifice interpretability and fine-grained controllability.
  • Explicit-memory training includes reinforcement learning with sparse trajectory-level rewards and multi-dimensional construction with global rewards shared across operations.The latter does not account for differing downstream impacts of individual memory operations.

3 Methodology

MemBuilder organizes long-term dialogue information into four specialized memory types and trains their construction with session-level QA rewards and contribution-aware updates. The architecture preserves temporal evolution while reinforcement learning optimizes memory operations for downstream question answering.

  • Problem: The task is long-term dialogue question answering over information distributed across sessions that cannot typically fit into one context.An external memory bank compresses and organizes historical information for selective retrieval.
  • Memory Architecture: The memory bank contains Core, Episodic, Semantic, and Procedural components processed simultaneously from each new session.Core stores persistent profiles, Episodic records timestamped events, Semantic stores user-specific facts, and Procedural stores routines and workflows.
  • Memory Architecture: Core Memory is always included in context, while Episodic, Semantic, and Procedural memories are stored in a vector database and retrieved by semantic similarity.Core is maintained as a fixed block with automatic compression when capacity is exceeded.
  • Memory Operations: UPDATE preserves temporal history by creating a timestamped entry that references its predecessor instead of overwriting it.MERGE synthesizes related events across a time range while retaining references to the original events as evidence.
  • Dense Rewards: Synthetic session-level QA evaluates memory quality at each step, replacing a single end-of-trajectory signal with denser learning feedback.For each session, retrieved memories support expert-generated question-answer pairs; rollouts are evaluated by answer correctness against ground truth.

4 Experiments

MemBuilder is evaluated across three long-term dialogue benchmarks, with ablations examining training stages, attributed gradient weighting, reward density, answer-model transfer, question types, and memory-operation behavior. The method achieves strong benchmark performance, generalizes beyond its training source, and benefits from dense rewards and contribution-aware weighting.

  • Main Results: MemBuilder achieves SOTA performance across LongMemEval, LoCoMo, and PerLTQA, with accuracies of 85.75%, 84.23%, and 93.14%, respectively.On LoCoMo, it surpasses MIRIX by 6.75 percentage points and Claude 4.5 Sonnet as the memory construction model.
  • Main Results: Training exclusively on LongMemEval still yields strong performance on LoCoMo and PerLTQA, whose dialogue structures and question types differ from the training data.These benchmarks serve as out-of-distribution test sets.
  • Training Ablation: SFT raises performance from 68.07% to 81.74%, while subsequent RL increases it to 84.23%; RL without SFT reaches only 79.31%.The ablation identifies complementary roles for valid multi-dimensional outputs and downstream QA optimization.
  • Gradient Weighting: Gradient weighting improves LoCoMo performance from 82.82% at α = 1 to 84.04% at α = 4, while larger coefficients degrade performance through gradient imbalance.The optimal coefficient is α = 4, and retrieval frequency matches the actual information source in 95% of human-evaluated cases.
  • Question Types: The model gains most over Memory-R1 on MultiHop and Temporal questions, reaching 77.88% versus 53.01% and 71.71% versus 51.55%, respectively.It also achieves 90.58% on Adversarial questions.

5 Conclusion

MemBuilder trains a reinforcement-learning framework for multi-dimensional memory construction in long-term dialogues. With ADRPO, Qwen3-4B achieves 84.23% on LoCoMo and generalizes effectively to out-of-distribution benchmarks.

  • MemBuilder is a reinforcement-learning framework for multi-dimensional memory construction in long-term dialogues.
  • 84.23% on LoCoMo is achieved by Qwen3-4B with ADRPO.
  • The trained lightweight open-source model surpasses prompting-based frameworks using expensive closed-source models.
  • The method generalizes effectively to out-of-distribution benchmarks.

Limitations

The evaluation spans three long-term dialogue benchmarks with different formats, domains, and data organization. Key limitations include reliance on a fixed closed-source answer model and possible inaccuracies in synthetic QA pairs.

  • Limitations: Evaluation relies on a fixed closed-source Claude 4.5 Sonnet answer model, while synthetic QA pairs may contain occasional inaccuracies or ambiguities.
  • The evaluation uses three long-term dialogue benchmarks differing in dialogue format, domain, and data organization.
  • Benchmark scope: LongMemEval contains 500 independent questions averaging 40 sessions and approximately 115K tokens, testing five memory abilities.
  • Benchmark scope: LoCoMo contains 10 human-human dialogues averaging 27 sessions and 14K tokens, with 1,986 questions across five types.
  • Benchmark scope: PerLTQA contains 141 fictional characters and 8,593 questions requiring integration of Episodic and Semantic Memories.
  • Benchmark scope: LoCoMo and PerLTQA serve as out-of-distribution test sets relative to LongMemEval.

B.2 RL Training

ADRPO trains a multi-dimensional memory system using contribution-aware gradient weighting and session-level rewards. The system assigns role-specific agents to Core, Episodic, Semantic, and Procedural memories with structured validity checks and retrieval operations.

  • RL optimization: ADRPO extends verl with contribution-aware gradient weighting and session-level reward computation.
  • Retrieval: The framework retrieves top-20 memories for construction, top-10 memories for QA, and top-5 chunks for RAG baselines.
  • Memory components: Core memory is a 5,000-character block that the policy compresses when capacity is exceeded, preserving essential profile information.
  • Memory components: Episodic, Semantic, and Procedural memories are stored as independently embedded entries in a vector database without an explicit size limit.
  • Agent actions: Each memory type receives a role-specific prompt, and the model outputs JSON specifying the action type and content.
  • Agent actions: Actions are valid only when JSON structure, action type, required fields, and referenced entries satisfy the defined criteria.

D.3 Training-based Methods

The training-based comparison reports that MemBuilder performs strongly against Memory-R1 under the stated evaluation setup. An ablation varies contribution-aware weighting on a reduced training set to study its effect.

  • 82.00% on LoCoMo is achieved by Qwen3-4B as both construction and answer model, versus Memory-R1’s reported 62.74%.
  • The comparison is not identical across methods because Memory-R1 uses Llama-3.1-8B-Instruct for both memory construction and answering.
  • The contribution-aware weighting ablation varies α ∈{1, 2, 4, 8, 16} on 10 dialogues sampled from 50 RL training dialogues.

E.2 Reward Density Ablation

The ablation varies the density of session-level task rewards while keeping training conditions comparable. It also documents the pipeline’s separate data-preparation and reinforcement-learning cost stages.

  • The ablation uses the full training set of 50 dialogues.
  • Reward density is varied by randomly skipping task-reward computation for a fraction of sessions.At density 1/d, each session independently receives a task reward with probability 1/d.
  • All configurations train for the same number of epochs, while unrewarded sessions retain format-validity and length-penalty signals.

F.1 Data Preparation Costs (One-Time)

The pipeline separates one-time data preparation from per-run reinforcement-learning costs. Preparation includes expert demonstrations and synthetic questions, while training evaluates generated answers and judgments at scale.

  • Expert trajectory generation covers 50 LongMemEval conversations and 2,400 sessions, producing 9,600 memory-agent API calls.Each session invokes Core, Episodic, Semantic, and Procedural memory agents.
  • Synthetic question generation creates 5 QA pairs per session across 2,400 sessions, costing $70.The generation uses Claude 4.5 Opus and produces approximately 8.5M input and 1.1M output tokens.
  • GRPO training performs 96,000 session-rollouts and 480,000 QA evaluations, with a per-run total cost of $581.Answer generation costs $461 and LLM judging costs $120.

G.1 Core Memory: Precision Evolution

Across illustrative memory cases, RL learns targeted updates, temporal versioning, and separation of user preferences from conceptual knowledge. These behaviors preserve information that append, rewrite, or overwrite strategies can lose.

  • Core Memory: RL replaces only the relevant Core Memory field, preserving unrelated Family and Hobbies fields.The targeted replacement updates dietary preference, plant-milk preference, and pancake information without rewriting the whole block.
  • Episodic Memory: The RL model constructs a temporal chain by preserving trip planning, adding completion information, and marking status transitions.Base creates a disconnected entry, while SFT overwrites the planning details.
  • Semantic Memory: RL separates impressionist-art preference from Monet’s Water Lilies as distinct entries for more precise retrieval.The analysis connects this separation to answering questions about artistic taste versus factual artwork information.
  • Procedural Memory: RL versioned the new pour-over method instead of overwriting the prior French press procedure.The versioned entry preserves the original method and references it as replaced.
  • Attribution Validation: Retrieval frequency aligns with the actual information source in 95% of sampled cases.Three graduate students evaluated 100 synthetic QA pairs, with pairwise agreement ranging from 92.0% to 96.0%.

H.2 Synthetic QA Quality Evaluation

Synthetic QA quality is assessed through human judgments of reasonableness, correctness, and specificity, with results reported for 100 sampled pairs. The evaluation also controls memory quality when comparing answer models.

  • QA Quality Criteria: Human evaluators assess 100 sampled QA pairs for question reasonableness, answer correctness, and answer specificity.Final labels are determined by majority vote among three graduate students.
  • Error Analysis: Among 8 incorrect answers, 6 involve temporal calculation errors and 2 involve unsupported over-inference across sessions.
  • Answer Model Comparison: Qwen3-30B-A3B performs comparably to GPT-4.1-mini when both use memory constructed by Claude 4.5 Sonnet.The controlled setup isolates answer-model capability under identical retrieval settings.
Loading 2601.05488v4…