MemBuilder: Reinforcing LLMs for Long-Term Memory Construction via Attributed Dense Rewards
This addresses the challenge of long-term memory in dialogues for AI systems, though it appears incremental as it builds on existing memory-augmented frameworks with novel training methods.
The paper tackled the problem of maintaining consistency in long-term dialogues for LLMs by introducing MemBuilder, a reinforcement learning framework that uses attributed dense rewards for memory construction, resulting in a 4B-parameter model outperforming state-of-the-art closed-source baselines.
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.