Technische Samenvatting: Over de Positiebias van On-Policy Distillatie

Probleemstelling

On-Policy Distillatie (OPD) verbetert de leer-efficiëntie van Reinforcement Learning (RL) voor Large Language Models (LLM's) door dichte, op token-niveau gebaseerde supervisie te bieden van een teacher-model aan een student-model dat zijn eigen rollouts genereert. De standaard OPD-doelstellingen middelen echter de token-level verliezen uniform (meestal via KL-divergentie), wat impliciet uitgaat van een gelijke waarde voor alle tokens in een sequentie.

De auteurs identificeren een kritieke Positiebias in deze aanpak: naarmate de student-rollouts langer worden, wijkt de gegenereerde trajectorie steeds meer af van de distributie van de teacher. Bij gevolg is de kwaliteit van de teacher-supervisie aan latere posities minder. Empirische observaties tonen aan dat:

1. Asymmetrie in de waarde van supervisie: Training met alleen de eerste 30% van de tokens (prefix) levert een prestatie op die vergelijkbaar is met of zelfs beter is dan het gebruik van alle tokens, terwijl training op de laatste 30% (suffix) verwaarloosbare leervoordelen biedt.
2. Aanhoudende divergentie: Zelfs na convergentie blijft de KL-divergentie tussen de teacher en de student significant, wat suggereert dat OPD de student alleen optimaliseert binnen een klein lokaal gebied van de policy-ruimte.
3. Inefficiëntie: Het uniform wegen van alle tokens verspilt het beperkte gradiënt-updatebudget aan latere tokens waar de student al is afgedreven naar Out-of-Distribution (OOD) regio's ten opzichte van de teacher, waardoor het supervisiesignaal ruisachtig of onbruikbaar wordt.

Methodologie: Importance-Weighted On-Policy Distillation (IW-OPD)

Om deze bias aan te pakken, stellen de auteurs Importance-Weighted On-Policy Distillation (IW-OPD) voor, een methode die token-gewichten dynamisch aanpast op basis van de geaccumuleerde discrepantie tussen de student- en teacher-distributies.

Theoretische Fundering

De auteurs kaderen OPD als een geconstreind optimalisatieprobleem (een lokale projectie) waarbij de student-policy wordt bijgewerkt binnen een trust region die wordt gedefinieerd door een KL-divergentie budget $\rho$.

• Optimale Policy Afleiding: Het oplossen van de geconstreinde minimalisatie van $D_{KL}(q \| \pi_T)$ onder de voorwaarde $D_{KL}(q \| \pi_\theta) \leq \rho$ levert een gesloten vorm van de optimale policy $q^\star$ op.
• Likelihood Ratio Reweighting: De oplossing $q^\star$ herweegt de student-policy $\pi_\theta$ door de likelihood ratio $(\pi_T / \pi_\theta)^\alpha$. Dit impliceert dat tokens waar de waarschijnlijkheid van de teacher aanzienlijk hoger is dan die van de student (hoge compatibiliteit) een hogere weging moeten krijgen, terwijl tokens waar de student is afgedreven (lage ratio) een lagere weging moeten krijgen.

Praktische Implementatie

Het direct gebruiken van de ruwe likelihood ratio is numeriek instabiel vanwege de exponentiële accumulatie van waarschijnlijkheden over lange sequenties. IW-OPD introduceert een stabiele, ongesigneerde cumulatieve prefix-discrepantie als een proxy voor het belang-gewicht:

1. Geaccumuleerde Discrepantie: In plaats van de ruwe ratio, somt de methode de absolute token-level advantages ($|A^{OPD}_k|$) op vanaf het begin van de sequentie tot aan de huidige token $t$. Dit vangt de totale "drift" op van het geprefereerde pad van de teacher.
2. Normalisatie: De geaccumuleerde discrepantie wordt binnen elke sample gestandaardiseerd naar een bereik van $[0, 1]$, wat zorgt voor een stabiele schaling.
3. Interpolatie: Om de volledige verwijdering van supervisie voor afgedreven tokens te voorkomen, is het uiteindelijke gewicht een interpolatie tussen het standaard OPD-signaal (gewicht 1) en het belang-gewicht afgeleid van de discrepantie.
   • Formule: $\tilde{r}^{IW-OPD}_t = 1 + \gamma \cdot (1 - \frac{\sum_{k<t} |A^{OPD}_k|}{\sum_{k<T} |A^{OPD}_k|})$.
   • Hierbij controleert $\gamma$ de bijdrage van de importance weighting.

Deze aanpak geeft van nature een hogere weging aan vroege tokens (waar de prefix compatibel is met de teacher) en een lagere weging aan latere tokens (waar de geaccumuleerde drift groot is), waardoor het gradiëntbudget effectief wordt geheralloceerd naar de meest informatieve delen van de trajectorie.

Belangrijkste Bijdragen

1. Identificatie van Positiebias: De paper toont empirisch aan en verklaart theoretisch dat OPD lijdt aan een positiebias waarbij vroege tokens aanzienlijk waardevollere supervisie bieden dan latere tokens vanwege de distributieve drift.
2. Perspectief van Geconstreinde Optimalisatie: De auteurs bieden een fundamenteel begrip van OPD als een geconstreerde lokale projectie, waarbij zij afleiden dat de optimale update-strategie inherent een ratio-gebaseerde herweging vereist.
3. IW-OPD Algoritme: Zij stellen IW-OPD voor, een praktisch en stabiel algoritme dat de belangrijkheid van tokens schat met behulp van een ongesigneerde cumulatieve prefix-discrepantie. Het vereist geen extra teacher-evaluaties buiten de standaard OPD.
4. Empirische Validatie: Uitgebreide experimenten laten zien dat IW-OPD sneller convergeert en een hogere uiteindelijke prestatie bereikt dan standaard OPD over diverse model-schalen en teacher-student configuraties heen.

Experimentele Resultaten

De auteurs evalueerden IW-OPD met behulp van Qwen3-modellen (0.6B, 1.7B, 4B, 30B, en 235B) als teachers en students op benchmarks voor wiskundige redenering (AIME-2024, AIME-2025, HMMT-2025) en code-generatie (HumanEval+, MBPP+).

• Prestatiewinst: IW-OPD presteert consistent beter dan standaard OPD.
  • In de Qwen3-1.7B student setting met een Qwen3-30B teacher, verbeterde IW-OPD de AIME-2025 nauwkeurigheid met +6,9 punten bij stap 10 en met +1,7 punten bij convergentie.
  • Over alle student-schalen (0.6B tot 30B) en teacher-settings heen, verbeterde IW-OPD de gemiddelde score.
• Sample Efficiëntie: IW-OPD convergeert aanzienlijk sneller. In verschillende gevallen evenaarde of overtrof de IW-OPD checkpoint bij stap 10 de uiteindelijke prestatie van standaard OPD.
• Schaaleffecten:
  • Sterkere Teachers: IW-OPD maakt sterkere teachers meer sample-efficiënt. Terwijl standaard OPD aanvankelijk moeite heeft met grotere teacher-student gaten (meer stappen nodig om te profiteren van de sterkere teacher), benut IW-OPD het signaal van de sterkere teacher onmiddellijk.
  • Kleinere Students: De relatieve winst van IW-OPD neemt toe naarmate het student-model kleiner wordt (bijv. +9,5% relatieve winst voor de 0.6B student versus +1,8% voor de 4B student), wat aangeeft dat het bijzonder effectief is in cross-scale distillatie waar de drift heviger is.

Betekenis en Claims

Het artikel claimt dat effectieve on-policy supervisie rekening moet houden met de trajectorie-context, en niet alleen met de mismatch op token-niveau. Door te erkennen dat de teacher-supervisie niet uniform betrouwbaar is langs een door de student gegenereerde trajectorie, heralloceert IW-OPD het eindige gradiëntbudget naar "teacher-compatibele" prefixes.

De auteurs concluderen dat deze eenvoudige modificatie — het lager wegen van tokens waar de student is afgedreven van de teacher-distributie — zowel de sample-efficiëntie als de uiteindelijke prestatie van on-policy distillatie aanzienlijk verbetert. Dit inzicht suggereert dat toekomstige on-policy methoden de cumulatieve historie van de rollout (de prefix) moeten overwegen bij het toekennen van credit of supervisie-gewichten, in plaats van alle tokens in een traject als onafhankelijk en even waardevol te behandelen.