MemReward: Graph-Based Experience Memory for LLM Reward Prediction with Limited Labels

MemReward is een grafgebaseerd ervaringsgeheugenframework dat door het modelleren van rollouts als een heterogeen graf en het gebruik van een GNN voor reward-propagatie, de prestaties van reinforcement learning bij grote taalmodellen aanzienlijk verbetert met slechts een beperkt aantal reward-labels.

De kern

Stel je voor dat je een zeer slimme, maar nog onervaren robot wilt leren om complexe problemen op te lossen, zoals wiskundige raadsels, code te schrijven of vragen te beantwoorden. Om dit te doen, moet de robot oefenen door duizenden pogingen te doen. Maar hier is het probleem: om te weten of een poging goed of fout was, heb je een menselijke expert nodig om elke oplossing te controleren.

In de echte wereld is het echter veel te duur en te tijdrovend om duizenden experts in te huren om elke antwoord van een robot te beoordelen. Het is alsof je een school wilt laten slagen, maar je hebt maar geld voor één leraar om 1000 leerlingen te controleren. De meeste leerlingen krijgen dan geen feedback en leren niet.

MemReward is de oplossing die de auteurs van dit paper hebben bedacht. Het is een slimme manier om die schaarste aan "menselijke feedback" op te lossen. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Gedachtenkrant" (Experience Memory)

Stel je voor dat de robot eerst een grote stapel met zijn eigen gedachten en antwoorden verzamelt. Dit noemen ze "rollouts".

• De vraag: Het probleem dat de robot moet oplossen.
• Het denkproces: Hoe de robot stap voor stap redeneert.
• Het antwoord: Het eindresultaat.

Normaal gesproken zou je nu alleen de antwoorden van de 20% van de vragen controleren waarvoor je een menselijke expert hebt. De andere 80% zou je gewoon weggooien. MemReward doet dit niet. Het bewaart alles in een speciale "gedachtenkrant".

2. De "Slimme Netwerkkaart" (Heterogeneous Graph)

Dit is het hart van de uitvinding. In plaats van dat de robot antwoorden als losse eilanden ziet, bouwt MemReward een enorm, complex netwerk (een grafiek) van al die gedachten.

• De knooppunten: De vragen, de denkprocessen en de antwoorden zijn allemaal punten op dit netwerk.
• De lijntjes: Het systeem verbindt punten die op elkaar lijken.
  • Als twee vragen qua inhoud erg op elkaar lijken, krijgt ze een lijntje.
  • Als een vraag leidt tot een bepaald denkproces, krijgt ze een lijntje.
  • Als een denkproces leidt tot een antwoord, krijgen ze een lijntje.

Het is alsof je een gigantische stamboom maakt van alle problemen die de robot heeft gezien. Als je weet dat vraag A goed was, en vraag B lijkt heel erg op vraag A, dan is de kans groot dat vraag B ook goed is.

3. De "Boodschapper" (GNN)

Nu komt de magie. Het systeem gebruikt een speciaal type kunstmatige intelligentie (een GNN, of Graph Neural Network) dat fungeert als een boodschapper.

• De boodschapper kijkt naar de 20% van de vragen die wel door een mens zijn gecontroleerd (de "goede" en "slechte" voorbeelden).
• Vervolgens "stroomt" deze kennis door het netwerk. De boodschapper zegt: "Hey, deze vraag hier is heel erg vergelijkbaar met die vraag daar, die we al weten dat goed was. Dus deze nieuwe vraag is waarschijnlijk ook goed!"
• Zo krijgt de robot voor de 80% van de vragen die geen menselijke controle kregen, toch een soort van "virtuele feedback" op basis van wat soortgelijke vragen eerder hebben opgeleverd.

4. Het Resultaat: Leren met weinig geld

Het mooie aan MemReward is dat het de robot bijna net zo goed laat leren als wanneer je 100% van de antwoorden door mensen had laten controleren, terwijl je maar 20% betaalde.

• Voorbeeld uit de wiskunde: Als de robot een moeilijke vergelijking oplost die lijkt op een andere die hij al goed heeft opgelost, krijgt hij automatisch een "groen licht" (beloning), zelfs zonder dat een mens het heeft gezien.
• Voorbeeld uit code: Als de code die de robot schrijft lijkt op code die eerder werkte, krijgt hij een positieve feedback.

Waarom is dit zo belangrijk?

Stel je voor dat je een nieuwe taal wilt leren.

• De oude manier: Je moet elke zin die je schrijft laten controleren door een leraar. Als de leraar er niet is, stop je met oefenen.
• De MemReward manier: Je schrijft alles op. Je laat de leraar alleen de eerste paar zinnen controleren. Daarna kijk je naar je eigen notities en zeg je: "Deze zin lijkt veel op die vorige die de leraar goedkeurde, dus deze is waarschijnlijk ook goed." Je kunt blijven oefenen en verbeteren, zelfs als de leraar weg is.

Kort samengevat:
MemReward is een slimme "herinneringsmachine" die de ervaringen van een AI-netwerk met elkaar verbindt. Door te kijken naar wat er eerder goed ging bij vergelijkbare situaties, kan de AI zichzelf belonen en straffen zonder dat er voor elke stap een mens nodig is. Hierdoor kunnen we slimme robots veel goedkoper en sneller leren, zelfs met een heel klein budget voor menselijke controle.

Technische samenvatting

Probleemstelling

Het trainen van Large Language Models (LLM's) voor complex redeneren via versterkingsleren (Reinforcement Learning, RL) vereist vaak reward labels die aangeven of gegenereerde antwoorden (rollouts) correct zijn. Het verkrijgen van deze labels op grote schaal is echter een grote bottleneck:

• Kosten en tijd: Het vereist vaak dure menselijke annotatie of tijdrovende verificatieprocedures (bijvoorbeeld het beoordelen van wiskundige bewijzen door experts).
• Gebrek aan ground truth: Bij open-ended vragenontwikkeling ontbreekt vaak een definitief correct antwoord.
• Beperkingen: Wanneer reward labels schaars zijn, wordt de effectiviteit van RL-finetuning beperkt. Bestaande methoden die slechts een klein deel van de data labelen (bijv. 20%), presteren aanzienlijk slechter dan volledig gesuperviseerde methoden.

De kernvraag is: Hoe kunnen we gestructureerd geheugen van gelabelde redeneerervaringen gebruiken om rewards te voorspellen voor ongelabelde rollouts, zodat effectieve RL-finetuning mogelijk is onder label-schaarste?

Methodologie: MemReward

MemReward is een raamwerk dat een graf-gebaseerd ervaringsgeheugen introduceert om rewards te propageren van gelabelde naar ongelabelde data. De methode bestaat uit de volgende fasen:

1. Opbouw van het Heterogene Graafgeheugen

Voordat de RL-optimalisatie begint, genereert een initiële LLM-policy voor elke query een "rollout" bestaande uit een denkproces (thinking process) en een eindantwoord (answer). Deze worden opgeslagen als ervaringen in een heterogene graaf met drie soorten knopen en drie soorten randen:

• Knopen:
  • Query-knopen: De oorspronkelijke vraag.
  • Thinking-knopen: Het redeneerproces van het model.
  • Answer-knopen: Het gegenereerde antwoord.
• Randen:
  • Query-Query: Verbindt semantisch vergelijkbare vragen (gebaseerd op cosine-相似iteit van embeddings).
  • Query-Thinking: Verbindt een vraag met zijn bijbehorende denkproces.
  • Thinking-Answer: Koppelt een denkproces één-op-één aan het corresponderende antwoord.

2. Training van de GNN (Warm-up Fase)

Een Heterogene Graph Neural Network (GNN) wordt getraind op de gelabelde subset van de data.

• De GNN gebruikt type-specifieke gewichten en attentie-mechanismen om informatie te aggregeren over de verschillende knoop- en randtypes.
• Het doel is om een reward-score te voorspellen voor elke rollout door de relationele boodschappen tussen de knopen te verwerken.
• De loss-functie is de binaire cross-entropy, waarbij de GNN leert om correcte antwoorden (reward = 1) te onderscheiden van incorrecte (reward = 0).

3. Online Policy Optimalisatie met GRPO

Tijdens het online trainen (met GRPO - Group Relative Policy Optimization) wordt een hybride reward-strategie toegepast:

• Gelabelde queries: Ontvangen hun echte ground-truth reward.
• Ongelabelde queries: Worden tijdens de inferentie verbonden met de warm-up graaf via de top-k meest vergelijkbare gelabelde queries. De getrainde GNN voorspelt vervolgens de reward voor deze ongelabelde rollouts.
• Deze voorspelde rewards worden gebruikt om de advantage-waarden te berekenen en de policy te updaten, waardoor het model kan leren van de volledige dataset, niet alleen van het kleine gelabelde deel.

Belangrijkste Bijdragen

1. Graf-gebaseerde Reward Propagatie: MemReward is de eerste methode die een heterogene graaf (query, denken, antwoord) gebruikt om reward signalen te propageren in RL voor LLM's, in plaats van elke ervaring onafhankelijk te behandelen.
2. Effectiviteit bij Label-Schaarste: Het raamwerk maakt het mogelijk om bijna even goed te presteren als volledig gesuperviseerde methoden ("Oracle") met slechts 20% ground-truth labels.
3. Cross-Domain Generalisatie: De GNN is getraind op een mix van domeinen (wiskunde, vragenbeantwoording, codering) en kan deze patronen overdragen naar onbekende domeinen zonder domeinspecifieke labels.
4. Architecturale Inzichten: Het paper toont aan dat het onderscheiden van knooptypes (heterogeen) en het behouden van het denkproces (thinking nodes) cruciaal zijn voor succes, vooral bij complexe redeneertaken.

Resultaten

De methode is geëvalueerd op Qwen2.5-3B en Qwen2.5-1.5B over 13 benchmarks in drie domeinen: Wiskunde, Vragenbeantwoording (QA) en Code Generatie.

• Prestatie met 20% Labels:
  • MemReward bereikt 97,3% van de Oracle-prestatie op het 3B-model en 96,6% op het 1,5B-model.
  • Het overtreft aanzienlijk de baseline "R1-p" (die alleen 20% labels gebruikt en de rest verworpt), met verbeteringen van +1,35 punten (3B) en +5,38 punten (1,5B).
• Buitendomein Generalisatie (Out-of-Domain):
  • Opvallend is dat MemReward op buitendomein-taken (zoals NuminaMath, SIQA, PIQA) zelfs beter presteert dan de volledig gesuperviseerde Oracle.
  • Voor het 3B-model: 66,96% (MemReward) vs. 66,07% (Oracle).
  • Dit suggereert dat de GNN-predicties generalisatie verbeteren door semantisch vergelijkbare patronen uit het trainingsgeheugen te halen.
• Schaalbaarheid:
  • De prestaties schalen soepel met het label-budget. Bij 70% labels bereikt het systeem 99,4% van de Oracle-prestatie.
• Ablatie Studies:
  • Het verwijderen van de grafstructuur (MLP-baseline) leidt tot de grootste daling, vooral bij QA.
  • Het verwijderen van "thinking nodes" heeft een negatief effect op wiskunde en codering, wat aantoont dat het redeneerproces essentieel is voor het voorspellen van de correctheid.

Significantie

MemReward biedt een praktische oplossing voor het grootste obstakel in RL voor LLM's: de afhankelijkheid van enorme hoeveelheden menselijke labels.

• Efficiëntie: Het stelt onderzoekers in staat om RL-finetuning uit te voeren met tot 80% minder labels zonder significante prestatieverlies.
• Democratisering: Door de kosten van annotatie te verlagen, wordt RL-training toegankelijker voor onderzoekers met beperkte budgetten.
• Generalisatie: De bevinding dat graf-gebaseerde reward-predictie zelfs beter kan generaliseren dan volledige supervisie op onbekende taken, opent nieuwe richtingen voor het ontwerpen van robuuste RL-systemen die vertrouwen op gestructureerd geheugen in plaats van louter brute-force annotatie.

Kortom, MemReward bewijst dat door de structurele afhankelijkheden tussen redeneerervaringen te modelleren, de kwaliteit van de reward-signalen drastisch kan worden verbeterd, zelfs wanneer de grondwaarheid schaars is.