Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Technische Samenvatting: Het overbruggen van redeneertrajecten in On-Policy Distillatie via Nabije-Toekomst-Begeleiding

Probleemstelling
On-Policy Distillatie (OPD) is een hoeksteen van moderne post-training voor Large Language Models (LLM's), ontworpen om redeneren te verbeteren door een studentmodel te trainen op trajecten die gesampled zijn uit zijn eigen beleid onder toezicht van een docent. Standaard OPD werkt volgens een "getraject-gesampled maar token-geleerd" mechanisme: het identificeert afwijkingen via tokens met een hoge loss en herstelt deze middels lokale reverse-KL divergentie correctie.

De auteurs identificeren een fundamentele mismatch in deze aanpak, waarbij zij beargumenteren dat token-niveau supervisie structureel niet is afgestemd op de vereisten van traject-niveau redeneercorrectie. Deze mismatch manifesteert zich in twee kritieke faalmodi:

1. Ruisende Vals Alarmen: Tokens met een hoge loss komen niet altijd overeen met echte redeneer-divergentie. Empirische analyse onthult dat ongeveer 30% van de tokens met een hoge loss valt binnen een regime van lage divergentie, wat staat voor oppervlakkige vorm-mismatches (bijv. stilistische voorkeuren zoals "en" versus "daarna") in plaats van logische splitsingen. Het behandelen van deze als kritieke fouten levert ruisende supervisie die actief kan interfereren met de redeneercorrectie.
2. De Token-voor-Token Leerval: Zelfs wanneer genu genuine divergentiepunten worden geïdentificeerd, is single-token reverse-KL correctie vaak onvoldoende om het gehele redeneepad te heroriënteren. Redeneerfouten evolueren doorgaans als distributieve drift met een korte horizon. Het corrigeren van één token lokaal garandeert niet dat het model zal uitlijnen met het toekomstige traject van de docent; het model wijkt vaak direct na de lokale correctie weer af naar foutieve takken.

Methodologie: Traject-bewuste OPD (TOPD)
Om deze beperkingen aan te pakken, stellen de auteurs Trajectory-aware OPD (TOPD) voor, die de kloof tussen lokale token-correctie en globale traject-evolutie overbrugt door nabije-toekomst trajectinformatie te incorporeren. De methode bestaat uit twee primaire stadia:

1. Divergentie Detectie:

   • Voor hoog-loss kandidaatposities stapt het systeem een stap terug naar het voorgaande token om een gedeelde prefix vast te stellen.
   • Zowel de docent als de student genereren korte-venster continuaties (lengte $K=50$) vanaf deze prefix.
   • Het systeem meet de traject-niveau discrepantie met behulp van de Optimal Transport (OT) afstand tussen de continuaties van de docent en de student.
   • Echte divergente punten worden geïdentificeerd als posities die zowel een hoge token-niveau loss als een hoge korte-venster OT-divergentie vertonen. Dit filtert benigne oppervlakkige mismatches eruit.

2. Injectie van Traject-niveau Docent-signaal:

   • In plaats van alleen de huidige token te optimaliseren, breidt TOPD de supervisie uit naar het volledige korte-venster traject.
   • Gebruikmakend van het OT-transportplan ($\gamma$), construeert de methode een "traject-bewuste zachte doelstelling" voor elke toekomstige student-positie. Deze doelstelling stemt de toekomstige tokens van de student af op het traject van de docent op basis van de optimale transportmapping.
   • De student wordt getraind om deze OT-afgestemde zachte doelstellingen over het venster te matchen, waardoor het effectief leert hoe het zijn nabije-toekomst redeneepad naar het pad van de docent kan sturen, in plaats van alleen te leren hoe het de huidige token moet corrigeren.
   • De uiteindelijke objectieve functie combineert de standaard OPD loss met een traject-loss term ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

Belangrijkste Bijdragen

• Identificatie van Mismatch: Het werk stelt vast dat token-niveau supervisie een "bot proxy" is voor traject-niveau redeneercorrectie, die faalt in het onderscheiden van genuine logische splitsingen van benigne oppervlaktevariaties.
• Empirische Karakterisering van Faalmodi: De auteurs leveren bewijs dat ~30% van de hoge-loss tokens lage-divergentie vals alarmen zijn en demonstreren de "token-voor-token leerval" waarbij lokale correcties falen om het algehele redeneepad te herstellen.
• Traject-niveau Begeleidingsprincipe: Het artikel introduceert een principe waarbij nabije-toekomst trajectinformatie over meerdere tokens wordt verdeeld. Door gebruik te maken van korte-venster OT-afstemming, injecteert TOPD docent-student pad-discrepanties direct in de loss-functie, waardoor de student in staat is om traject-verschuivingen te leren in plaats van geïsoleerde token-bewerkingen.

Experimentele Resultaten
Experimenten werden uitgevoerd met Qwen3-modellen (het distilleren van een 30B docent naar een 4B student) op competitieve wiskunde benchmarks (AIME24, AIME25, HMMT25-Feb).

• Prestatiewinsten: Standaard OPD verbeterde de gemiddelde nauwkeurigheid van het warm-start model van 38,9% naar 47,8%. TOPD verhoogde dit verder naar 52,2%.
• Specifieke Benchmarks: TOPD verbeterde de AIME24 nauwkeurigheid van 60,0% naar 63,3% en AIME25 van 46,7% naar 53,3%.
• Ablatie-studies:
  • Het onderdrukken van lage-OT hoge-loss tokens (vals alarm) in standaard OPD leverde een bescheiden winst op (+0,4%), wat bevestigt dat deze tokens ruisend zijn.
  • Het onderdrukken van hoge-OT hoge-loss tokens (echte divergenties) veroorzaakte een significante prestatiedaling, wat hun noodzakelijkheid bevestigt.
  • Analyse van lokale traject-correctie toonde aan dat terwijl standaard OPD de korte-venster OT afstand gedeeltelijk verminderde (van 0,538 naar 0,350), TOPD een substantiële reductie bereikte (naar 0,204), wat een superieure bekwaamheid demonstreert om nabije-toekomst trajecten te heroriënteren.
• Efficiëntie: TOPD introduceert een bescheiden computationele overhead, waarbij de trainingstijd 1,41$\times$ die van standaard OPD is, aangezien OT alleen wordt berekend op geselecteerde hoge-loss vensters.

Significantie en Claims
Het artikel claimt dat redeneerfouten in OPD fundamenteel traject-niveau fenomenen zijn in plaats van geïsoleerde token-mismatches. Door te verschuiven van puntwijze token-bewerking naar globale, transitie-bewuste afstemming, maakt TOPD effectievere redeneer-afstemming mogelijk. De auteurs positioneren dit werk als een stap naar meer traject-bewuste distillatie-objectieven, waarbij zij beargumenteren dat succesvolle redeneercorrectie vereist dat de student niet alleen leert welke token te corrigeren, maar ook hoe het daaropvolgende redeneepad te navigeren om consistent te blijven met de docent.

Beperkingen
De auteurs erkennen dat TOPD extra computationele kosten met zich meebrengt door OT-berekeningen, zich primair richt op korte-horizon divergentie (waardoor potentieel fouten met een lange context gemist kunnen worden), en momenteel voornamelijk gevalideerd is op wiskundige redeneer-benchmarks in plaats van code of open-ended taken. Daarnaast meet de OT-afstand de nabijheid van trajecten, maar legt het mogelijk geen semantische equivalentie vast waarbij verschillende paden tot hetzelfde correcte antwoord leiden.