Self-Distilled Policy Gradient

Technische Samenvatting: Self-Distilled Policy Gradient (SDPG)

Probleemstelling
Reinforcement Learning met Verifieerbare Beloningen (RLVR) is een standaard geworden voor het post-trainen van Large Language Models (LLM's) op redeneertaken zoals wiskunde en codegeneratie. Algoritmen zoals Group Relative Policy Optimization (GRPO) optimaliseren tegen regelgebaseerde uitkomstbeloningen (bijv. correct/incorrect antwoorden) zonder dat daar menselijke voorkeurannotaties voor nodig zijn. Echter, RLVR kampt met twee primaire beperkingen:

1. Sparse Rewards (IJle beloningen): Het beloningssignaal wordt toegepast op sequentieniveau, wat het moeilijk maakt om krediet toe te wijzen aan specifieke redeneerstappen (tokens), wat leidt tot inefficiënte exploratie.
2. Trainingsinstabiliteit: De vroege trainingsstadia lijden vaak aan negatieve advantages, wat instabiliteit en potentiële degradatie van fundamentele taalvaardigheden veroorzaakt.

Hoewel On-Policy Distillation (OPD) is voorgesteld om dichte, token-niveau supervisie te bieden, vertrouwen traditionele benaderingen op externe, grotere leermodellen (teacher models), wat hoge geheugenkosten met zich meebrengt en mismatch in distributies introduceert. Recente "zelf-distillatie" methoden, waarbij het model zowel student als teacher fungeert met behulp van geprivilegieerde context (bijv. grondwaarheid-oplossingen), pakken het geheugenprobleem aan maar lopen het risico op het versterken van lokaal plausibele tokens op globaal incorrecte trajecten als dit niet goed wordt gecontroleerd. Bovendien kan pure zelf-distillatie leiden tot mode collapse en verminderde exploratie.

Methodologie: SDPG-framework
Het artikel stelt Self-Distilled Policy Gradient (SDPG) voor, een framework dat exacte full-vocabulary on-policy zelf-distillatie integreert met KL-geregulariseerde policy optimalisatie. De kernopdracht combineert drie componenten:

1. Outcome-Based Policy Gradient ($L_{out}$):
   SDPG behoudt het ijle binaire beloningssignaal van een verifier. Het maakt gebruik van een group-relative advantage ($A_{out}$) die binnen een groep van rollouts wordt genormaliseerd om het leren te stabiliseren, vergelijkbaar met GRPO, maar geformuleerd als een directe REINFORCE-stijl surrogate loss zonder de PPO-stijl importance ratio clipping.

2. Exact Full-Vocabulary On-Policy Distillation ($L_{OPD}$):
   In plaats van tokens te samplen om het teacher-signaal te benaderen, berekent SDPG de exacte reverse Kullback-Leibler (KL) divergentie tussen de student-distributie ($p_t$) en een geprivilegieerde teacher-distributie ($q_t$), geconditioneerd op aanvullende context $c$ (bijv. oplossingen of redeneerpaden).

   • Gradient Identity: De auteurs leiden een cruciaal theoretisch resultaat af (Propositie 3.1): de gradiënt van de reverse-KL full-vocabulary loss met betrekking tot de student-policy is lokaal identiek aan een detached-sampling policy gradient update. De "advantage" in deze update is een gecentreerde log-ratio van de teacher- en student-waarschijnlijkheden.
   • Gating Mechanisme: Om te voorkomen dat de geprivilegieerde teacher incorrecte trajecten versterkt, wordt de distillatie-loss gegate door de outcome advantage. Distillatie wordt alleen toegepast wanneer de rollout als correct is geverifieerd ($A_{out} > 0$).

3. Reference Policy Regularization ($L_K$):
   Om te voorkomen dat de policy te ver afwijkt van zijn initiële capaciteiten of instort, verankert SDPG de bijgewerkte policy aan een vaste reference policy ($\pi_{ref}$). Het artikel gebruikt Unnormalized KL (UKL) divergentie (forward of reverse) voor deze regularisatie, wat elegantere en symmetrische gradiënten biedt vergeleken met standaard genormaliseerde KL, vooral bij het werken met ongenormaliseerde maten.

Stabilisatiestrategieën
Om de ruis van het geprivilegieerde signaal en de afweging tussen exploratie en exploitatie te beheersen, introduceert SDPG twee stabilisatoren:

• Positive Advantage Gating: Zorgt ervoor dat distillatie alleen plaatsvindt op trajecten die door de verifier worden onderschept.
• Warmup-Decay Schedule ($\beta(k)$): De distillatiecoëfficiënt $\beta$ wordt geleidelijk verhoogd (warmup) om het model te laten de juiste trajecten vinden voordat het vertrouwt op distillatie, en vervolgens verlaagd (decay) tegen het einde van de training. Dit voorkomt dat het geprivilegieerde signaal het model te sterk beperkt zodra het de noodzakelijke kennis heeft geïnternaliseerd, wat de finale fase van exploratie stimuleert.

Belangrijkste Bijdragen

1. Theoretische Afleiding: Het artikel leidt een exacte lokale policy-gradient vorm af voor reverse-KL full-vocabulary on-policy distillatie, waarbij wordt aangetoond dat dit gelijk is aan een detached-sampling update met een gecentreerde log-ratio advantage.
2. Nieuw Doel (Objective): Het stelt een verenigde doelstelling voor die binaire outcome rewards, gated full-vocabulary geprivilegieerde distillatie en unnormalized KL regularisatie tegen een vaste reference combineert.
3. Stabilisatietechnieken: De introductie van positive-advantage gating en een warmup-decay schema voor de distillatiecoëfficiënt om de balans te vinden tussen dichte supervisie en exploratie.

Experimentele Resultaten
De auteurs evalueerden SDPG op wiskundige redeneerbenchmarks (AIME2024, AIME2025, AMC23) met behulp van Qwen3-4B.

• Prestaties: Beide SDPG-varianten (gebruikmakend van Unnormalized Forward KL en Unnormalized Reverse KL) presteerden beter dan de baseline GRPO en de zelf-distillatie baseline RLSD op alle benchmarks.
• Stabiliteit: SDPG-varianten bereikten significanter hogere reward-plateaus dan GRPO. Opvallend genoeg behield SDPG een hogere actor-entropie tijdens de training vergeleken met RLSD, dat een entropie-instorting vertoonde (een teken van mode collapse) rond stap 250.
• Efficiëntie: SDPG-methoden produceerden responslengtes die voldoende waren voor meerstaps-redeneren, maar compacter dan de verbale outputs van GRPO.

Significantie
Het artikel claimt dat SDPG erin slaagt de kloof te overbruggen tussen sparse-reward RLVR en dichte-supervisie zelf-distillatie. Door zelf-distillatie te behandelen als een policy gradient update en het geprivilegieerde signaal strikt te controleren via gating en scheduling, bereikt SDPG superieure stabiliteit en prestaties. Het demonstreert dat dichte, token-niveau supervisie afgeleid van de eigen geprivilegieerde context van het model, effectief redeneermodellen kan begeleiden zonder de geheugenoverhead van externe teachers of de instabiliteit van pure zelf-distillatie.