Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models

Probleemstelling

Grote Taalmodellen (LLMs) kunnen complexe taken oplossen door lange ketens van redenering (Chain-of-Thought of CoT) te genereren. Om de prestaties verder te verbeteren, wordt Test-Time Scaling (TTS) gebruikt, waarbij meerdere varianten van tussenstappen worden gegenereerd, geverifieerd en de beste worden geselecteerd (bijv. via Best-of-N of Beam Search).

De huidige staat van de kunst voor het verifiëren van deze redeneerstappen zijn Process Reward Models (PRMs). Deze modellen evalueren de kwaliteit van elke stap in de redenering. PRMs hebben echter drie grote nadelen:

1. Hoge rekencost: Ze zijn vaak zeer grote modellen (1,5 tot 8 miljard parameters), wat leidt tot aanzienlijke overhead in GPU-geheugen en rekentijd tijdens inferentie.
2. Data-efficiëntie: Het trainen van PRMs vereist vaak dure menselijke annotatie of complexe Monte-Carlo-rollouts en consensus-mechanismen.
3. Generalisatie: PRMs zijn vaak gespecialiseerd in specifieke domeinen (zoals wiskunde) en presteren minder goed op onbekende domeinen (Out-of-Domain).

Methodologie: ReProbe

De auteurs stellen ReProbe voor, een lichtgewicht alternatief voor PRMs dat de interne staten van het LLM zelf gebruikt om de geloofwaardigheid van redeneerstappen te schatten.

• Principe: In plaats van een extern model dat de gegenereerde tekst leest (zoals een PRM), "kijkt" ReProbe naar de interne signalen van het genererende model (de "frozen" LLM) terwijl het genereert.
• Input Features: ReProbe gebruikt features zoals attention-weights, hidden states (verborgen staten) en logits van het basis-LLM.
• Architectuur: ReProbe is een zeer klein transformer-model (minder dan 10 miljoen parameters) dat als een plug-in module bovenop het gefroste LLM werkt. Het projecteert de token-level features naar een stap-level score die de waarschijnlijkheid aangeeft dat een stap correct is.
• Training:
  • Het model wordt getraind met een binaire cross-entropy loss.
  • De labels (correct/incorrect) kunnen worden gegenereerd door een groter LLM (zoals DeepSeek-R1) of in een zelf-supervisie setting waarbij het model zijn eigen gegenereerde stappen beoordeelt.
  • De training is extreem efficiënt: het kost slechts ongeveer $200 aan annotatiekosten en 4 GH200 GPU-uren.

Belangrijkste Bijdragen

1. Efficiëntie: ReProbe is tot 810 keer kleiner dan de beste bestaande PRMs (minder dan 10M parameters vs. 1,5-8B parameters), wat leidt tot een aanzienlijke daling in inferentie-tijd en geheugengebruik.
2. Generalisatie: In tegenstelling tot domein-specifieke PRMs, generaliseert ReProbe uitstekend naar nieuwe domeinen (zoals planning en algemene kennisvragen) zonder over te fitten op trainingsdata.
3. Zelf-supervisie: Het paper toont aan dat ReProbe effectief kan worden getraind zonder menselijke labels, door gebruik te maken van het model zelf voor annotatie, wat de kosten drastisch verlaagt.
4. Complementaire Krachten: Het onderzoek toont aan dat ReProbe en PRMs complementaire signalen vangen (interne zekerheid vs. externe kennis), en dat het combineren van beide de prestaties verder verbetert.

Resultaten

De auteurs hebben ReProbe geëvalueerd op diverse domeinen: wiskunde (In-Domain), planning en algemene kennisvragen (Out-of-Domain).

• Stap-voor-stap verificatie: ReProbe presteert vergelijkbaar met of beter dan PRMs die tot 810 keer groter zijn, vooral op Out-of-Domain taken waar PRMs vaak falen.
• Test-Time Scaling: Bij het toepassen van Best-of-N en Beam Search strategieën, zorgt ReProbe voor een hogere nauwkeurigheid van het uiteindelijke antwoord dan de beste PRMs, zelfs op complexe taken.
• Efficiëntie: ReProbe biedt een snelheidswinst van 2,6x tot 25x ten opzichte van state-of-the-art PRMs tijdens de inferentie.
• Native Thinking Mode: ReProbe werkt ook effectief op modellen die in "native thinking mode" werken (waarbij de redenering niet expliciet gestructureerd is in stappen), wat aantoont dat het niet afhankelijk is van specifieke prompt-engineering.

Betekenis en Conclusie

Het paper "ReProbe" markeert een belangrijke verschuiving in hoe we redenering in LLMs verifiëren. Het bewijst dat de interne staten van een model rijke signalen bevatten over de betrouwbaarheid van zijn eigen redenering, zonder dat zware externe verificatiemodellen nodig zijn.

Dit maakt Test-Time Scaling veel toegankelijker en schaalbaarder, omdat de kosten voor verifiëren drastisch dalen. Het opent de weg naar meer introspectieve LLM-systemen die zichzelf kunnen controleren tijdens het genereren, wat essentieel is voor het oplossen van complexe, multi-stap problemen in de praktijk. De combinatie van lichtgewicht probes met bestaande PRMs biedt bovendien een veelbelovende hybride aanpak voor toekomstige redeneringssystemen.