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

1. Problem Statement

Large Language Models (LLMs) struggle to maintain consistency in long-term dialogues due to the inability of standard retrieval mechanisms (like RAG) to capture the temporal evolution of historical states. While memory-augmented frameworks exist, current approaches face three critical limitations:

1. Open-Loop Operation: Prompting-based systems (e.g., MemGPT, Mem0) operate without feedback on whether the constructed memories actually benefit downstream tasks.
2. Model Dependency: They rely heavily on expensive, closed-source models for memory construction, making them cost-prohibitive at scale.
3. Training Bottlenecks: Existing training-based methods suffer from:
   • Sparse Trajectory-Level Rewards: A single reward at the end of a long dialogue provides no signal for which specific memory operations in intermediate sessions contributed to the final success.
   • Multi-Dimensional Attribution Ambiguity: In multi-component memory systems (e.g., Episodic, Semantic), a global reward fails to distinguish the specific contribution of each component, leading to noisy gradient updates.

2. Methodology: MemBuilder Framework

MemBuilder introduces a Reinforcement Learning (RL) framework designed to train lightweight models (specifically a 4B parameter model) to orchestrate multi-dimensional memory construction.

A. Multi-Dimensional Memory Architecture

The system decomposes conversation history into four specialized memory types, each managed by a role-specific prompt:

1. Core Memory ($M_{core}$): Persistent user profile (identity, preferences, relationships). Stored as a fixed-size block with automatic compression.
2. Episodic Memory ($M_{epi}$): Time-stamped event records ("What happened and when").
3. Semantic Memory ($M_{sem}$): Factual knowledge about entities (people, places, concepts).
4. Procedural Memory ($M_{proc}$): Step-by-step processes and routines.

Action Spaces: The model learns specific operations for each type (e.g., APPEND, REPLACE, REWRITE for Core; ADD, UPDATE, MERGE for Episodic). Notably, the UPDATE operation creates a new entry referencing the old one rather than overwriting, preserving temporal history.

B. Training Pipeline

The training process consists of two stages:

1. Supervised Fine-Tuning (SFT): Uses expert trajectories (generated by Claude 4.5 Sonnet) to teach the 4B model valid JSON formatting and operation structures, solving the "cold-start" problem where RL alone fails due to invalid outputs.
2. Attributed Dense Rewards Policy Optimization (ADRPO): The core RL algorithm addressing the two bottlenecks mentioned above.

C. Key Technical Innovations in ADRPO

1. Dense Session-Level Rewards via Synthetic QA:

   • Instead of waiting for a final reward, the system generates synthetic question-answer (QA) pairs for each session based on the current memory state.
   • An LLM judge evaluates the accuracy of answers retrieved from the candidate memory bank.
   • This provides a dense learning signal ($r_{task}$) after every session, allowing the model to learn which memory operations were beneficial immediately.
   • Regularization: The final reward includes a validity gate (zero reward for malformed JSON) and a length penalty to prevent verbose, uninformative memory storage.

2. Contribution-Aware Gradient Weighting:

   • Since all memory components share a global session reward, the system cannot inherently know which component drove the success.
   • The method dynamically weights gradient updates based on retrieval frequency. If a specific memory type (e.g., Episodic) is retrieved frequently to answer questions, its gradient updates are amplified ($\alpha > 1$).
   • This resolves the credit assignment problem in cooperative multi-agent settings, ensuring that components contributing most to downstream QA performance receive stronger reinforcement.

3. Key Contributions

• ADRPO Algorithm: A novel RL framework combining dense session-level rewards and contribution-aware gradient weighting to optimize long-term memory construction.
• Multi-Dimensional Memory Design: A structured approach (Core, Episodic, Semantic, Procedural) that shifts computational burden from processing entangled tokens to retrieving "pre-digested" fragments.
• Performance with Lightweight Models: Demonstrates that a trained 4B open-source model can outperform state-of-the-art closed-source prompting frameworks.
• Synthetic QA for Dense Rewards: A mechanism to generate intermediate learning signals in long-horizon tasks where traditional sparse rewards fail.

4. Experimental Results

The framework was evaluated on three long-term dialogue benchmarks: LoCoMo, LongMemEval, and PerLTQA.

• Performance:
  • MemBuilder (4B model) achieved 84.23% on LoCoMo, surpassing the best prompting-based baseline (MIRIX at 77.48%) and even outperforming the closed-source Claude 4.5 Sonnet (82.61%) used as a memory constructor.
  • It achieved 85.75% on LongMemEval and 93.14% on PerLTQA.
• Generalization: Despite being trained exclusively on LongMemEval, the model generalized effectively to out-of-distribution (OOD) benchmarks with different dialogue structures (human-human vs. user-assistant).
• Ablation Studies:
  • Gradient Weighting: Optimal performance was achieved with a weighting coefficient $\alpha=4$, confirming that amplifying updates for dominant memory types improves results.
  • Reward Density: Sparse rewards (rewarding only 1/8 of sessions) degraded performance below the SFT baseline, proving the necessity of dense session-level feedback.
  • Model Efficiency: The trained 4B model significantly outperformed Memory-R1 (which used an 8B model) and closed-source baselines.

5. Significance

MemBuilder represents a paradigm shift in long-term memory management for LLMs:

1. Cost Efficiency: It proves that lightweight, open-source models can replace expensive closed-source APIs for memory construction if trained with the right RL objectives.
2. Scalability: By solving the sparse reward and attribution problems, it enables the training of robust memory systems for extended dialogues without relying on "open-loop" prompting.
3. Interpretability & Control: The structured, multi-dimensional memory approach offers better interpretability and fine-grained control over what information is stored and how it evolves, compared to latent memory or raw RAG approaches.

In summary, MemBuilder establishes that with Attributed Dense Rewards, small models can learn to construct high-quality, structured long-term memories that outperform larger, untrained, or prompting-based systems.