Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation

Dit paper introduceert het Generalized On-Policy Distillation (G-OPD)-framework, dat door middel van beloningsextrapolatie en een flexibele referentiemodelkeuze de prestaties van studentenmodellen verbetert en zelfs de prestatiegrenzen van docentenmodellen kan doorbreken.

De kern

Samenvatting: "Leren van een Meester, maar dan nog Beter"

Stel je voor dat je een jonge, ambitieuze kok bent (de student) die wil leren koken van een beroemde, geprezen chef-kok (de meester).

In de wereld van kunstmatige intelligentie (AI) proberen wetenschappers vaak deze jonge kok te trainen door hem de recepten van de meester na te laten maken. Maar hoe doe je dat het beste?

Het oude probleem: "Kijk en nabootsen"

Vroeger was de methode simpel: de meester maakt een gerecht, en de student kijkt er naar en probeert het exact na te maken. Dit heet off-policy distillation.

• Het nadeel: De student leert alleen wat de meester deed, maar leert niet waarom het lekker is. Als de student later zelf een nieuw gerecht moet bedenken, faalt hij omdat hij niet heeft geoefend met zijn eigen fouten en successen.

De nieuwe methode: "On-Policy Distillation" (OPD)

Deze paper introduceert een slimme truc: On-Policy Distillation (OPD).
Hierbij maakt de student zelf een gerecht (een antwoord), en de meester kijkt er dan naar en zegt: "Nee, diep in je hart wist je dat dit ingrediënt beter was. Probeer het weer, maar denk aan mijn advies."
De student leert dus van zijn eigen pogingen, maar krijgt direct feedback van de meester. Dit werkt al heel goed, maar de auteurs van deze paper dachten: "Kunnen we dit niet nog slimmer maken?"

De grote doorbraak: G-OPD (De "Super-Kookcursus")

De auteurs (Wenkai Yang en collega's) hebben een nieuwe methode bedacht genaamd G-OPD. Ze hebben twee magische ingrediënten toegevoegd aan het recept:

1. De "Versterkings-knop" (Reward Extrapolation)

Stel je voor dat de meester zegt: "Dit gerecht is een 8/10."

• Normale methode: De student probeert een 8 te halen.
• Deze paper's methode (ExOPD): De student krijgt een knop die hij op 1,25 kan zetten. Hij denkt dan: "Als de meester een 8 vindt, dan moet ik proberen om een 10 te halen!"

Dit noemen ze Reward Extrapolation. In plaats van alleen te doen wat de meester doet, leert de student om beter te zijn dan de meester.

• Het resultaat: In tests bleek dat de studenten die deze "versterkings-knop" gebruikten, zelfs betere wiskundige oplossingen en code schreven dan de oorspronkelijke meester-chef. Ze hebben de grenzen van de meester doorbroken!

2. De "Referentie-kookboek" (Reference Model)

Bij het leren is het belangrijk om te weten: "Vanuit welk startpunt kom ik?"

• Standaard: De student vergelijkt zijn nieuwe gerecht met zijn oude, beginnende recepten.
• De slimme truc: Als de student een heel groot meester kopieert, is het soms beter om te vergelijken met het oorspronkelijke recept van de meester (voordat hij zijn eigen trucs leerde). Dit heet Reward Correction.
• Het voordeel: Het geeft een scherpere, zuivere feedback. Het is alsof je niet vergelijkt met je eigen beginnende kookkunsten, maar met de pure basis van de meester. Dit werkt nog beter, maar kost wel meer tijd en energie om te berekenen.

Wat hebben ze bewezen?

De auteurs hebben dit getest op twee moeilijke taken:

1. Wiskunde: Het oplossen van complexe wiskundepuzzels.
2. Programmeren: Het schrijven van computercode.

De resultaten:

• Meesters samenvoegen: Stel je hebt een meester die goed is in wiskunde en een andere die goed is in programmeren. Met hun methode konden ze één student maken die beide vaardigheden beter beheerste dan de twee aparte meesters.
• Groot naar Klein: Ze konden ook een heel slimme, grote AI (de meester) "in een klein pakketje" (de student) stoppen, zodat de kleine AI bijna net zo slim werd als de grote, maar dan nog sneller.

Conclusie in één zin

Deze paper laat zien dat je AI-modellen niet alleen kunt leren door ze te laten nabootsen, maar dat je ze kunt trainen om beter te zijn dan hun leraar, door een slimme "versterkings-knop" te gebruiken die hen uitdaagt om de grenzen van het mogelijke te verleggen.

Het is alsof je een leerling niet alleen laat zien hoe je een fiets rijdt, maar hem leert hoe je een racefiets rijdt, en hem zelfs de vaardigheden geeft om de wereldkampioen te verslaan.

Technische samenvatting

Probleemstelling

On-Policy Distillation (OPD) is een effectieve post-trainingparadigma voor Large Language Models (LLM's), waarbij een studentmodel leert van de logit-distributie van een docentmodel op trajecten die door de student zelf zijn gegenereerd. Hoewel OPD empirisch superieur presteert aan off-policy methoden (zoals Supervised Fine-Tuning op docentdata) en vaak beter is dan traditionele Reinforcement Learning (RL), blijft het mechanistische inzicht beperkt.
De huidige OPD-methoden hebben twee belangrijke beperkingen:

1. Ze zijn strikt gebonden aan een vaste weging tussen de beloningsterm (reward) en de Kullback-Leibler (KL) regularisatie (altijd 1:1).
2. Ze gebruiken doorgaans een vaste referentiemodel (de initiële staat van de student), wat de flexibiliteit beperkt bij het optimaliseren van de leerdoelen.
   Dit beperkt het potentieel van OPD om de prestaties van de docent te overtreffen, vooral in scenario's zoals het samenvoegen van kennis van meerdere domeinspecialisten of het distilleren van een grote docent naar een kleinere student.

Methodologie: Generalized On-Policy Distillation (G-OPD)

De auteurs stellen eerst een theoretische link vast tussen OPD en dense RL met een KL-beperking. Ze tonen aan dat standaard OPD een speciaal geval is van RL waarbij de beloning en de KL-regularisatie altijd gelijk gewogen zijn en de referentie willekeurig gekozen kan worden.

Op basis hiervan introduceren ze het G-OPD-framework, dat het standaard OPD-doelwit generaliseert door twee nieuwe componenten toe te voegen:

1. Een schaalbaarheidsfactor voor de beloning ($\lambda$): Dit controleert de relatieve weging van de beloningsterm ten opzichte van de KL-regularisatie.
2. Een flexibel referentiemodel ($\pi_{ref}$): Dit kan elk model zijn, niet beperkt tot de initiële student.

De geoptimaliseerde doelstelling wordt als volgt geformuleerd:
$$J_{G-OPD}(\theta) = \max_{\theta} \mathbb{E}_{x \sim D, y \sim \pi_\theta(\cdot|x)} \left[ \lambda \log \frac{\pi^*(y|x)}{\pi_{ref}(y|x)} - D_{KL}(\pi_\theta(y|x) \parallel \pi_{ref}(y|x)) \right]$$

Waarbij:

• $\pi^*$ het docentmodel is.
• $\pi_\theta$ het studentmodel is.
• $\pi_{ref}$ het referentiemodel is.
• $\lambda$ de schaalfactor is.

Belangrijke regimes binnen G-OPD:

• Reward Interpolation ($0 < \lambda < 1$): Het studentmodel gedraagt zich als een lineaire interpolatie tussen het referentiemodel en de docent.
• Reward Extrapolation ($\lambda > 1$): Dit regime, genaamd ExOPD, stimuleert het studentmodel om verder te gaan dan de log-probabiliteiten van de docent door een extra verschuivingsterm toe te voegen. Dit kan leiden tot prestaties die de grenzen van de docent overstijgen.
• Reward Correction (voor Strong-to-Weak distillatie): In het geval van het distilleren van een grote docent naar een kleine student, stellen de auteurs voor om als referentiemodel niet de basis van de student, maar de pre-RL-versie van de docent te gebruiken. Dit vermindert ruis in het beloningssignaal veroorzaakt door de kenniskloof tussen de twee modellen.

Belangrijkste Bijdragen

1. Theoretische Unificatie: Het aantonen dat OPD een speciaal geval is van dense RL en het generaliseren ervan tot een flexibeler framework (G-OPD).
2. Extrapolatie van Beloningen (ExOPD): Het ontdekken dat het instellen van $\lambda > 1$ (extrapolatie) consequent leidt tot betere prestaties dan standaard OPD. Dit stelt studenten in staat om de prestatiegrens van de docent te doorbreken.
3. Multi-Teacher Samenvoeging: Het aantonen dat ExOPD in staat is om een uniek studentmodel te creëren dat de prestaties van meerdere domeinspecialisten (docenten) combineert en deze allemaal overtreft.
4. Correctie bij Strong-to-Weak Distillatie: Het introduceren van een "reward correction" techniek waarbij het pre-RL-model van de docent als referentie dient, wat de distillatie-efficiëntie van grote naar kleine modellen aanzienlijk verbetert.

Resultaten

De auteurs hebben uitgebreide experimenten uitgevoerd op taken voor wiskundig redeneren (bijv. AIME, HMMT) en codegeneratie (bijv. HumanEval, MBPP).

• Same-Sized Distillatie (Student = Docent):

  • Extrapolatie werkt: Met $\lambda = 1.25$ (ExOPD) overtrof het studentmodel consequent zowel de standaard OPD als de domeindocent in wiskunde en code.
  • Stabiliteit: Te hoge waarden van $\lambda$ (bijv. 1.5) leidden tot instabiliteit, maar een matige extrapolatie was optimaal.
  • Multi-Teacher: Bij het samenvoegen van een wiskunde-docent en een code-docent (beide afgeleid van dezelfde basis), slaagde ExOPD erin een uniek model te trainen dat op alle benchmarks beter presteerde dan beide afzonderlijke docenten. Standaard OPD en SFT faalden hierin of bleven beperkt tot de prestaties van de docenten.

• Strong-to-Weak Distillatie (Grote Docent $\to$ Kleine Student):

  • ExOPD presteerde aanzienlijk beter dan standaard OPD en SFT bij het distilleren van een 30B-parameter model naar een 1.7B of 4B model.
  • Reward Correction: Het gebruik van de pre-RL-versie van de docent als referentiemodel (in plaats van de student-basis) leverde verdere verbeteringen op, hoewel dit meer rekenkracht vereist.

Betekenis en Impact

Dit werk biedt een fundamenteel nieuw perspectief op on-policy distillatie. Het toont aan dat distillatie niet beperkt hoeft te zijn tot het "nabootsen" van een docent, maar dat het door het manipuleren van de beloningsweging (extrapolatie) mogelijk is om beyond-teacher performance te bereiken.

De belangrijkste implicaties zijn:

• Efficiëntie: Het biedt een manier om de kennis van meerdere gespecialiseerde RL-modellen efficiënter te combineren dan met traditionele methoden.
• Schalbaarheid: Het maakt het mogelijk om de capaciteiten van zeer grote modellen over te dragen naar kleinere, efficiëntere modellen met hogere nauwkeurigheid dan eerder mogelijk was.
• Theoretische Diepgang: Het verbindt het veld van distillatie en RL dichter bij elkaar, wat nieuwe richtingen opent voor toekomstig onderzoek naar beloningsschaling en referentiemodel-selectie.

Kortom, ExOPD demonstreert dat door slimme aanpassing van de trainingsdoelstellingen, studenten niet alleen docenten kunnen kopiëren, maar zelfs kunnen overtreffen.