Technische Samenvatting: Redeneren in Geheugen (RiM)

Probleemstelling

Huidige methoden voor het verbeteren van de redeneercapaciteiten van Large Language Models (LLM's) vertrouwen sterk op het schalen van de rekentijd tijdens het testen door het genereren van tussenliggende tokens vóór het definitieve antwoord, zoals bij Chain-of-Thought (CoT) redeneren. Deze benadering koppelt echter interne berekening direct aan autoregressieve tekstgeneratie. Bijgevolg wordt het model gedwongen om "hardop na te denken", waarbij een aanzienlijk deel van zijn rekenbudget wordt besteed aan het genereren van grammaticaal vloeiende natuurlijke taal in plaats van aan puur interne redenering.

Hoewel recente benaderingen voor "latente redenering" proberen de beperkingen van natuurlijke taal te omzeilen door het gebruik van continue representaties, behouden ze grotendeels het autoregressieve generatieparadigma van stap-voor-stap. Bij deze methoden worden tussenliggende berekeningen nog steeds sequentieel geëxternaliseerd voordat de volgende stap zich hierop kan baseren. Dit artikel betoogt dat menselijke cognitie een superieur ontwerpprincipe biedt: het werkgeheugen. Mensen kunnen taakrelevante informatie intern vasthouden en manipuleren zonder elke tussenliggende gedachte te externaliseren. De auteurs stellen dat LLM's kunnen worden getraind om een vergelijkbare interne werkruimte te gebruiken, waardoor redeneren wordt ontkoppeld van de sequentiële bottleneck van token-generatie.

Methodologie: Redeneren in Geheugen (RiM)

De auteurs introduceren Redeneren in Geheugen (RiM), een methode voor latente redenering die de autoregressieve generatie van redeneerstappen vervangt door vaste geheugenblokken.

Kernmechanisme

In plaats van een reeks redeneertokens te genereren, krijgt het model $K$ vaste geheugenblokken, $m_{1:K}$. Elk blok bestaat uit een vooraf gedefinieerde reeks speciale tokens:
$$m_k = [<b>, <m>, \dots, <m>, </b>]$$
waarbij $<b>$ en $</b>$ de blokgrenzen afbakenen en $<m>$-tokens de latente werkruimte vormen.

• Vaste Invoer: Deze blokken worden als vaste tokens aan de invoervraag toegevoegd.
• Enkele Forward Pass: Omdat de blokken vaste invoer zijn in plaats van gegenereerde output, wordt de volledige verrijkte reeks (vraag + geheugenblokken) in één enkele forward pass verwerkt. Dit elimineert de sequentiële decoderingsbottleneck die geassocieerd is met autoregressieve generatie.
• Gestudeerde Representaties: Hoewel tokenidentiteiten en posities vastliggen, worden hun contextuele representaties zo geleerd dat ze taakafhankelijk zijn en tussenliggende redeneerstappen coderen.

Tweestaps Curriculum

Om het model te leren deze vaste blokken als werkgeheugen te gebruiken, hanteert RiM een tweestaps trainingscurriculum:

Fase 1: Supervisie van Redeneerstappen

• Doel: De geheugenblokken grondvesten door ze een rekenrol toe te kennen.
• Mechanisme: Voor een redeneersporen met $T$ stappen wordt het model getraind met $T$ geheugenblokken. Na elk blok $m_t$ moet het model de volgende expliciete redeneerstap $r_{t+1}$ voorspellen.
• Beperking: Een aangepast attentiemasker zorgt ervoor dat de uitlezing voor stap $t+1$ zich kan richten op de vraag en eerdere geheugenblokken, maar niet op eerdere expliciete redeneerstappen. Dit dwingt het model om de nodige tussenberekening te coderen binnen de latente werkruimte van de geheugenblokken.
• Supervisie: Gebruikt een gewogen negatieve log-waarschijnlijkheidsdoelstelling met een afkoelingschema om de supervisiesterkte geleidelijk te verminderen.

Fase 2: Verfijning van het Definitieve Antwoord

• Doel: De grondveste geheugenblokken omzetten in een latente werkruimte voor de voorspelling van het definitieve antwoord.
• Mechanisme: Het expliciete redeneersporen wordt verworpen. Het model gebruikt een vast aantal $K$ geheugenblokken. Na elk blok wordt het model getraind om het definitieve antwoord direct te voorspellen.
• Doelstelling: Dit verschuift de supervisie van het voorspellen van tussenstappen naar het geleidelijk verfijnen van het definitieve antwoord naarmate meer geheugenblokken (latente berekening) beschikbaar komen.
• Inferentie: Tijdens het testen verwerkt het model de vraag en geheugenblokken in één pass, en de uitlezing van het laatste blok levert het antwoord.

Attentie Maskering

Een aangepast causaal attentiemasker (Figuur 2 in het artikel) is cruciaal voor de methode. Het voorkomt informatielekken door ervoor te zorgen dat:

1. Geheugenblokken zich richten op de vraag en eerdere geheugenblokken.
2. Doel-uitlezingen (redeneerstappen in Fase 1, definitief antwoord in Fase 2) zich richten op de vraag en beschikbare geheugenblokken, maar niet op andere doel-uitlezingen.
   Dit staat toe dat alle doelen in één enkele forward pass worden voorspeld zonder dat het model "valt" door toekomstige grondwaarheid te lezen.

Belangrijkste Bijdragen

1. Redeneren in Geheugen (RiM): Een nieuwe methode voor latente redenering die de werkgeheugencapaciteit van LLM's ontsluit door een toegewijde, vaste latente werkruimte voor tussenberekening te bieden, waardoor deze wordt ontkoppeld van autoregressieve generatie.
2. Tweestaps Curriculum: Een gestructureerde trainingsaanpak die eerst geheugenblokken grondvest via supervisie van expliciete stappen en deze vervolgens verfijnt voor de voorspelling van het definitieve antwoord, waardoor effectief gebruik van de latente ruimte mogelijk wordt zonder complexe auxiliaire doelstellingen.
3. Efficiëntie en Prestaties: Empirisch bewijs dat aantoont dat RiM bestaande methoden voor expliciete en latente redenering evenaart of overtreft, terwijl het de inferentiesnelheid van modellen voor directe antwoorden behoudt (geen sequentiële generatie-overhead).

Experimentele Resultaten

De auteurs hebben RiM geëvalueerd op de wiskundige redeneerbenchmarks GSM8K (in-distributie) en GSM-Hard (out-of-distributie) met behulp van GPT-2 en Llama-3.2-modellen (1B en 3B parameters).

• Prestaties: RiM presteerde consequent beter dan de sterkste latente redeneringsbaseline (Coconut) en Supervised Fine-Tuning (SFT) voor directe antwoorden.
  • Op GSM8K (Llama-3.2-1B) bereikte RiM 42,1% greedy-accuraatheid, vergeleken met 36,9% voor Coconut en 23,9% voor SFT zonder CoT.
  • Op GSM-Hard bereikte RiM 10,5%, wat Coconut (8,5%) en SFT zonder CoT (5,3%) overtreft.
  • RiM toonde ook superieure Pass@8-scores, wat aangeeft dat correcte oplossingen met hoge waarschijnlijkheid over de latente trajecten zijn verdeeld.
• Latentie: RiM bereikte inferentiesnelheden die vergelijkbaar waren met SFT zonder CoT (Tijd tot Eerste Token ~16ms voor Llama-3.2-1B), terwijl Coconut ongeveer 7x langzamer was en SFT met CoT ongeveer 27x langzamer vanwege autoregressieve generatie.
• Representatie-analyse: PCA-projecties van representaties van geheugenblokken toonden aan dat het model leert om onderscheidende, steekproefafhankelijke latente werkruimten te creëren. De representaties evolueren systematisch tijdens de training, wat bevestigt dat de blokken niet slechts vaste placeholders zijn, maar actieve computationele toestanden.
• Robuustheid: De prestaties bleven stabiel bij variërende rekenbudgetten tijdens de inferentie (aantal blokken), wat suggereert dat de methode robuust is voor implementatiebeperkingen.

Betekenis en Claims

Het artikel beweert dat LLM's kunnen worden getraind om werkgeheugen te gebruiken als een effectief mechanisme voor latente redenering.

De primaire betekenis ligt in de ontkoppeling van berekening van communicatie. Door het genereren van tussenliggende gedachten te vervangen door vaste geheugenblokken, stelt RiM modellen in staat om dichte, paralleliseerbare interne berekeningen uit te voeren zonder de overhead van het genereren van natuurlijke taal of sequentiële continue representaties. Dit resulteert in een methode die:

1. Rekenkundig Efficiënt is: Het vermijdt de sequentiële bottleneck van autoregressieve generatie, waardoor inferentie in één pass mogelijk is.
2. Effectief is: Het bereikt state-of-the-art prestaties onder methoden voor latente redenering op wiskundige benchmarks.
3. Schaalbaar is: Het werkt over verschillende modelfamilies en -groottes zonder complexe recurrente architecturen of trainingsdoelstellingen met meerdere paden te vereisen.

De auteurs concluderen dat deze aanpak een haalbare weg aantoont naar efficiëntere redeneersystemen die de cognitieve efficiëntie van mensen nabootsen door het redeneerproces te internaliseren in plaats van het te externaliseren als tekst.