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

Probleemstelling

Het handhaven van consistentie in lange-termijn dialogen blijft een fundamentele uitdaging voor Large Language Models (LLMs). Bestaande oplossingen kampen met twee hoofdproblemen:

1. Beperkingen van RAG: Retrieval-Augmented Generation (RAG) behandelt informatie als statische, onafhankelijke stukken ("chunks"). Dit faalt in het vastleggen van de temporele evolutie van historische toestanden of het begrijpen welke feiten verouderd zijn.
2. Inefficiënte training van bestaande frameworks: Bestaande memory-augmented frameworks (zoals MemGPT, Mem0) vertrouwen vaak op prompting van dure, gesloten bronmodellen zonder feedbacklus. Trainingsgebaseerde benaderingen (zoals Memory-R1 of Mem-α) lijden onder:
   • Sparse Rewards: In lange dialogen wordt vaak pas één beloning gegeven aan het einde van een traject. Dit maakt het voor het model onmogelijk om te bepalen welke specifieke sessie of actie bijdroeg aan het succes, wat leidt tot ruis in de gradient-updates.
   • Gebrek aan attributie: Bestaande methoden geven een globale beloning voor alle geheugentypes (bijv. episodisch, semantisch), zonder rekening te houden met het feit dat sommige componenten meer bijdragen aan de uiteindelijke taakprestatie dan anderen.

Methodologie: MemBuilder

MemBuilder is een Reinforcement Learning (RL) framework dat is ontworpen om modellen te trainen voor het opbouwen van multidimensionaal geheugen met behulp van toegeschreven dichte beloningen (attributed dense rewards).

1. Multidimensionale Geheugenarchitectuur

Het systeem deconstrueert conversaties in vier gespecialiseerde geheugentypen, elk beheerd door een specifieke prompt voor hetzelfde LLM:

• Core Memory: Persistente gebruikersprofielinformatie (identiteit, voorkeuren, relaties). Dit wordt bewaard als een vast tekstblok.
• Episodic Memory: Tijdstempelgedreven gebeurtenisregistraties ("wat" en "wanneer").
• Semantic Memory: Feitelijke kennis over entiteiten in het leven van de gebruiker (mensen, plaatsen, concepten).
• Procedural Memory: Stap-voor-stap processen, routines en workflows.

Het model voert acties uit zoals ADD, UPDATE, MERGE, REPLACE en SKIP. Belangrijk is dat UPDATE en MERGE de temporele geschiedenis behouden door verwijzingen naar eerdere entries te maken in plaats van deze te overschrijven.

2. Attributed Dense Rewards Policy Optimization (ADRPO)

De kerninnovatie is de ADRPO-algoritme, die twee technische uitdagingen aanpakt:

• Dichte Sessie-niveau Beloningen (Dense Session-Level Rewards):
  In plaats van te wachten op het einde van een lange conversatie voor een beloning, genereert het systeem synthetische vraag-antwoordparen (QA) per sessie. Een expert-model genereert vragen gebaseerd op de huidige sessie en het bestaande geheugen. Tijdens het RL-trainingstraject wordt de kwaliteit van het gegenereerde geheugen direct getest door een model deze vragen te laten beantwoorden. De beloning is gebaseerd op de nauwkeurigheid van deze antwoorden per sessie, wat zorgt voor een continu leersignaal.

• Bijdrage-bewuste Gradient Weging (Contribution-Aware Gradient Weighting):
  Omdat alle vier geheugentypen een gezamenlijke beloning delen, is het onduidelijk welke component de prestaties het meest beïnvloedde. MemBuilder lost dit op door de gradient-updates te schalen op basis van de retrieval-frequentie.

  • Als een bepaald geheugentype (bijv. Episodisch) vaker wordt opgehaald om een vraag te beantwoorden, krijgt dit type een hogere gewichtsfactor ($\alpha > 1$) in de gradient-berekening.
  • Dit zorgt ervoor dat de modelupdates sterker worden versterkt voor de componenten die daadwerkelijk bijdragen aan het succes van de downstream-taak (QA).

3. Trainingspipeline

• Supervised Fine-Tuning (SFT): Eerst wordt het model (Qwen3-4B) getraind op expert-trajecten gegenereerd door Claude 4.5 Sonnet. Dit lost het "cold-start"-probleem op door geldige outputformaten (JSON) te garanderen.
• RL met ADRPO: Vervolgens wordt het model verfijnd met RL, waarbij de dichte sessie-beloningen en de bijdrage-bewuste weging worden toegepast om de downstream QA-accuraatheid te maximaliseren.

Belangrijkste Resultaten

Experimenten zijn uitgevoerd op drie benchmarks: LongMemEval, LoCoMo en PerLTQA.

• Prestaties: Het getrainde Qwen3-4B model bereikte een score van 84,23% op de LoCoMo-benchmark. Dit is een significante verbetering ten opzichte van state-of-the-art prompting-frameworks (zoals MIRIX: 77,48%) en zelfs beter dan het gebruik van een gesloten bronmodel (Claude 4.5 Sonnet) voor het geheugenbouwen (82,61%).
• Generalisatie: Het model is uitsluitend getraind op LongMemEval, maar toont sterke generalisatie naar out-of-distribution benchmarks (LoCoMo en PerLTQA) met verschillende dialoogstructuren.
• Ablatie Studies:
  • Zonder SFT presteert RL slecht (79,31% vs 81,74% voor SFT alleen), wat aantoont dat SFT essentieel is voor het oplossen van het cold-start-probleem.
  • Dichte beloningen zijn cruciaal; bij het verminderen van de beloningsdichtheid daalt de prestatie onder de SFT-baseline.
  • De optimale wegingsfactor voor bijdrage-bewuste gradienten bleek $\alpha = 4$ te zijn.

Betekenis en Impact

MemBuilder demonstreert dat:

1. Lichtgewicht open-source modellen (4B parameters) gespecialiseerde, gesloten bronmodellen kunnen vervangen voor complexe geheugenconstructie-taken, mits ze correct worden getraind.
2. Reinforcement Learning effectief kan worden ingezet voor geheugenmanagement door het probleem van sparse rewards op te lossen via synthetische sessie-evaluaties en door credit assignment (toewijzing van verdiensten) te verfijnen via attributie.
3. Een gestructureerde, multidimensionale aanpak voor geheugen (Episodisch, Semantisch, Proceduraal, Core) superieur is aan statische RAG-benaderingen voor lange-termijn dialogen, vooral bij vragen die temporeel redeneren en multi-hop inferentie vereisen.

De studie biedt een schaalbare en kosteneffectieve route voor het bouwen van agenten met robuust lange-termijn geheugen, zonder afhankelijk te zijn van dure API-aanroepen voor de geheugenconstructie zelf.