Technische Samenvatting: ARKD – Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation

1. Probleemstelling

Knowledge Distillation (KD) is een primaire techniek voor het comprimeren van Large Language Models (LLM's) om de inferentiekosten te verlagen terwijl de prestaties behouden blijven. Traditionele KD-methoden die vertrouwen op een enkele Kullback-Leibler (KL) divergentie-objectief staan echter voor een fundamentele afruil in open-ended tekstgeneratie:

• Forward KL (FKL): Stimuleert "mode-covering" gedrag, waarbij de student probeert alle modi van de leraar-distributie te dekken. Dit leidt vaak tot het overschatten van waarschijnlijkheden in gebieden met lage betrouwbaarheid en een lange staart (long-tail), wat de generatiekwaliteit verslechtert.
• Reverse KL (RKL): Stimuleert "mode-seeking" gedrag, waarbij de focus ligt op dominante modi en gebieden met een lage waarschijnlijkheid worden genegeerd. Hoewel dit overschatting van de lange staart voorkomt, gaat dit ten koste van diversiteit en dekking, wat potentieel kan leiden tot mode collapse.

Bestaande statische interpolatiemethoden (bijv. een vaste 0.5:0.5 weging) slagen er niet in de dynamische aard van de student-leraar kloof tijdens de training te vangen. Een statische coëfficiënt kan niet zich aanpassen aan de verschuivende vereisten van distributie-alignement, die verschillende balansen van FKL en RKL vereisen tijdens verschillende trainingsfasen.

2. Methodologie: ARKD

De auteurs stellen ARKD voor (Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation), een framework dat dynamisch gewichten toekent aan FKL en RKL met behulp van een reinforcement learning (RL) agent.

Theoretische Fundering

Het paper stelt eerst de theoretische consistentie van FKL en RKL vast. Analyse toont aan dat beide objectieven convergeren naar dezelfde oplossing ($q_\theta = p$) wanneer de gradiënten nul zijn. Dit bevestigt dat beide divergenties valide zijn voor distributie-matching, maar complementaire gedragingen vertonen tijdens het optimalisatiepad, wat het gezamenlijk gebruik ervan rechtvaardigt.

RL-gebaseerd Framework

ARKD herformuleert het gewichtstoewijzingsprobleem als een sequentieel beslissingsproces (Markov Decision Process):

• Toestand ($s$): Een 6-dimensionale vector die de huidige trainingsdynamiek vastlegt, inclusief leraar/student entropieën ($H_T, H_S$), varianties ($\sigma^2_T, \sigma^2_S$) en huidige FKL/RKL loss-waarden ($L_{FKL}, L_{RKL}$).
• Actie ($\alpha$): Een policy netwerk ($\pi_\phi$) geeft een scalaire gewicht $\alpha \in (0, 1)$ uit, waarbij $\alpha$ de FKL-term weegt en $(1-\alpha)$ de RKL-term weegt.
• Beloning ($r$): Gedefinieerd als de negatieve distillatie loss ($-L_{distil}$) voor de huidige batch. Het doel is om de cumulatieve beloning te maximaliseren (de totale loss te minimaliseren) over het trainingspad.
• Optimalisatie:
  • Student Update: Minimaliseert de gewogen distillatie loss plus een taalmodellering loss.
  • Policy Update: Gebruikt het REINFORCE-algoritme met een Exponential Moving Average (EMA) baseline en entropie-regularisatie om de training te stabiliseren en boundary saturatie te voorkomen.

Trainingsdynamiek

Het policy netwerk leert adaptief de afruil te navigeren:

1. Exploratiefase: In het begin van de training verkent de policy een breder bereik van $\alpha$ (bijv. 0.3–0.6), met nadruk op FKL om brede dekking te garanderen en mode collapse te voorkomen.
2. Convergentiefase: Naarmate de student-distributie de leraar-distributie nadert, verschuift de policy naar lagere $\alpha$-waarden (RKL-dominant) om te focussen op precieze alignement en het voorkomen van overschatting van de lange staart.
3. Stabilisatiefase: De policy stabiliseert op een specifieke $\alpha$ (bijv. $\approx 0.18$) voor de uiteindelijke verfijning.

3. Belangrijkste Bijdragen

1. Theoretische Analyse: Het paper biedt een theoretische analyse die de complementariteit en convergentie-consistentie van FKL en RKL vaststelt, wat de basis vormt voor adaptieve distillatie.
2. ARKD Framework: Het introduceert een nieuw RL-gebaseerd gezamenlijk optimalisatieframework waarbij een policy netwerk dynamisch FKL en RKL weegt op basis van toestand-afhankelijke kenmerken, waarmee het verder gaat dan statische heuristieken.
3. Empirische Validatie: Uitgebreide experimenten tonen aan dat ARKD de generatiekwaliteit, distributie-fitting en generalisatie consistent verbetert over meerdere benchmarks en model-schalen.

4. Experimentele Resultaten

De auteurs evalueerden ARKD op instructie-tuning taken met GPT-2 (1.5B $\to$ 120M/340M) en LLaMA (13B $\to$ 7B) configuraties.

• In-Distribution (ID) Prestaties: Op de DollyEval benchmark behaalde ARKD de hoogste BertScore en competitieve of superieure Rouge-L scores vergeleken met baselines (SFT, SeqKD, FKL, RKL, FKL+RKL, en MiniLLM). Bijvoorbeeld, op GPT-2 1.5B $\to$ 120M behaalde ARKD een Rouge-L van 24.5 en een BertScore van 0.815, waarmee het de sterke MiniLLM baseline overtrof.
• Out-of-Distribution (OOD) Generalisatie: ARKD demonstreerde superieure generalisatie op onbekende datasets (SelfInst, VicunaEval, S-NI, UnNI), waarbij het consequent alle baselines overtrof binnen beide modelfamilies.
• Ablatie-studies:
  • vs. Statische/Lineaire Scheduling: ARKD presteerde significant beter dan statische weging (vaste $\alpha$) en lineaire scheduling strategieën (tijd-gebaseerde decay/toename), die er niet in slaagden complexe alignement-dynamiek te vangen.
  • vs. Greedy Heuristieken: ARKD overtrof greedy heuristieken (het selecteren van de divergentie met de laagste loss op elk stap) met 0.4–0.6 punten. De greedy benadering bleek myoop en degradeerde vaak naar pure RKL, terwijl ARKD een voordelige kleine FKL-component behield door continue mixing.
  • Trajectanalyse: Het geleerde $\alpha$-traject bevestigde een intelligente strategie: vroege nadruk op FKL voor dekking, overgaand naar RKL voor precisie, een patroon dat niet haalbaar is door handmatig ontwerp.

5. Betekenis en Claims

Het paper claimt dat ARKD de staat van de techniek in adaptieve kennisdistillatie vooruit helpt door de beperkingen van statische KL-gebaseerde methoden aan te pakken. Door gebruik te maken van reinforcement learning om de balans tussen mode-covering en mode-seeking gedrag dynamisch te beheren, bereikt ARKD een betere distributie-alignement die zowel de generatiekwaliteit als de generalisatie verbetert.

De auteurs positioneren ARKD als een principieelere aanpak die de noodzaak van dynamische, toestand-bewuste weging boven vaste heuristieken valideert. Zij merken op dat hoewel de huidige policy steunt op een handmatig geconstrueerde toestand-vector en directe batch-beloningen, het framework compatibel is met bredere $f$-divergence families en on-policy paradigma's. Het werk legt een solide fundament voor efficiënte LLM-compressie door aan te tonen dat adaptieve, RL-gestuurde optimalisatie consistent beter presteert dan statische en greedy baselines over diverse taken en model-schalen.