GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

1. Het Probleem: Beperkingen van Heuristische Perturbatie

Inference-Time Scaling (ITS) is een methode om de redeneerprestaties van Large Language Models (LLMs) te verbeteren door tijdens het inferentieproces meerdere redeneertrajecten te genereren en de beste te selecteren.

• Discrete vs. Continue Ruimte: In traditionele, op tokens gebaseerde modellen gebeurt dit via probabilistische sampling (bijv. temperature scaling) over een expliciete kansverdeling. In continue latent reasoning-modellen (waar redenering plaatsvindt in verborgen staten zonder tussenliggende teksttokens) ontbreekt deze expliciete verdeling.
• Huidige Aanpak: Bestaande methoden voor ITS in deze continue ruimten vertrouwen op heuristische perturbaties, zoals dropout of het toevoegen van vaste Gaussische ruis.
• De Kernproblematiek: De auteurs tonen aan dat het simpelweg verhogen van de "stochasticiteit" (ruis) niet leidt tot betere resultaten.
  • Sterkere perturbaties veroorzaken vaak grote verschuivingen in de verdeling (distribution shifts) zonder nuttige redeneerpaden te genereren.
  • Deze methoden leiden tot onder-exploratie (te weinig variatie) of over-exploratie (te veel ruis die de beslissing verstoort).
  • Er ontbreekt een expliciete, conditionele steekproefverdeling, waardoor het proces niet gecontroleerd of geoptimaliseerd kan worden.

2. Methodologie: De Gaussian Thought Sampler (GTS)

Om deze beperkingen op te lossen, stellen de auteurs GTS voor. Dit is een lichtgewicht module die latent exploratie herschrijft als het steekproeven uit een geleerde conditionele verdeling.

• Architectuur: GTS wordt toegevoegd aan een bevroren (frozen) latent reasoning backbone. Het bestaat uit twee lichte MLP-heads die per stap in het redeneerproces de mean ($\mu$) en log-standard deviation ($\log \sigma$) voorspellen voor een diagonale Gaussische verdeling.
• Conditionele Sampling: In plaats van vaste ruis toe te voegen, leert GTS een context-afhankelijke verdeling $q_\phi(z_k | c_k)$ over de perturbaties $z_k$. De perturbatie wordt berekend als:
  $$\tilde{h}_k = h^{det}_k + z_k$$
  waarbij $z_k \sim \mathcal{N}(\mu_\phi(c_k), \text{diag}(\sigma^2_\phi(c_k)))$.
• Training met Versterkingslering (RL): Omdat de kwaliteit van een traject afhangt van het eindresultaat (een niet-differentieerbare functie), wordt GTS getraind met GRPO (Group Relative Policy Optimization).
  • Reward Functie: De beloning combineert correctheid (correct vs. incorrect antwoord) met een "confidence shaping" term. Dit stimuleert niet alleen correcte antwoorden, maar ook antwoorden met een hoge zekerheid.
  • Doel: De sampler leert welke perturbaties de kans vergroten dat het model een correct antwoord vindt binnen een beperkt inferentiebudget.

3. Belangrijkste Bijdragen

1. Diagnose van Heuristiek: De auteurs identificeren dat effectieve ITS in continue ruimten niet gaat over het injecteren van meer ruis, maar over het controleren van de steekproefkwaliteit. Ze tonen aan dat heuristische methoden vaak falen omdat ze niet kunnen onderscheiden tussen nuttige en schadelijke exploratie.
2. GTS Framework: Ze introduceren een nieuwe, leerbare benadering die heuristische perturbatie vervangt door een expliciete, optimaliseerbare probabilistische beleidsstrategie.
3. Empirisch Bewijs: Ze demonstreren dat GTS consistent beter presteert dan dropout en vaste Gaussische ruis op meerdere benchmarks en architecturen, zonder de basis-LLM parameters aan te passen.

4. Resultaten en Analyse

De experimenten zijn uitgevoerd op twee latent reasoning-modellen: COCONUT (gebaseerd op GPT-2) en CODI (gebaseerd op LLaMA-3.2-1B), voornamelijk op de GSM8K wiskundetaak.

• Schalingsgedrag (Pass@N): GTS toont een sterkere schaling naarmate het inferentiebudget ($N$, het aantal getrokken trajecten) toeneemt. Waar heuristische methoden (zoals dropout met $p=0.5$) vaak prestatieverlies vertonen bij grotere budgets door over-exploratie, blijft GTS stabiel en verbeterd.
• Steekproefkwaliteit:
  • Sampling Gain (SG): GTS bereikt een hogere "Sampling Gain", wat betekent dat de getrokken trajecten de waarschijnlijkheid van het juiste antwoord significant verhogen ten opzichte van de deterministische baseline.
  • Verdelingsverschuiving: Heuristische methoden veroorzaken vaak grote verdelingsverschuivingen (gemeten via Jensen-Shannon divergentie) zonder de kwaliteit te verbeteren. GTS behoudt een gecontroleerde verschuiving die gericht is op verbetering.
• Generalisatie: GTS presteert ook beter op out-of-distribution datasets (MultiArith, SVAMP, GSM8K-Hard), wat aantoont dat de geleerde perturbatiestrategie robuust is en niet alleen op de trainingsdata is overgefit.
• Adaptiviteit: Analyse van het Signaal-Ruisverhouding (SNR) toont aan dat GTS dynamisch aanpast aan de interne dynamiek van het model. Bij COCONUT volgt het een "zaagtand"-patroon (afwisselend deterministisch en stochastisch), terwijl CODI eerder een "front-loaded" deterministische aanpassing toont. GTS leert deze patronen na te bootsen in plaats van een starre schema te hanteren.

5. Betekenis en Conclusie

Dit paper biedt een fundamentele verschuiving in hoe we naar inferentie-schaling in continue latent spaces kijken.

• Van Heuristiek naar Principieel Beleid: Het bewijst dat "meer ruis" niet gelijkstaat aan "beter redeneren". In plaats daarvan vereist effectieve exploratie een gecontroleerde, leerbare en context-afhankelijke steekproefstrategie.
• Efficiëntie: GTS is een lichtgewicht toevoeging (slechts ~1-2% extra parameters) die geen hertraining van de basis-LLM vereist, maar wel aanzienlijke prestatiewinst biedt tijdens de inferentie.
• Toekomstige Richting: De resultaten suggereren dat toekomstige systemen voor continue redenering expliciete probabilistische modellen moeten integreren om de exploratie-ruimte efficiënter te navigeren, in plaats van te vertrouwen op handgemaakte perturbaties.

Kortom, GTS transformeert latent exploration van een willekeurig proces in een geoptimaliseerde zoekstrategie, wat leidt tot betrouwbaardere en schaalbaardere redeneerprestaties in moderne LLM-architecturen.