Hybrid Policy Distillation for LLMs

Titel: Hybrid Policy Distillation (HPD) voor Large Language Models (LLMs)

Auteurs: Wenhong Zhu, Ruobing Xie, Rui Wang, Pengfei Liu.

---

1. Het Probleem

De prestaties van Large Language Models (LLMs) zijn sterk toegenomen door het vergroten van parameteraantallen en data, maar dit brengt aanzienlijke kosten met zich mee voor opslag en rekenkracht. Knowledge Distillation (KD) is een veelbelovende techniek om deze grote "leraar"-modellen te comprimeren naar kleinere "student"-modellen.

Echter, bestaande KD-methoden kampen met fundamentele uitdagingen die voortvloeien uit de complexe interactie tussen drie ontwerpaspecten:

1. Richting van divergentie:
   • Forward KL (FKLD): Bevordert het dekken van alle modi (mode coverage) van de leraar, maar leidt vaak tot te gladde (over-gesmoorde) voorspellingen.
   • Reverse KL (RKLD): Richt zich op het vinden van de hoogste waarschijnlijkheden (mode-seeking), maar kan instabiel worden als het verschil tussen leraar en student groot is, en negeert geldige, maar minder waarschijnlijke output.
2. Optimalisatiestrategie: Het gebruik van verliesfuncties (loss) versus beloningssignalen (reward) beïnvloedt de gradiëntbias en dynamiek.
3. Data-regime:
   • Off-policy: De student leert van een vast dataset (vaak gegenereerd door de leraar). Dit is efficiënt maar kan leiden tot een mismatch tussen trainings- en inferentiegedrag.
   • On-policy: De student genereert zijn eigen data (rollouts) om te leren. Dit lost de mismatch op, maar is computatief zeer duur en introduceert variabiliteit.

Bestaande methoden kiezen vaak voor één richting of één data-regime, wat leidt tot suboptimale prestaties, instabiliteit tijdens training of inefficiëntie.

---

2. Methodologie: Hybrid Policy Distillation (HPD)

De auteurs introduceren een unified view van distillatie als een hergewogen log-likelihood-doel op token-niveau. Op basis hiervan stellen ze Hybrid Policy Distillation (HPD) voor.

Kernconcepten van HPD:

HPD combineert de voordelen van Forward en Reverse KL, evenals off-policy en light-weight on-policy data, door divergenties te behandelen als token-level reward signalen.

1. Unificatie via Reweighting:
   In plaats van een vaste divergentie te minimaliseren, wordt het trainingsdoel herschreven als:
   $$L(\theta) = \min_{\theta} -E_{(s_t, a_t) \sim D_\pi} [w(a_t | s_t) \log q_\theta(a_t | s_t)]$$
   Waarbij $w(a_t | s_t)$ een wegingsfactor is die de discrepantie tussen student en leraar kwantificeert.

2. Hybride Forward-Reverse KL Mechanisme:
   HPD gebruikt een negatieve Reverse K1-schatting (een Monte Carlo schatter van de KL-divergentie) om de gap tussen leraar en student te meten.

   • Expert Token (Ground Truth/Leraar): Als de student de waarschijnlijkheid van het expert-token onderschat ($k_1 > 0$), wordt een Forward-KL-achtige weging toegepast om de kans te verhogen. Als de student het overschat, wordt de weging gemaskeerd om conflicterende gradiënten te voorkomen.
   • Student-ge-samplede Token: De student samplet ook een niet-expert token. Als dit token onredelijk is (de leraar geeft een lage kans, de student een hoge), wordt dit bestraft met een negatieve weging (Reverse KL gedrag). Dit zorgt ervoor dat waarschijnlijkheidsmassa wordt verplaatst weg van foutieve samples en terug naar de expert-token.

3. Hybride Data Regime:

   • HPD gebruikt voornamelijk off-policy data (offline dataset) voor efficiëntie.
   • Het introduceert echter een lichtgewicht on-policy component: de student samplet één token per stap tijdens de training om de verdeling te evalueren en te corrigeren. Dit is veel goedkoper dan volledige sequentie-rollouts (zoals in MiniLLM of GKD).

4. Het Algorithmische Proces:
   Voor elke stap in een offline traject:

   • Bereken de gap voor het expert-token.
   • Sample een token van de student.
   • Bereken de gap voor het gesamplede token.
   • Pas dynamische wegingen toe: versterk het expert-token als de student het onderschat, en onderdruk het gesamplede token als het onredelijk is.

---

3. Belangrijkste Bijdragen

1. Unificerend Kader: De auteurs presenteren een theoretisch kader dat bestaande methoden (SFT, FKLD, RKLD) verenigt onder één reweighted log-likelihood doel. Dit biedt inzicht in hoe divergentierichting, optimalisatie en data-regime samenwerken.
2. HPD Algoritme: Een nieuwe methode die Forward en Reverse KL hybride combineert via token-level reweighting, zonder extra hyperparameters. Het behoudt de efficiëntie van one-hot supervisie (zoals bij SFT) maar voegt de rijkdom van distributie-matching toe.
3. Efficiëntie en Stabiliteit: HPD vermijdt de hoge kosten van volledige on-policy rollouts terwijl het toch de train-inference mismatch aanpakt die bij puur off-policy training optreedt.

---

4. Resultaten

HPD is getest op diverse taken: wiskundig redeneren (langere generatie), dialoog (korte generatie) en coderen.

• Wiskundig Redeneren (Off-policy):

  • HPD overtreft consistent SFT, SeqKD, RKLD en JSD op benchmarks zoals AIME, AMC en MATH.
  • Opmerkelijk: HPD laat kleine modellen (bijv. 3B parameters) prestaties behalen die vergelijkbaar zijn met veel grotere modellen. Bijvoorbeeld, Qwen 2.5 3B verbeterde met 41% en LLaMA 3 3B met 77,9% ten opzichte van de baseline.
  • Stabiliteit: In tegenstelling tot SFT, dat vaak leidt tot "entropy collapse" (de student wordt te zeker en overfittet), behoudt HPD een stabiele entropie en convergeert soepel naar de leraar-distributie.

• Personalisatie & Dialoog (Off-policy):

  • Op benchmarks zoals MT-Bench en AlpacaEval behaalde HPD de hoogste scores, met name in het behouden van coherentie in multi-turn gesprekken.

• Coderen (Off-policy):

  • HPD toonde de meest stabiele en betrouwbare verbeteringen op HumanEval en MBPP, met minder variatie tussen verschillende modellen dan andere KD-methoden.

• On-policy Data & Initialisatie:

  • HPD fungeert als een sterke initialisatie voor verdere on-policy training (OPD). Wanneer HPD wordt gevolgd door OPD, worden de beste resultaten behaald, wat aantoont dat HPD een robuuste basis legt zonder de kosten van volledige on-policy training in de vroege fasen.

---

5. Significantie en Impact

• Efficiëntie: HPD biedt een manier om de prestaties van distillatie te maximaliseren zonder de rekenkosten van volledige on-policy rollouts (zoals in PPO of GKD) te hoeven betalen.
• Robuustheid: De methode lost het fundamentele compromis op tussen "mode coverage" (FKLD) en "mode seeking" (RKLD) door deze dynamisch te combineren op basis van de huidige staat van de student.
• Toepasbaarheid: De resultaten tonen aan dat HPD werkt over verschillende modelgroottes en domeinen (redeneren, chat, code), wat het een veelzijdige tool maakt voor de praktische implementatie van geoptimaliseerde LLM's.
• Toekomst: De auteurs suggereren dat HPD potentieel heeft om al in eerdere trainingsfasen (zoals mid-training) te worden toegepast om de basis van het model te versterken.

Kortom, HPD is een geavanceerde distillatiestrategie die de efficiëntie van SFT combineert met de expressiviteit van distributie-based distillatie, wat leidt tot stabielere training en superieure prestaties voor kleine taalmodellen.