Technische Samenvatting: Temperatuur-geschaalde On-Policy Self-Distillatie (TS-OPSD)

1. Probleemstelling

Reinforcement Learning from Verifiable Rewards (RLVR) heeft de redeneervaardigheden van Large Language Models (LLMs) aanzienlijk verbeterd, met name in wiskundige en logische domeinen. Deze vooruitgang stuit echter op een kritieke flessenhals die bekend staat als entropie-instorting (entropy collapse). Naarmate de beleidsoptimalisatie vordert, wordt de outputdistributie van het model steeds geconcentreerder (scherper), wat leidt tot:

• Verminderde Rollout-diversiteit: Het model genereert repetitieve trajecten, waardoor de exploratie van de oplossingsruimte wordt beperkt.
• Afnemende Leersignalen: Effectieve gradiëntinformatie concentreert zich op een klein fractie van hoog-entropische tokens bij cruciale vertakkingen in het redeneerproces. Zodra het beleid instort, verdwijnen deze signalen, wat verdere verbetering blokkeert.

Bestaande remedies proberen dit te mitigeren door ofwel:

1. Het RL-doel te wijzigen: Door gebruik te maken van entropie-regularisatie of adaptieve clipping. Deze methoden leggen externe beperkingen op die mogelijk ten koste gaan van de beloningsmaximalisatie.
2. De Sampling Temperatuur aan te passen: Het verhogen van de temperatuur tijdens het verzamelen van rollouts. Hoewel dit het gedragsbeleid (behavior policy) verandert, verandert het de onderliggende geparametriseerde distributie van het model niet, wat een mismatch creëert tussen het gedragsbeleid en het geoptimaliseerde beleid.

Deze benaderingen blijven extern aan de modelparameters. Het artikel stelt dat het exploratieve effect van temperatuur direct in de modelgewichten moet worden geïnternaliseerd om een betere initialisatie voor voortgezette RL te bieden.

2. Methodologie: TS-OPSD

De auteurs stellen Temperature-Scaled On-Policy Self-Distillation (TS-OPSD) voor, een lichtgewicht "policy reheating"-methode die ontworpen is om entropie te herstellen zonder externe docenten, geprivilegieerde data of extra inferentiekosten.

Kernmechanisme

TS-OPSD werkt als een interventie tussen RL-stadia. Vertrekkend vanuit een entropie-geïnstorte checkpoint ($\theta_{RL-collapse}$), construeert de methode een zelf-docent (self-teacher) en distilleert een gladdere distributie terug naar het studentmodel ($\theta_{reheat}$).

1. Constructie van de Zelf-Docent: Gegeven de huidige logits van de student $z_\theta$, wordt een zelf-docent-distributie $p_T$ gegenereerd door een hoog-temperatuur softmax-transformatie ($T > 1$) toe te passen op de logits. Cruciaal is dat de docent-logits een stop-gradient hebben en dat er geen extern model wordt gebruikt.
   $$p_T(v | x_{<t}) = \frac{\exp(z_{\theta,v}/T)}{\sum_u \exp(z_{\theta,u}/T)}$$
2. On-Policy Distillatie: De student wordt getraind om de Kullback-Leibler (KL) divergentie te minimaliseren tussen zijn huidige distributie $q_\theta$ en de hoog-temperatuur zelf-docent $p_T$, gebruikmakend van on-policy rollouts die uit $q_\theta$ zelf zijn gesampled.
   $$\mathcal{L} = D_{KL}(p_T \parallel q_\theta) \quad \text{(Forward KL)}$$
   $$\mathcal{L} = D_{KL}(q_\theta \parallel p_T) \quad \text{(Reverse KL)}$$

Structurele Eigenschappen

• Rangbehoud (Rank Preservation): Omdat $p_T$ een monotone machtstransformatie is van $q_\theta$ ($p_T(v) \propto q_\theta(v)^{1/T}$), blijft de token-rangorde behouden. De docent en de student delen dezelfde lokale ondersteuning en semantische voorkeursvolgorde; alleen de scherpte van de distributie verschilt.
• Internalisering: In tegen tegenover de standaard aanpassingen van de decoding-temperatuur, werkt TS-OPSD de modelgewichten bij om de gladdere, hogere-entropie gedragingen permanent te coderen. Dit elimineert de mismatch tussen training en inferentie.
• Gradiëntdynamiek:
  • Forward KL (FKL): De gradiënt neemt de vorm aan $q_\theta(v) - p_T(v)$. Dit voert "gestructureerde smoothing" uit, waarbij onderdrukte alternatieven opnieuw worden gewogen op basis van de eigen eerdere oordelen van het model.
  • Reverse KL (RKL): Onder zelf-temperatuurkoppeling is het minimaliseren van RKL lokaal gelijk aan het maximaliseren van de entropie van de student, waarbij de temperatuur slechts fungeert als een scalaire vermenigvuldiger op de gradiëntgrootte.

3. Belangrijkste Bijdragen

1. Policy Reheating Framework: Introduceert TS-OPSD, een methode om het exploratieve effect van temperatuur te internaliseren in de modelparameters, dienend als een post-collapse interventie.
2. Theoretisch Inzicht: Toont aan dat onder zelf-temperatuurkoppeling de mismatch tussen docent en student wordt geminimaliseerd (rangbehoud), en dat RKL reduceert tot een directe entropie-maximaliserende update.
3. Empirische Validatie: Toont aan dat het reheaten van een entropie-geïnstorte checkpoint een superieure initialisatie biedt voor voortgezette RL vergeleken met standaard voortgezette training of temperatuur-aanpassingen op rollout-niveau.
4. Analyse van Dynamiek: Karakteriseert het "entropie-sprong" fenomeen, waarbij entropieherstel in twee fasen plaatsvindt (langzame smoothing gevolgd door een scherpe discontinue sprong), en analyseert de omkeerbaarheid van deze parameter-updates.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op Qwen3-4B-Base en Qwen3-8B-Base met de DAPO-Math-17K dataset en geëvalueerd op MATH-500, AIME24/25/26, en OlymMATH.

• Prestatiewinst: TS-OPSD (specifiek met Forward KL, PR-FKL) presteerde consequent beter dan zowel de standaard DAPO-baseline als Rollout Reheating (DAPO w/ RR).
  • Op Qwen3-8B-Base bereikte PR-FKL een gemiddelde pass@8 van 54,74, vergeleken met 51,60 voor standaard DAPO en 52,66 voor Rollout Reheating.
  • De winst was het meest uitgesproken op moeilijkere benchmarks (AIME, OlymMATH).
• Entropieherstel: Na het reheaten konden de modellen de RL hervatten vanaf een initialisatie met hogere entropie en bereikten ze een hogere uiteindelijke beloning dan standaard DAPO, dat onderdrukt bleef.
• Vergelijking met Rollout Reheating: Hoewel het verhogen van de sampling-temperatuur tijdens rollouts (RR) niet-triviale winsten opleverde, leverde Policy Reheating (PR) superieure resultaten op. Dit suggereert dat het voordeel voortkomt uit de gladdere beleidsdistributie zelf die als een beter fundament voor optimalisatie dient, in plaats van enkel door ruisiger sampling.
• Behoud van Capaciteiten: Analyses toonden aan dat TS-OPSD ondersteunings-behoudende waarschijnlijkheidsherverdeling uitvoert:
  • Intermediaire Representaties: Lineaire CKA-scores tussen de base en de reheated lagen bleven nabij de 1,0, wat aangeeft dat de interne representaties grotendeels onveranderd bleven.
  • Redeneervaardigheid: De taaknauwkeurigheid op MATH-500 en GPQA bleef stabiel tijdens het hele distillatieproces, wat bevestigt dat het herstel van entropie de redeneervaardigheden niet heeft verslechterd.
  • Kandidaten-sets: De overlap in de top-10 tokens en de retentie van de waarschijnlijkheidsmassa waren nabij de 1,0, wat aangeeft dat de semantische voorkeuren van het model behouden bleven terwijl de marges werden verzacht.

5. Betekenis en Claims

Het artikel beweert dat entropieherstel kan worden ontkoppeld van het RL-doel en kan worden ingezet als een modulaire, tussenfase-operatie wanneer een beleid te deterministisch wordt.

• Modulaire Interventie: TS-OPSD biedt een eenvoudig, lichtgewicht mechanisme om de redeneringsgerichte RL uit te breiden voorbij de entropie-instortingsflessenhal zonder het kern-RL-algoritme te wijzigen.
• Stabiliteit in Data-arme Settings: In experimenten met verminderde hoeveelheden data (2048 prompts) werd standaard RL-training instabiel (gradiëntpieken, degradatie van beloning) door ernstige entropie-instorting. TS-OPSD fungeerde als een stabiliserend mechanisme, waardoor een langere, stabiele training mogelijk werd.
• Computationele Efficiëntie: De reheating-interventie voegt ongeveer 15–17% toe aan de totale RL-trainingstijd. Wanneer genormaliseerd voor batchgrootte, is de kosten per sample vergelijkbaar met of zelfs lager dan standaard RL-stappen.
• Beperkingen: De auteurs merken op dat de experimenten beperkt zijn tot specifieke modelmaten (4B/8B) en dat de stopcriteria voor reheating (entropie-drempelwaarde) momenteel empirisch in plaats van principieel zijn. Ze suggereren dat toekomstig werk de mogelijkheden kan verkennen van gekoppelde schema's van afwisselende OPSD en RL.

Samenvattend laat TS-OPSD zien dat het "internaliseren" van temperatuur via zelf-distillatie effectief de exploratieve capaciteit van ingestorte beleidsregels herstelt, waardoor voortgezette RL een hoger prestatieplafond kan bereiken terwijl de tijdens de training verworven redeneerstructuren behouden blijven.