$\mathcal{X}$-KD: General Experiential Knowledge Distillation for Large Language Models

Probleemstelling

Bestaande methoden voor Kennisdistillatie (Knowledge Distillation - KD) bij Large Language Models (LLMs) richten zich voornamelijk op het imiteren van het gedrag van het lerende model (de "teacher"). Hierbij wordt de student getraind om de output-distributie van de teacher na te bootsen. Een fundamenteel tekortkoming van deze aanpak is dat ze vaak het oorspronkelijke leeromgeving negeren waarin de teacher zijn kennis heeft verworven.

Volgens de Experiential Learning Theory en onderzoek in Inverse Reinforcement Learning (IRL) presteert een imitator beter als hij niet alleen het gedrag nabootst, maar ook de onderliggende beloningsfunctie (reward function) van de omgeving leert begrijpen. Bestaande methoden zoals Behavioral Cloning (BC) missen dit inzicht. Hoewel recente werken zoals MiniLLM en GKD stappen zetten in de richting van IRL, hebben ze beperkingen:

• Ze gebruiken vaak beleidsgradienten (policy gradients) die afwijken van het standaard toezicht-gebaseerde (supervised learning) trainingsparadigma.
• Ze zijn complex om te integreren in bestaande frameworks.
• Ze zijn niet flexibel genoeg om verschillende divergentie-maatstaven of sequence-level distillatie voor black-box modellen te ondersteunen.

Methodologie: X-KD

De auteurs stellen X-KD (Experiential Knowledge Distillation) voor, een algemeen raamwerk dat studenten toelaat om te leren in de oorspronkelijke leeromgeving van de teacher. De kern van de methode is het gezamenlijk modelleren van de teacher's oorspronkelijke beloningsfunctie en het uitvoeren van beleidsdistillatie.

Technische Kernen:

1. Bayesische Inverse Reinforcement Learning (BIRL): X-KD baseert zich op het AVRIL (Approximated Variational Reward Imitation Learning) framework. In plaats van alleen de teacher's output te kopiëren, schat X-KD de posterior-verdeling van de beloningsfunctie $p(R|\pi)$ die de teacher heeft gebruikt.
2. Gecombineerde Doelfunctie: De trainingsdoelstelling bestaat uit twee delen:
   • Een distillatie-component (bijv. sequence-level of divergence-based) die de student leert de teacher na te bootsen.
   • Een experientiële regularisatie-component die de student dwingt consistent te zijn met de geschatte beloningsfunctie van de teacher.
3. Formulering:
   • Het probleem wordt gemodelleerd als een Markov Decision Process (MDP) voor Natural Language Generation (NLG).
   • De student leert een Q-waarde decoder ($Q_\theta$) die de beloning schat.
   • De totale loss-functie combineert de standaard KD-loss met een term die de TD-error (Temporal Difference error) minimaliseert onder de geschatte beloningsverdeling. Dit zorgt ervoor dat het gedrag van de student indirect wordt geconditioneerd op de beloning van de oorspronkelijke omgeving.
4. Flexibiliteit: X-KD is ontworpen om compatibel te zijn met:
   • Sequence-level KD: Distillatie op het niveau van volledige zinnen.
   • Supervised KD: Token-level distillatie (forward-KL).
   • Generalized KD (GKD): Gebruik van on-policy data en diverse divergentie-maatstaven (zoals JS-divergentie).
   • Black-box en White-box: Het werkt zowel wanneer de teacher's interne distributie bekend is (white-box) als wanneer alleen de output beschikbaar is (black-box).

Belangrijkste Bijdragen

1. Nieuw Raamwerk: X-KD is een innovatief KD-framework dat Bayesische IRL integreert om studenten te laten leren in de teacher's oorspronkelijke leeromgeving, in plaats van alleen gedrag te imiteren.
2. Eenvoud en Flexibiliteit: Het framework behoudt het standaard supervised learning paradigm (geen complexe policy gradients nodig) en kan naadloos worden geïntegreerd met bestaande KD-methoden (SeqKD, SKD, GKD) zonder de fundamentele trainingsstructuur te veranderen.
3. Empirisch Bewijs: Uitgebreide experimenten tonen aan dat X-KD superieur presteert op meerdere Natural Language Generation (NLG) taken.

Resultaten

De auteurs hebben X-KD getest op drie taken: abstractive summarization (XSum), machine translation (WMT14), en arithmetic reasoning (GSM8k).

• Prestatie: X-KD overtreft consistent de baselines (zoals GKD, MiniLLM, en standaard KD) in zowel "large-to-small" als "large-to-base" modelconfiguraties.
• Trade-off tussen Prestatie en Diversiteit: X-KD bereikt een betere balans tussen outputkwaliteit en diversiteit. Terwijl de prestaties van andere methoden sterk dalen bij het verhogen van de diversiteit (sampling temperature), behoudt X-KD zijn kwaliteit beter. Dit suggereert dat de student meer generaliseerbare strategieën leert.
• Data-efficiëntie: X-KD bereikt vergelijkbare prestaties als baselines die op 100% van de data zijn getraind, maar met slechts 50% tot 75% van de trainingsdata. Dit toont aan dat het leren van de onderliggende beloningsfunctie efficiënter is dan het bloot imiteren van data.
• Black-box Distillatie: In experimenten met Llama-2-7b als black-box teacher, presteerde SeqXKD beter dan standaard SeqKD, beoordeeld door GPT-4.

Betekenis en Conclusie

X-KD markeert een verschuiving in kennisdistillatie van louter "gedragsimitatie" naar "experientiële learning". Door de student te plaatsen in de context van de teacher's oorspronkelijke beloningsstructuur, leert het model fundamentele principes van de taak in plaats van oppervlakkige patronen.

De belangrijkste implicaties zijn:

• Efficiëntie: Minder trainingsdata is nodig voor hoogwaardige resultaten.
• Robuustheid: Betere prestaties bij het genereren van diverse outputs.
• Toepasbaarheid: Het framework is eenvoudig te implementeren in bestaande LLM-pipelines omdat het het supervised learning paradigma behoudt.

De auteurs erkennen wel dat de prestaties gevoelig zijn voor de instelling van de "experiential weight" (hyperparameter) en dat toekomstig werk zich meer moet richten op systematische studies in black-box scenario's.