Technische Samenvatting: On-Policy Distillatie met Best-of-N Teacher Rollout Selectie (BRTS)

1. Probleemstelling

On-policy distillatie (OPD) is een data-efficiënte post-training methode voor Large Language Models (LLMs) naar voren gekomen, waarbij een studentmodel wordt gesuperviseerd op zijn eigen gesamplede trajecten met behulp van de per-token log-kansen van de leraar. Hoewel OPD blootstellingsbias vermindert door supervisie af te stemmen op de staten die de student daadwerkelijk bezoekt, lijden standaardimplementaties aan supervisiemsignalen met hoge variantie.

De geïdentificeerde kernbeperking is tweeledig:

1. Ruizige Studentcontexten: Supervisie wordt berekend op voorvoegsels gegenereerd door een imperfecte student, die kunnen afdwalen naar ruizige redeneerstaten waar lokale lerarenfeedback minder informatief is.
2. Stochastische Teacher Rollouts: Standaard OPD vertrouwt doorgaans op één enkele teacher rollout per prompt. Omdat de leraar stochastisch is, kan een enkele gesamplede trajecton onjuist, oninformatief of slecht afgestemd zijn op het huidige redeneergedrag van de student. Dit resulteert in onbetrouwbare begeleiding, vooral bij moeilijke prompts waar de gesamplede outputs van de leraar aanzienlijk variëren in juistheid en stijl.

Bijgevolg loopt standaard OPD het risico foutieve leraren signalen te propageren of de student niet bloot te stellen aan volledige, correcte redeneerpaden, wat zijn effectiviteit op uitdagende redeneerbenchmarks beperkt.

2. Methodologie: BRTS Framework

De auteurs stellen BRTS (Best-of-N Rollout Teacher Selection) voor, een framework ontworpen om betrouwbare teacher trajecten te selecteren voordat deze worden gebruikt voor supervisie binnen de OPD-lus. BRTS breidt standaard student-context distillatie uit met een teacher-context supervisie tak die is opgebouwd uit een geselecteerde, hoogwaardige teacher traject.

2.1 Juistheids- en Afstemmingsbewuste Traject Selectie

Voor elke prompt samplet BRTS een pool van $N$ onvoorwaardelijke teacher rollouts ($\{y_{T,i}\}_{i=1}^N$). Het selectieproces volgt een prioriteitshiërarchie:

1. Prioriteit 1: Juistheid. Het systeem filtert op rollouts waarbij het geëxtraheerde eindantwoord overeenkomt met de grondwaarheid ($y^*$).
2. Prioriteit 2: Student Afstemming. Als meerdere correcte rollouts bestaan, selecteert BRTS die met de hoogste token-niveau overlap met de top-$K$ kandidaatset van de student. Dit zorgt ervoor dat de geselecteerde trajecton compatibel is met de huidige redeneerverdeling van de student.
3. Grondwaarheid-Gestuurde Herstel (Tier-2). Als alle onvoorwaardelijke teacher samples falen (d.w.z. geen enkele is correct), voert BRTS een herstelpas uit. Het construeert een gewijzigde prompt die stilzwijgend het grondwaarheidsantwoord injecteert als een validatiesignaal, terwijl de leraar wordt geïnstrueerd om een natuurlijke afleiding te produceren. Er wordt een nieuwe rollout gesampled; indien correct, wordt deze geselecteerd. Als dit ook faalt, grijpt het systeem terug naar de onvoorwaardelijke rollout met de hoogste student overlap.

2.2 Teacher-Context Supervisie

Het trainingsdoel combineert twee verliescomponenten:

• Student-Context Verlies ($L_{stu-ctx}$): Het standaard OPD-doel, waarbij de reverse KL-divergentie tussen de student- en lerarenverdelingen wordt geminimaliseerd, voorwaardelijk op het door de student gegenereerde voorvoegsel. Dit corrigeert de student op de staten die het daadwerkelijk bezoekt.
• Teacher-Context Verlies ($L_{tea-ctx}$): Een auxiliair verlies berekend op de geselecteerde teacher trajecton ($y'$). Hier worden zowel de student- als lerarenverdelingen voorwaardelijk op het voorvoegsel van de leraar. Dit stelt de student bloot aan een coherent, correct redeneerpad dat buiten het huidige hoogstwaarschijnlijkheidsgebied van de student kan liggen.

Het totale verlies wordt gedefinieerd als:
$$L_{total} = L_{stu-ctx} + \lambda L_{tea-ctx}$$
waarbij $\lambda$ een wegingscoëfficiënt is (empirisch ingesteld op 10).

2.3 Top-K Kandidaat Richting

Om deze verliezen efficiënt te implementeren, maakt BRTS gebruik van top-$K$ aggregatie:

• Voor de student-context tak wordt de kandidaat-tokenset afgeleid van de top-$K$ tokens van de student onder het student-voorvoegsel.
• Voor de teacher-context tak wordt de kandidaatset afgeleid van de top-$K$ tokens van de leraar onder het geselecteerde teacher-voorvoegsel. Dit stelt de auxiliaire tak in staat tokens te superviseren die door de leraar worden geprefereerd maar die de student momenteel misschien niet overweegt, terwijl een compacte kandidaatset wordt behouden.

3. Belangrijkste Bijdragen

Het paper vat zijn bijdragen als volgt samen:

1. Identificatie van Supervisie met Hoge Variantie: De auteurs benadrukken dat standaard student-context supervisie in OPD wordt beperkt door ruizige student-voorvoegsels en stochastische teacher rollouts, waardoor het leraren signaal onbetrouwbaar wordt op moeilijke prompts.
2. Introductie van BRTS: Een lichtgewicht framework voor het selecteren en distilleren van betrouwbare teacher trajecten. Het hanteert een "eerst correctheid, dan afstemming" selectieregel en bevat een grondwaarheid-gestuurde fallback voor gefaalde onvoorwaardelijke samples.
3. Empirische Verbetering: BRTS presteert aanzienlijk beter dan standaard OPD op moeilijke redeneerbenchmarks. Het paper biedt diagnostische analyses die aantonen hoe de kwaliteit van teacher kandidaten, hersteldekking en backbone-keuzes de prestaties beïnvloeden.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op AIME 2024, AIME 2025 en AMC 2023 met een 1.5B studentmodel (DeepSeek-1.5B) en diverse teacher configuraties (JustRL-1.5B en DeepSeek-R1-Distill-Qwen-7B).

• Prestatiewinst: BRTS verbetert consequent ten opzichte van de baseline die alleen de student gebruikt. De winst is het meest uitgesproken op de moeilijkere AIME 2025 dataset. Bijvoorbeeld, op AIME 2024 behaalde het gebruik van vier teacher kandidaten een gemiddelde nauwkeurigheid van 0.4000, vergeleken met 0.3917 voor de baseline.
• Impact van Tier-2 Herstel: De grondwaarheid-gestuurde herstelpas (Tier-2) biedt aanzienlijke verbeteringen op moeilijke prompts waar onvoorwaardelijke sampling faalt. Op AIME 2025 verhoogde het toevoegen van Tier-2 de meerderheidsstem-nauwkeurigheid van 0.2748 (alleen Tier-1) naar 0.3133.
• Vroege Trainingspieken: BRTS bereikt een hogere prestatiepiek in vroege trainingsfasen in vergelijking met baselines, wat wijst op snellere convergentie en een betere compute-prestatie trade-off.
• Robuustheid: De methode blijft effectief bij het wisselen van teacher backbones (bijvoorbeeld het gebruik van een 7B leraar voor een 1.5B student) en is niet breekbaar bij lichte prompt-perturbaties, die daadwerkelijk helpen teacher samples te decorreleren.
• Rekenkosten: De overhead is bescheiden; het vergroten van de teacher kandidatenpool van 1 naar 2 voegt ongeveer 59 seconden per trainingsstap toe op 8×B200 GPUs, wat als acceptabel wordt beschouwd gezien de prestatiewinsten.

5. Betekenis en Claims

Het paper claimt dat BRTS een kritieke bottleneck in on-policy distillatie aanpakt: de onbetrouwbaarheid van stochastische leraren signalen op moeilijke taken. Door de teacher supervisie te structureren om zowel juistheidsbewust (fouten filteren) als afstemmingsbewust (afstemmen op de verdeling van de student) te zijn, transformeert BRTS stochastische teacher rollouts naar gestructureerde, hoogwaardige supervisie.

De auteurs positioneren BRTS niet als een vervanging voor OPD, maar als een noodzakelijke uitbreiding die studenten in staat stelt te leren van volledige, betrouwbare redeneerpaden zonder de catastrofale vergetelheid die vaak geassocieerd wordt met standaard supervised fine-tuning, of de hoge variantie van standaard OPD. De methode wordt gepresenteerd als een schaalbare, lichtgewicht uitbreiding die redeneercapaciteiten verbetert, met name in domeinen waar grondwaarheidsantwoorden beschikbaar zijn voor verificatie.