Technische Samenvatting: Koppeling van KL en Trajecten Ontkoppelen in LLM-distillatie

Probleemstelling

Kennisdistillatie is een hoeksteen geworden van het post-training van Large Language Models (LLM), met name voor redeneertaken, en werkt vaak samen met Versterkend Leren (RL). Het ontwerpruimte van distillatie blijft echter slecht begrepen. Huidige pijplijnen volgen doorgaans een van twee paradigma's: off-policy distillatie (training op door de leraar gegenereerde traces, vaak via Supervised Fine-Tuning) of on-policy distillatie (OPD) (training op door de student gegenereerde rollouts met feedback van de leraar).

Het artikel betoogt dat deze heersende paradigma's impliciet twee orthogonale ontwerpkiezen koppelen:

1. Prefixbron: Of supervisie wordt toegepast op prefixes gegenereerd door de leraar (off-policy) of door de student (on-policy).
2. Token-niveau KL-richting: Of het doel de forward KL ($KL(p_{teacher} \parallel q_{student})$) minimaliseert of de reverse KL ($KL(q_{student} \parallel p_{teacher})$).

De conventionele praktijk koppelt teacher-prefixes aan forward KL en student-prefixes aan reverse KL. De auteurs betogen dat deze koppeling niet inherent is aan het token-niveau en dat het ontkoppelen van deze assen een bredere ontwerpruimte onthult met vier geldige doelen, die elk corresponderen met verschillende leerregimes.

Methodologie

Theoretisch Kader: De Assen Ontkoppelen

De auteurs ontleden de KL-divergentie op sequentieniveau met behulp van de autoregressieve kettingregel. Zij tonen aan dat:

• Forward KL ($KL(p_T \parallel q_S)$) van nature gepaard gaat met teacher-prefixes (off-policy), omdat de verwachting wordt genomen over de trajectieverdeling van de leraar.
• Reverse KL ($KL(q_S \parallel p_T)$) van nature gepaard gaat met student-prefixes (on-policy), omdat de verwachting wordt genomen over de trajectieverdeling van de student.

Door deze als onafhankelijke assen te behandelen, definiëren de auteurs vier doelen:

1. Teacher Prefix + Forward KL: Correspondeert met off-policy SFT (standaard distillatie).
2. Student Prefix + Forward KL: Correspondeert met DAgger-stijl on-policy SFT.
3. Teacher Prefix + Reverse KL: Correspondeert met een offline RL-stijl distillatiedoel.
4. Student Prefix + Reverse KL: Correspondeert met On-Policy Distillatie (OPD), gezien als dense-reward on-policy RL.

Het artikel vestigt identiteiten op gradientenniveau die aantonen dat:

• Forward KL een SFT-stijl cross-entropy-gradiënt oplevert met teacher-soft targets.
• Reverse KL een RL-stijl policy gradient oplevert met een dense beloning gedefinieerd door de log-ratio tussen leraar en student ($r = \log p_T - \log q_S$).

Experimentele Opstelling

De auteurs voeren een gecontroleerde studie uit op wiskundig redeneren met Qwen3-modellen (0.6B student, 4B/8B leraren). Zij evalueren de vier doelen in twee regimes:

1. Standalone Distillatie: Vergelijking van de uiteindelijke prestaties van modellen die uitsluitend met elk doel zijn getraind.
2. Distillatie-volgende-RL: Het gebruik van de gedistilleerde checkpoints als initialisatie voor Group Relative Policy Optimization (GRPO).

Zij volgen nauwkeurigheid, voorspellende entropie en responslengte over verschillende sequentielengtes (128 vs. 4096 tokens) en benchmarks (MATH500, GSM8K, AIME24, AMC23).

Belangrijkste Bijdragen

1. Identificatie van Drie Trade-offs

De empirische analyse onthult drie fundamentele trade-offs die het distillatieontwerp beheersen:

• Nauwkeurigheid–Entropie Trade-off (KL-richting): Reverse KL bereikt consistent hogere standalone-nauwkeurigheid (Avg@k), maar drijft de voorspellende entropie naar instorting, wat de diversiteit vermindert (lagere Pass@k) en het model minder stabiel maakt voor daaropvolgende RL. Forward KL behoudt entropie, wat een stabielere RL-verbetering ondersteunt, zij het met lagere initiële standalone-nauwkeurigheid.
• Kwaliteit–Berekeningskosten Trade-off (Prefixbron): Student-prefixes (on-policy) leveren hogere nauwkeurigheid op bij gelijkgestelde trainingsstappen door staten te superviseren die de student daadwerkelijk bezoekt. Echter, teacher-prefixes (off-policy) zijn rekenkundig efficiënter bij gelijkgestelde FLOPs, omdat ze het hergebruik van gecachte teacher-logits en offline trajecten mogelijk maken, waardoor de kosten van online studentgeneratie worden vermeden.
• Nauwkeurigheid–Stabiliteit Trade-off (Trainingslengte): Lange-sequentie distillatie (4096 tokens) verbetert de redeneernauwkeurigheid, maar veroorzaakt, met name onder reverse KL, ernstige entropie-instorting en "lengte-inflatie" (antwoorden die de generatielimiet bereiken). Korte-sequentie training is stabieler maar minder nauwkeurig.

2. Voorgestelde Methoden om Trade-offs in Evenwicht te Brengen

Gedreven door deze bevindingen, stellen de auteurs twee praktische methoden voor:

• KL-menging: Een gewogen combinatie van forward en reverse KL bij elke prefix. De auteurs vinden dat bij lange-sequentie distillatie een forward-dominante mengeling cruciaal is. Dit behoudt de nauwkeurigheidsvoordelen van reverse KL, terwijl het entropie-instorting en lengte-inflatie voorkomt en de trainingsdynamiek stabiliseert.
• Entropie-gedreven Lengtecurriculum: Een trainingschema dat begint met korte horizonnen en de sequentielengte alleen verhoogt zolang de voorspellende entropie boven een stabiliteitsdrempel blijft. Deze aanpak voorkomt de entropie-instorting die geassocieerd wordt met vaste lange-horizon training.

Resultaten

Standalone Distillatie

• Reverse KL presteert over het algemeen beter dan Forward KL in standalone nauwkeurigheid (Avg@k), met gemiddeld +2,45 punten winst.
• Student-prefixes presteren beter dan Teacher-prefixes qua nauwkeurigheid, vooral in lange-sequentie instellingen (+3,55 Avg@k winst).
• Entropie-instorting: Pure reverse KL in lange-sequentie instellingen drijft de entropie dicht bij nul en zorgt ervoor dat responslengtes infleren tot het maximale limiet (bijv. 8192 tokens), wat Pass@k verslechtert.

Distillatie-volgende-RL (GRPO)

• Initialisatiestabiliteit: Modellen geïnitieerd met Forward KL (zelfs als ze lager scoren in standalone nauwkeurigheid) tonen een betrouwbaardere verbetering tijdens GRPO omdat ze hogere entropie behouden.
• Reverse KL-valkuil: Modellen geïnitieerd met Reverse KL beginnen vaak met hogere nauwkeurigheid, maar degraderen of plateau tijdens RL vanwege lage entropie en verminderde exploratiecapaciteit.
• Curriculum Succes: Het entropie-gedreven lengtecurriculum verbetert de prestaties aanzienlijk. In vergelijking met vaste 4096-token training verbeterde het Avg@k met 3,6 punten en Pass@k met maximaal 5,8 punten, terwijl de gemiddelde responslengte met ongeveer 3× werd verlaagd (bijv. van ~8185 naar ~2688 tokens).

KL-menging Resultaten

In lange-sequentie instellingen behield een mengeling met een hoge forward KL-weigh (bijv. 0,8 Forward / 0,2 Reverse) het grootste deel van de nauwkeurigheid van reverse KL, terwijl het entropie stabiliseerde en lengte-inflatie voorkwam.

Betekenis en Aanspraken

Het artikel claimt een unificerend theoretisch kader te bieden dat uiteenlopende trainingsrecepten (SFT, DAgger, Offline RL, OPD) verbindt als specifieke instantiaties van twee orthogonale ontwerpkiezen. Door prefixbron en KL-richting te ontkoppelen, betogen de auteurs dat practitioners voorbij rigide paradigma's kunnen gaan om redeneerdistillatiedoelen te ontwerpen die expliciet in evenwicht brengen:

• Nauwkeurigheid (via reverse KL en lange sequenties),
• Diversiteit/Exploratie (via forward KL en entropiebehoud),
• Berekeningskosten-efficiëntie (via teacher-prefixes/caching), en
• Trainbaarheid voor downstream RL (via stabiele entropie).

De voorgestelde methoden (KL-menging en entropie-gedreven curriculum) worden gepresenteerd als praktische hulpmiddelen om de geïdentificeerde trade-offs te navigeren, waardoor het trainen van kleinere redeneermodellen mogelijk is die zowel nauwkeurig als stabiel zijn voor verdere policy-optimalisatie. De auteurs benadrukken dat het "beste" standalone doel niet per se de beste initialisatie is voor een RL-pijplijn, wat de noodzaak onderstreept van een holistisch beeld van het post-training proces.