TEMPO: Scaling Test-time Training for Large Reasoning Models

Probleemstelling

Grote redeneringsmodellen (Large Reasoning Models - LRMs) tonen indrukwekkende prestaties in complexe taken zoals wiskunde en logica, vaak door uitgebreide redeneerketens tijdens de inferentie te genereren. Echter, hun parameters blijven statisch na het trainen, wat betekent dat ze geen kennis kunnen integreren die tijdens de testfase wordt opgedaan.

Test-time training (TTT) probeert dit op te lossen door modelparameters bij te werken op ongelabelde testgegevens. Bestaande TTT-methoden (zoals TTRL en EMPO) vertrouwen op zelfgegenereerde beloningssignalen (bijv. entropie, meerderheidsstemming, zelf-consistentie) zonder grondwaarheid (ground-truth). Dit leidt tot twee fundamentele beperkingen:

1. Prestatieplateaus: De beloningssignalen worden intrinsiek beperkt door de initiële capaciteiten van het model. Zodra het model zelfverzekerd wordt over een beperkt aantal redeneerpatronen, overschatten de heuristische signalen de kwaliteit, wat leidt tot een zelfversterkende cyclus en stagnatie.
2. Diversiteitsinstorting (Diversity Collapse): Modellen convergeren naar een smal scala aan redeneerpaden om de gemiddelde prestatie te maximaliseren, waardoor de variatie in output afneemt en de algehele redeneerkwaliteit verslechtert.

De kernoorzaak is het ontbreken van grondwaarheid tijdens de testfase, waardoor het beloningssignaal "drift" vertoont naarmate het beleid (policy) evolueert.

Methodologie: TEMPO

De auteurs stellen TEMPO (Test-time Expectation-Maximization Policy Optimization) voor, een TTT-framework dat beloningsgeneratie en beleidsoptimalisatie ontkoppelt via een afwisselend actor-critic ontwerp. Het framework baseert zich op het Expectation-Maximization (EM) algoritme en bestaat uit twee fasen die periodiek worden afgewisseld:

1. E-stap (Expectation): Critic Recalibration

   • Een critic-model wordt periodiek getraind op een gelabelde dataset ($D_L$) om de posterior-verdeling van correcte antwoorden te schatten.
   • De critic leert token-niveau waarde-schattingen ($V_\phi$) om de waarschijnlijkheid van een correct antwoord te voorspellen.
   • Dit zorgt ervoor dat het beloningssignaal "gegrondeerd" blijft in externe supervisie en niet driftt.

2. M-stap (Maximization): Policy Refinement

   • Het actor-model (het beleid $\pi_\theta$) genereert redeneertrajecten op ongelabelde testvragen ($D_u$).
   • Het beleid wordt geoptimaliseerd tegen de beloningen van de gecalibreerde critic, zonder gebruik te maken van grondwaarheid voor deze specifieke vragen.
   • Dit gebeurt via een beleidsgradient (policy gradient) methode waarbij de critic-scores als advantage-signalen worden gebruikt.

Theoretische Inzicht:
De auteurs formaliseren dit als het maximaliseren van een Evidence Lower Bound (ELBO). Bestaande methoden worden geïnterpreteerd als "degeneratieve" varianten van het EM-algoritme die alleen de M-stap uitvoeren (beleidsoptimalisatie) en de cruciale E-stap (posterior-schatting/critic-calibratie) missen. Door de E-stap periodiek te herintroduceren, wordt de ELBO strakker gehouden, wat leidt tot duurzame verbetering.

Belangrijkste Bijdragen

1. TEMPO Framework: Een nieuw TTT-framework dat duurzame prestatieverbeteringen bereikt door afwisselende actor-critic optimalisatie, waardoor diversiteitsinstorting en plateaus worden voorkomen.
2. Unificerende Theoretische Analyse: Het paper karakteriseert bestaande TTT-methoden als onvolledige EM-procedures. Het identificeert het ontbreken van de E-stap (periodieke recalibratie) als de worteloorzaak van schaalbaarheidsproblemen en biedt een principieel remedie.
3. Uitgebreide Experimenten: Validatie over meerdere model-families (Qwen3, OLMO3), schalen (7B, 8B, 14B) en diverse redenerings-domeinen (wiskunde, logica, STEM).

Resultaten

TEMPO toont aanzienlijke verbeteringen ten opzichte van bestaande baselines (RLVR, TTRL, EMPO) op benchmarks zoals AIME 2024/2025 en Beyond AIME:

• Prestatieverbetering:
  • OLMO3-7B: Verbetering op AIME 2024 van 33,0% naar 51,1%.
  • Qwen3-14B: Verbetering op AIME 2024 van 42,3% naar 65,8%.
  • De methoden blijven presteren verbeteren over 350+ trainingsstappen, terwijl baselines na een korte stijging platlopen.
• Behoud van Diversiteit:
  • In tegenstelling tot baselines die lijden aan diversiteitsinstorting (lage pass@K scores ondanks hoge gemiddelde scores), behoudt TEMPO hoge pass@K scores. Dit betekent dat het model diverse, hoogwaardige oplossingen blijft genereren in plaats van te convergeren naar één patroon.
• Generalisatie:
  • TEMPO werkt effectief op niet-wiskundige taken zoals logica-puzzels (ZebraLogic) en complexe STEM-vragen (GPQA-Diamond), wat aantoont dat het ontwerp domein-onafhankelijk is.
• Ablatie Studies:
  • Zonder de E-stap (vastgevroren critic) loopt de prestatie snel vast, wat bevestigt dat periodieke recalibratie op gelabelde data essentieel is voor langdurige zelfverbetering.

Betekenis en Impact

TEMPO biedt een fundamentele doorbraak in het veld van Test-time Training voor grote modellen. Het paper bewijst dat extra rekenkracht tijdens de inferentie (test-time compute) niet alleen nuttig is, maar schaalbaar kan worden vertaald naar prestaties, mits het beloningssignaal correct wordt onderhouden via externe supervisie.

De belangrijkste implicaties zijn:

• Doorbreken van het "Self-Reinforcement" valstrik: Het toont aan dat zelflerende systemen zonder externe kalibratie onvermijdelijk zullen instorten.
• Schaalbaarheid: Het biedt een route om modellen te laten "leren" van nieuwe, moeilijke problemen zonder dat er nieuwe gelabelde data nodig is voor de specifieke testvragen.
• Toekomstige Richting: Het werk legt de basis voor verdere onderzoek naar convergentiegaranties en de toepassing op agentische taken, en suggereert dat de EM-interpretatie een krachtig raamwerk is voor het analyseren van toekomstige TTT-methoden.

Kortom, TEMPO transformeert test-time training van een heuristische oefening die snel verzadigt, naar een gestructureerd, schaalbaar proces dat de grenzen van redeneringsmodellen continu verlegt.