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

TL;DR

MemBuilder is introduced, a reinforcement learning framework that trains models to orchestrate multi-dimensional memory construction with attributed dense rewards, and introduces contribution-aware gradient weighting that scales policy updates based on each component's downstream impact.

Abstract

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.

Figures (14)

• Figure 1: Sparse trajectory-level rewards (top) vs. our attributed dense session-level rewards (bottom). Dense rewards provide learning signals at each session rather than only at trajectory end.
• Figure 2: Multi-Dimensional Memory Architecture. Four memory types (Core, Episodic, Semantic, Procedural) are constructed during the Build Phase and selectively retrieved during the Answer Phase.
• Figure 3: ADRPO training pipeline. Each session's memory rollouts are evaluated via synthetic QA, with gradients weighted by each memory component's downstream contribution.
• Figure 4: Training curves with different gradient weighting coefficients $\alpha \in \{1, 2, 4, 8, 16\}$ on LoCoMo.
• Figure 5: Effect of reward density on LoCoMo accuracy. The x-axis indicates the fraction of sessions receiving task rewards during training.