Entropy-Aware On-Policy Distillation of Language Models

Deze paper introduceert Entropy-Aware On-Policy Distillation, een methode die de standaard reverse KL-divergentie aanvult met forward KL bij hoge teacher-entropie om de generatiediversiteit te behouden en de prestaties van taalmodellen op wiskundige redeneringstaken aanzienlijk te verbeteren.

De kern

Stel je voor dat je een jonge, slimme student (het student-model) hebt die wil leren wiskunde van een wereldberoemde professor (het leraar-model). De student moet niet alleen de juiste antwoorden leren, maar ook begrijpen hoe de professor denkt, inclusief de momenten waarop de professor twijfelt.

Dit artikel beschrijft een nieuwe manier om die kennis over te dragen, genaamd EOPD (Entropy-Aware On-Policy Distillation). Hier is hoe het werkt, vertaald naar alledaags taal:

1. Het oude probleem: De "Mode-Seeking" valkuil

Tot nu toe leerden studenten vaak door simpelweg te kijken naar het antwoord dat de professor het meest waarschijnlijk vond. In de wereld van AI noemen we dit "Reverse KL".

• De analogie: Stel je voor dat de professor een vraag krijgt: "Wat is de beste manier om naar het station te komen?"
  • Soms is het antwoord duidelijk: "Neem de trein." (De professor is 100% zeker).
  • Soms is het antwoord minder duidelijk: "Je kunt de bus nemen, de fiets, of lopen, afhankelijk van het weer." (De professor is onzeker en geeft meerdere goede opties).

Het oude systeem (Reverse KL) dwong de student om alleen de optie te kiezen die de professor het vaakst noemde. Als de professor twijfelde tussen bus, fiets en lopen, koos de student steevast voor de bus, omdat die net iets vaker werd genoemd.

• Het gevolg: De student werd een "one-trick pony". Hij verloor de diversiteit. Hij leerde niet dat er meerdere goede oplossingen zijn. In moeilijke wiskundeproblemen, waar vaak meerdere wegen naar het antwoord leiden, werd de student hierdoor dom en onzeker. Hij probeerde de professor na te bootsen, maar verloor de essentie van de twijfel.

2. Het nieuwe idee: EOPD (De slimme leraar)

De auteurs van dit artikel zeggen: "Wacht, we moeten kijken naar hoe zeker de leraar is."

Ze introduceren een slimme strategie die EOPD heet. Het werkt als een slimme coach die twee verschillende methodes gebruikt, afhankelijk van de situatie:

• Situatie A: De leraar is 100% zeker.
  • Analogie: De leraar zegt: "Het antwoord is 42."
  • Actie: De student leert hierdoor gewoon na te bootsen (Reverse KL). Dit is snel en efficiënt.
• Situatie B: De leraar is onzeker (hoge entropie).
  • Analogie: De leraar zegt: "Het kan A, B of C zijn, allemaal zijn ze mogelijk."
  • Actie: Hier schakelt de student over op een andere methode (Forward KL). In plaats van alleen naar het meest waarschijnlijke antwoord te kijken, leert de student nu om alle mogelijke goede opties te omarmen. Hij leert dat de leraar twijfelt, en dat die twijfel belangrijk is.

3. Waarom is dit zo goed?

Stel je voor dat je een spoorbaan bouwt.

• De oude methode bouwde maar één spoor, zelfs als er drie wegen naar de bestemming waren. Als die ene weg geblokkeerd was, kon de trein niet verder.
• De nieuwe methode (EOPD) bouwt meerdere sporen op de plekken waar de leraar twijfelde.

Dit heeft twee grote voordelen:

1. Meer creativiteit: De student leert dat er meerdere goede oplossingen zijn. Hij raakt niet vast in één denkpatroon.
2. Stabiel leren: Als de leraar twijfelt, geeft de oude methode de student een verwarrend signaal ("Kies maar wat!"). De nieuwe methode zegt duidelijk: "Oké, hier zijn drie goede opties, onthoud ze allemaal."

4. De resultaten in de praktijk

De auteurs hebben dit getest op wiskundeproblemen (zoals de AIME en AMC wedstrijden).

• Resultaat: De studenten die met de nieuwe methode (EOPD) leerden, werden niet alleen beter in het vinden van het juiste antwoord, maar ze waren ook veel creatiever.
• Ze konden meer verschillende manieren bedenken om een probleem op te lossen.
• Op de moeilijkste vragen (waar de leraar het meest twijfelde) scoorden de nieuwe studenten veel beter dan de oude studenten.

Samenvatting in één zin

Deze paper leert ons dat om een slimme AI te maken, we hem niet alleen de "juiste antwoorden" moeten geven, maar hem ook moeten leren omgaan met onzekerheid, zodat hij niet alleen één weg kent, maar het hele landschap van mogelijke oplossingen begrijpt.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het artikel "Entropy-Aware On-Policy Distillation of Language Models" in het Nederlands.

1. Het Probleem: Beperkingen van On-Policy Distillatie

Het artikel richt zich op On-Policy Distillatie, een methode waarbij een kleiner "student" model leert van een groter "leraar" model door te trainen op sequenties die door de student zelf zijn gegenereerd (in plaats van op statische data van de leraar). Dit lost het probleem van distributiemismatch op dat optreedt bij traditionele off-policy distillatie.

De huidige standaardmethode gebruikt echter Reverse KL-divergentie als trainingsdoel. Hoewel dit efficiënt is voor het kopiëren van de meest waarschijnlijke antwoorden (de "modi" van de leraar), heeft het twee kritieke nadelen:

1. Verlies van Diversiteit: Reverse KL is een "mode-seeking" doelstelling. Het straft de student af voor het genereren van tokens die de leraar onwaarschijnlijk vindt, maar negeert de leraar-modi die de student niet bezoekt. Dit leidt tot een instorting van de diversiteit, vooral in complexe redeneertaken waar meerdere geldige oplossingspaden bestaan.
2. Instabiliteit bij Hoge Entropie: Wanneer de leraar onzeker is (hoge entropie, bijvoorbeeld bij beslissingspunten met meerdere valide opties), levert Reverse KL instabiele leer signalen op. De student convergeert niet goed en faalt om de onzekerheidsstructuur van de leraar te behouden.

2. Methodologie: Entropy-Aware On-Policy Distillation (EOPD)

De auteurs introduceren EOPD, een raamwerk dat de stabiliteit van Reverse KL combineert met de diversiteitbehoudende eigenschappen van Forward KL, afhankelijk van de onzekerheid van de leraar.

Kernidee:
De methode past het trainingsdoel dynamisch aan op basis van de token-level entropie van de leraar:

• Bij lage entropie (Leraar is zeker): Er wordt Reverse KL gebruikt. Dit zorgt voor efficiënte, snelle convergentie naar de dominante antwoorden van de leraar.
• Bij hoge entropie (Leraar is onzeker): Er wordt Forward KL geactiveerd. Forward KL is een "mode-covering" doelstelling; het straft de student af als hij waarschijnlijkheden toekent aan tokens die de leraar waarschijnlijk vindt, zelfs als de student die tokens niet zelf genereert. Dit dwingt de student om de volledige spreiding van mogelijke antwoorden van de leraar te behouden.

Implementatie Details:

• Selectieve Toepassing: Een drempelwaarde ($\tau$) bepaalt wanneer Forward KL wordt toegepast. Alleen wanneer de entropie van de leraar ($H_{te}$) boven deze drempel ligt, wordt de Forward KL-term toegevoegd aan de loss-functie.
• Efficiëntie: Om de computationele kosten van Forward KL (die normaal gesproken de hele vocabulaire vereist) te beperken, wordt Forward KL alleen berekend over de top-k tokens van de leraar. Dit behoudt de belangrijkste waarschijnlijkheden zonder de "lange staart" van zeer onwaarschijnlijke tokens te hoeven modelleren.
• Integratie: De methode is gebaseerd op PPO (Proximal Policy Optimization) met een "clipped" Reverse KL loss, waarbij de Forward KL term conditioneel wordt toegevoegd.

3. Belangrijkste Bijdragen

1. Analyse van Diversiteitsverlies: Het artikel toont systematisch aan dat standaard on-policy distillatie (Reverse KL) de diversiteit drastisch vermindert. In experimenten behield de student slechts 6,8% van de hoge-entropie tokens van de leraar, terwijl de leraar zelf 18,5% had.
2. Instabiliteitsbewijs: Via een gecontroleerd "toy experiment" wordt aangetoond dat Reverse KL leidt tot fluctuerende top-10 sets en frequente veranderingen in de meest waarschijnlijke token bij hoge leraar-entropie, wat convergeert tot een instabiel leerproces.
3. Het EOPD-raamwerk: De introductie van een hybride loss-functie die de kracht van beide KL-divergenties benut zonder de computationele overhead van puur Forward KL.
4. Empirische Verbeteringen: Aantonen dat het expliciet modelleren van leraar-onzekerheid essentieel is voor effectieve kennisoverdracht in redeneertaken.

4. Resultaten

De methode werd getest op zes wiskundige redeneerbenchmarks (zoals AIME24, AIME25, AMC23, MATH500) met Qwen3-modellen (0.6B, 1.7B en 4B) als studenten en Qwen3-8B als leraar.

• Prestatieverbetering: EOPD overtreft consistent de baselines (KD, GRPO, en standaard OPD).
  • Voor het Qwen3-4B-Base model werd een toename van +5,05% in Pass@8 (kans op minstens één correct antwoord bij 8 pogingen) behaald ten opzichte van de standaard OPD.
  • Voor het Qwen3-1.7B-Base model was de winst +2,39% in Pass@8.
• Behoud van Diversiteit: EOPD behoudt aanzienlijk meer waarschijnlijkheidsmassa in hoge-entropie gebieden, wat dichter bij de verdeling van de leraar ligt dan bij standaard OPD.
• Out-of-Domain Generalisatie: Hoewel getraind op wiskundedata, presteerde EOPD ook beter op algemene redeneertaken (GPQA-Diamond, MMLU-Pro) en instructievolgende taken (AlpacaEval 2.0), wat suggereert dat het overbrengen van redeneerstrategieën en onzekerheid breder toepasbaar is.
• Pass@k Analyse: De prestatieverbetering van EOPD nam toe naarmate het aantal samples ($k$) toenam, wat aantoont dat het model effectiever diverse redeneerpaden exploreert.

5. Betekenis en Conclusie

De studie concludeert dat Reverse KL alleen onvoldoende is voor het distilleren van complexe redeneermodellen, omdat het de nuance en onzekerheid van de leraar verwijdert. Door Forward KL selectief in te zetten bij momenten van hoge onzekerheid, kan EOPD de diversiteit van de generatie behouden en stabiliteit garanderen.

Dit heeft grote betekenis voor de ontwikkeling van efficiënte, inzetbare taalmodellen:

• Het maakt het mogelijk om de redeneercapaciteiten van grote, dure modellen over te dragen naar kleinere, snellere modellen zonder dat deze "oververzekerd" worden of complexe oplossingspaden verliezen.
• Het biedt een nieuwe richting voor Reinforcement Learning en distillatie, waarbij het expliciet modelleren van onzekerheid net zo belangrijk is als het kopiëren van de juiste antwoorden.

Kortom, EOPD biedt een evenwichtige oplossing die de efficiëntie van on-policy training combineert met de robuustheid van het behoud van de leraar-distributie.