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

TL;DR

MemBuilder is a reinforcement learning framework that improves long-term memory in large language models by using dense rewards and contribution-aware training, leading to better dialogue consistency.

Contribution

It introduces a novel RL approach with dense rewards and attribution mechanisms to enhance memory construction in LLMs, outperforming existing baselines.

Findings

4B-parameter model surpasses state-of-the-art baselines

Enables strong generalization across long-term dialogue benchmarks

Addresses sparse rewards and multi-dimensional memory attribution

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.