Technische Samenvatting: Niet Alle Meningsverschillen Zijn Leerbaar: Token Leerbaarheid in On-Policy Distillatie

Probleemstelling

On-policy distillatie (OPD) traint een student-taalmodel op zijn eigen rollouts met behulp van token-level supervisie van een teacher-policy. Hoewel recente selectieve OPD-methoden proberen de efficiëntie te verbeteren door prioriteit te geven aan "saliente" tokens—zoals die met hoge entropie of hoge meningsverschillen tussen teacher en student (KL-divergentie)—vertrouwen deze benaderingen op de aanname dat hoge meningsverschillen gelijkstaan aan hoge leerwaarde.

Het paper identificeert een kritieke fout in deze aanname: ruwe KL-meningsverschillen zijn een grove proxy voor leren. Het vermengt twee kwalitatief verschillende scenario's:

1. Leerbaar Meningsverschil: De teacher wijst corrigerende waarschijnlijkheidsmassa toe aan tokens die al binnen de huidige top-$K$ voorspellende steun van de student vallen.
2. Incompatibel Meningsverschil: De teacher wijst massa toe aan tokens die ver buiten de huidige steun van de student vallen. Hoewel dit een grote KL-kloof creëert, kan de student deze correctie niet gemakkelijk absorberen, wat leidt tot zwakke of negatieve lokale leereffecten.

Het kernprobleem is het onderscheid maken tussen welke token-level teacher-signalen daadwerkelijk leerbaar zijn (d.w.z. in staat zijn lokale verbetering te induceren) versus die welke slechts salient of incompatibel zijn.

Methodologie

1. Diagnostic met Vaste Context

Om de leerwaarde van specifieke tokens te isoleren van rollout-ruis en verschuivingen in downstream-context, introduceren de auteurs een diagnostic met vaste context.

• Protocol: Door de student gegenereerde voorvoegsels worden verzameld en bevroren als een contextbank. Zowel het initiële (voor-update) als het getrainde studentenmodel worden opnieuw gescoord tegen dezelfde teacher-distributie op deze vaste contexten.
• Maatstaf: De Fixed-Context Token Gain ($G^{\text{fix}}_{i,t}$) wordt gedefinieerd als de reductie in KL-divergentie tussen de teacher en de student na training, gemeten op hetzelfde voorvoegsel.
• Doel: Dit stelt de auteurs in staat om pre-update signalen (zoals ruwe KL) te correleren met daadwerkelijke lokale verbetering, terwijl steekproefvariatie wordt gecontroleerd.

2. Formalisering van Token Leerbaarheid

Met behulp van de diagnostic decomponeren de auteurs meningsverschillen in twee componenten op basis van lokale steuncompatibiliteit:

• Lokale Steun ($S^S_t(K)$): De set van de top-$K$ tokens van de student.
• Compatibiliteitsmassa ($C_t$): De totale waarschijnlijkheidsmassa die de teacher toewijst aan de lokale steun van de student $S^S_t(K)$.
• Decompositie:
  • Leerbaar Meningsverschil ($D^L_t$): Hoog meningsverschil waarbij de correctie van de teacher is afgestemd op de steun van de student (Hoge $C_t$).
  • Incompatibel Meningsverschil ($D^I_t$): Hoog meningsverschil waarbij de correctie van de teacher buiten de steun van de student valt (Lage $C_t$).

De Leerbaarheidsscore wordt gedefinieerd als het product van genormaliseerd meningsverschil en compatibiliteit:
$$s^{\text{teach}}_t = \tilde{D}_t \cdot \tilde{C}_t$$
waarbij $\tilde{D}_t$ het genormaliseerde meningsverschil is en $\tilde{C}_t$ de genormaliseerde compatibiliteitsmassa.

3. Leerbaarheidsbewuste OPD (TA-OPD)

Gebaseerd op de bevinding dat $D^L_t$ de fixed-context gain beter voorspelt dan ruwe KL, stellen de auteurs TA-OPD voor, een lichtgewicht methode voor selectie van token-posities:

• Selectie: Voor een bepaald budget (bijv. 5% of 10% van de tokens) selecteert TA-OPD de top-$N$ posities met de hoogste leerbaarheidsscores ($s^{\text{teach}}_t$).
• Implementatie: Het vereist alleen token-probabiliteiten van teacher en student (specifiek top-$K$ log-probs) die beschikbaar zijn tijdens standaard OPD. Het vereist geen reward-modellen, verifiers of extra labels.
• Contrast: In tegenstelling tot eerdere methoden die selecteren op basis van entropie of ruwe divergentie (TIP), filtert TA-OPD expliciet op "steun-gealigneerde" correcties.

Belangrijkste Bijdragen

1. Diagnostic met Vaste Context: Een methode om KL-reductie in dezelfde context te meten, waarbij pre-update token-signalen worden gekoppeld aan lokale OPD-verbetering zonder verstorende rollout-ruis.
2. Formalisering van Token Leerbaarheid: Het aantonen dat ruwe KL-meningsverschillen leerbare en incompatibele signalen vermengt. Het paper bewijst dat steun-gealigneerd meningsverschil de betrouwbare predictor is van nuttig leren.
3. TA-OPD Algoritme: Een praktische, lichtgewicht selectiemethode die OPD-verlies alleen toepast op posities met hoge leerbaarheid, waardoor de efficiëntie wordt verbeterd zonder externe tools.
4. Empirische Validatie: Uitgebreide experimenten tonen aan dat selectie op basis van leerbaarheid betere downstream-prestaties oplevert dan selectie op basis van salientie, zelfs met aanzienlijk verminderde token-budgetten.

Experimentele Resultaten

De auteurs hebben TA-OPD geëvalueerd in vier teacher-student-configuraties met Qwen3 (4B, 8B, 14B) en Qwen2.5 (3B) modellen, met behulp van benchmarks zoals AIME, GPQA-Diamond, HumanEval en MATH-500.

• Prestaties vs. Volledige OPD: In veel situaties kwam TA-OPD overeen met of overtrof volledige-token OPD, terwijl het slechts 5% van de tokens gebruikte. Bijvoorbeeld, in de Qwen3-4B $\to$ 1.7B-configuratie behaalde TA-OPD een gemiddelde score van 44,89 tegenover 42,37 voor Volledige OPD, ondanks het gebruik van een 10%-budget (en het overtreffen van volledige OPD in specifieke metrieken).
• Vergelijking met Baselines: TA-OPD presteerde consequent beter dan entropie-gebaseerde, divergentie-gebaseerde en TIP (Token Importance in OPD) baselines.
  • In de cross-backbone-configuratie (DeepSeek-R1-Distill-Qwen-14B $\to$ Qwen2.5-3B) verbeterde TA-OPD de gemiddelde score naar 30,62, wat significant beter was dan het Basemodel (29,98) en Volledige OPD (28,76).
• Budgetgevoeligheid: De resultaten geven aan dat OPD-prestaties token-kwaliteit-ge limiteerd zijn in plaats van puur budget-ge limiteerd. Het verhogen van het token-budget boven 10% leverde vaak geen verdere winst op en degradeerde soms de prestaties, wat bevestigt dat niet alle extra supervisie gunstig is.
• Robuustheid: De positieve "high-support gain gap" (het prestatieverschil tussen tokens met hoge en lage leerbaarheid) bleef consistent over verschillende contexten (wiskunde, GSM8K-COT), teacherschalen (8B, 14B) en steungroottes ($K=8, 16, 32$).

Betekenis en Aanspraken

Het paper herformuleert selectieve on-policy distillatie niet als een probleem van het selecteren van "saliente" of "onzekere" tokens, maar als een probleem van het identificeren van lokaal leerbare teacher-signalen.

• Kernclaim: De winst uit OPD hangt af van de compatibiliteit tussen de correctie van de teacher en de huidige voorspellende steun van de student. Hoge KL-divergentie alleen is onvoldoende; de teacher moet massa toewijzen aan tokens die de student al als plausibel beschouwt.
• Praktische Impact: TA-OPD biedt een manier om de rekenkosten van distillatie drastisch te verminderen (door incompatibele tokens te filteren) terwijl de modelprestaties worden behouden of verbeterd. Dit wordt bereikt zonder de overhead van reward-modellen of verifiers, waardoor het toepasbaar is op standaard trainingspijplijnen.
• Beperkingen: De auteurs merken op dat hun analyse zich richt op zwaar wiskundig redeneren en Qwen-familie modellen. Hoewel de diagnostic model-agnostisch is, is bredere validatie op dialoog, code en niet-Qwen backbones nodig om generaliseerbaarheid te bevestigen. Bovendien verwijst het "token-budget" naar supervisie-dichtheid, niet noodzakelijk naar muurkloksnelheid, aangezien de methode posities selecteert in plaats van transformer-berekening te snoeien.

Concluderend suggereert het werk dat effectieve distillatie een compatibiliteitsbewuste supervisiestrategie vereist, waarbij de student het beste leert van teacher-signalen die zowel corrigerend als lokaal bereikbaar zijn.