Scaling Reasoning Efficiently via Relaxed On-Policy Distillation

Titel: Scaling Reasoning Efficiently via Relaxed On-Policy Distillation (REOPOLD)

Auteurs: Jongwoo Ko, Sara Abdali, Young Jin Kim, Tianyi Chen, Pashmina Cameron.

---

1. Het Probleem

Grote Taalmodellen (LLMs) hebben aanzienlijke redeneercapaciteiten ontwikkeld door middel van Reinforcement Learning (RL) en schaling tijdens het testen (test-time scaling), zoals geïllustreerd door modellen als OpenAI o1 en DeepSeek-R1. Het repliceren van dit succes in Kleine Taalmodellen (SLMs) is echter problematisch:

• Beperkte capaciteit: SLMs worstelen met directe beloningsoptimalisatie (RL), waardoor standaard RL-methoden vaak inefficiënt of onstabiel zijn.
• Instabiliteit bij On-Policy Distillatie: Bestaande methoden voor on-policy distillatie (waarbij een student leert van zijn eigen trajecten onder begeleiding van een leraar) lijden aan ernstige instabiliteit. Dit leidt vaak tot:
  • Negatieve overdracht (Negative Transfer): De prestaties van de student deterioreren in plaats van te verbeteren ten opzichte van het basismodel.
  • Vroege convergentie (Entropy Collapse): De entropie van het beleid daalt te snel, waardoor het model vastloopt in een smal scala aan antwoorden en niet meer exploreert.
  • Sensitiviteit: Standaard methoden vereisen vaak dat de leraar en student vergelijkbare maten hebben, wat de praktische bruikbaarheid beperkt.

2. Methodologie: REOPOLD

De auteurs introduceren REOPOLD (Relaxed On-Policy Distillation), een raamwerk dat de strikte imitatie-eisen van standaard distillatie versoepelt door het te interpreteren als een vorm van beleidsoptimalisatie (RL). De kernidee is dat de log-likelihood-ratio tussen leraar en student fungeert als een token-beloning.

REOPOLD lost de optimalisatieproblemen op via drie hoofdcomponenten:

A. Theoretische Basis: Stop-Gradient als Controlevariabele

De auteurs tonen theoretisch aan dat on-policy distillatie equivalent is aan een beleidsgradiëntmethode waarbij de beloning wordt gedefinieerd door de log-likelihood-ratio. Door een stop-gradient operatie toe te passen op de beloningsterm, wordt deze behandeld als een vaste, externe signaal. Dit fungeert als een controlevariabele die de variantie in de gradiënt-schattingen verlaagt en training stabiliseert.

B. Mixture-Based Reward Clipping (Beloningsclipping)

Standaard distillatie lijdt aan een "heavy-tailed" verdeling van negatieve beloningen: wanneer de leraar een token een zeer lage waarschijnlijkheid toekent, explodeert de negatieve beloning ($\to -\infty$), wat de training destabiliseert.

• Oplossing: REOPOLD gebruikt een theoretisch afgeleide ondergrens gebaseerd op een convexe mengselverdeling van leraar en student.
• Mechanisme: Negatieve beloningen worden afgekapt (geclipt) op een drempelwaarde $\log \frac{\lambda}{1-\lambda}$. Dit voorkomt dat extreme negatieve signalen de gradiënten domineren, terwijl het de "mode-seeking" aard van distillatie behoudt.

C. Entropy-Guided Token-Level Dynamic Sampling

De meeste tokens in een respons zijn deterministisch (laag entropie) en bieden weinig leerwaarde omdat leraar en student hier al overeenkomen.

• Oplossing: Het model selecteert dynamisch alleen tokens met hoge entropie (hoge onzekerheid) voor het berekenen van gradiënten.
• Mechanisme: Een masker filtert tokens met lage entropie eruit. Dit verhoogt de effectieve steekproefgrootte en voorkomt dat het model tijd verspillen aan het leren van triviaal bepaalde tokens.

D. Multi-Stage Training: Exploration-to-Refinement

REOPOLD implementeert een tweefasenstrategie om de balans tussen exploratie en exploitatie te beheersen:

1. Exploratiefase: In het begin worden sterk negatieve beloningen gefilterd (via reward clipping). Dit voorkomt vroege entropie-collaps en laat het model diverse oplossingspaden verkennen (vergelijkbaar met Supervised Fine-Tuning).
2. Refinement-fase: Later in de training wordt de strategie gewijzigd naar entropy-guided sampling. Hier worden negatieve feedbacksignalen weer toegestaan op kritieke, onzekere tokens om het beleid te verfijnen en te convergeren naar hoogwaardige trajecten.

3. Belangrijkste Bijdragen

1. Diagnose van Instabiliteit: Het paper identificeert zwaarstaartige negatieve beloningen en inefficiënte signalen als de primaire oorzaken van instabiliteit in on-policy distillatie.
2. REOPOLD Framework: Een unificerend raamwerk dat reward clipping, token-level dynamische sampling en multi-stage training integreert om de leerkracht van de leraar "temperend" en "selectief" toe te passen.
3. State-of-the-Art Efficiëntie: Empirisch bewijs dat REOPOLD superieure sample-efficiëntie biedt en test-time scaling mogelijk maakt voor kleine modellen.

4. Resultaten

De auteurs evalueren REOPOLD op wiskundig, visueel en agentisch redeneren:

• Sample Efficiëntie:

  • Op wiskundige benchmarks (AIME-25) bereikt REOPOLD met een 1.5B model een prestatie die 6.7x tot 12x efficiënter is dan bestaande RL-methoden (zoals ProRL) en standaard RKL.
  • Het overtreft recentere RL-benaderingen aanzienlijk.

• Test-Time Scaling:

  • Op visuele redeneertaken (Geometry3K, MathVerse) toont REOPOLD superieure schaalbaarheid.
  • Een 7B studentmodel met REOPOLD kan de prestaties van een 32B leraarmodel benaderen of evenaren, met een 3.32x versnelling in inferentiesnelheid.
  • In tegenstelling tot standaard RKL, die instabiel wordt bij grotere studenten (7B), blijft REOPOLD stabiel en verbetert het consistent.

• Robuustheid:

  • REOPOLD presteert consistent beter dan RKL, ongeacht de grootte van het leraarmodel (van 7B tot 32B).
  • Het voorkomt negatieve overdracht en entropie-collaps, wat vaak voorkomt bij standaard distillatie.

5. Significantie

Dit paper is significant omdat het een fundamentele beperking in het trainen van redenerende kleine modellen oplost.

• Paradigmaverschuiving: Het toont aan dat strikte imitatie (van standaard distillatie) schadelijk kan zijn voor redeneren en dat een "ontspannen" (relaxed) benadering, geleid door RL-principes, noodzakelijk is.
• Toegankelijkheid: Het maakt het mogelijk om de geavanceerde redeneercapaciteiten van enorme modellen (32B+) over te dragen naar compacte modellen (3B-7B) met veel minder rekentijd en data, wat cruciaal is voor de implementatie van redenerende AI in productieomgevingen met beperkte middelen.
• Stabiliteit: Het biedt een robuustere oplossing dan directe RL voor SLMs, die vaak te gevoelig zijn voor de variabiliteit van beloningen.

Kortom, REOPOLD biedt een efficiëntere, stabielere en schaalbaarder route om redeneercapaciteiten naar kleinere modellen te brengen door de leerprocessen van de leraar selectief en dynamisch toe te passen in plaats van blindelings te imiteren.