🤖 AI

Trajectory-Refined Distillation

Dit artikel introduceert Trajectory-Refined Distillation (TRD), een nieuwe on-policy distillatiemethode die het structurele "prefix failure"-probleem mitigeert door traject-niveau correcties toe te passen op student-rollouts onder begeleiding van een docent, waardoor de redeneernauwkeurigheid en dekking over diverse benchmarks en modelschalen worden verbeterd.

Oorspronkelijke auteurs: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Gepubliceerd 2026-06-09

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een student (de AI) leert hoe hij een complex wiskundig probleem moet oplossen. Je hebt een briljante leraar (een meer geavanceerde AI) die het juiste antwoord weet.

In de standaardmethode beschreven in het paper, genaamd On-Policy Distillation (OPD), werkt het proces als volgt:

De student probeert het probleem zelf op te lossen.
Terwijl de student zijn stappen opschrijft, kijkt de leraar toe en zegt: "Voor dit specifieke woord dat je net schreef, is dit de waarschijnlijkheid van wat er nu volgt."
De student probeert zijn brein aan te passen om de waarschijnlijkheden van de leraar voor elk woord te evenaren.

Het Probleem: "Prefix Failure" (De Foutieve Afslag)

Het paper identificeert een groot gebrek in deze standaardmethode, genaamd Prefix Failure.

Stel je voor dat de student een pad afloopt om het probleem op te lossen. Bij stap 3 slaat hij een verkeerde afslag. Hij is nu terechtgekomen in een doodlopende straat.

Het Dilemma van de Leraar: De leraar ziet dat de student verdwaald is. De leraar weet dat het juiste pad begint met "Wacht, eigenlijk..." en dat het daarna een compleet andere richting op gaat.
De Verwarring: Maar de leraar wordt gedwongen om naar de huidige locatie van de student te kijken (de doodlopende weg). De adviezen van de leraar worden daarom een verwarrende mix: "Blijf op deze doodlopende weg omdat je er al bent" EN "Draai onmiddellijk om."
Het Resultaat: De student krijgt een verwarrend signaal. Hij probeert te leren van de leraar, maar omdat de student vastzit op het verkeerde pad, raakt het advies van de leraar gefragmenteerd. De student leert "goed te zijn in verdwaald zijn" in plaats van te leren hoe hij weer op het juiste spoor komt. Het paper noemt dit een "bimodale mengeling"—de leraar probeert tegelijkertijd twee tegenovergestelde dingen te onderwijzen, en de student raakt in de war.

Bestaande oplossingen probeerden dit te verhelpen door simpelweg het "cijfer" (de loss function) voor specifieke woorden aan te passen, maar ze losten het feit niet op dat de student nog steeds op het verkeerde pad vastzat.

De Oplossing: Trajectory-Refined Distillation (TRD)

De auteurs stellen een nieuwe methode voor genaamd Trajectory-Refined Distillation (TRD). In plaats van alleen de slordige conceptversie van de student woord voor woord te beoordelen, veranderen ze de workflow:

Het Concept: De student maakt nog steeds een eerste poging tot het probleem (de "raw rollout").
Het Herschrijven: Voordat de beoordeling plaatsvindt, neemt de leraar dat slordige concept en schrijft het herschreven. De leraar corrigeert de verkeerde afslagen, herstelt de logica en produceert een zuivere, perfecte versie van de oplossing gebaseerd op het startpunt van de student.
De Les: De student leert vervolgens van deze zuivere, gecorrigeerde versie, en niet van de slordige originele versie.

De Analogie:

De Oude Manier: Je schrijft een verhaal met een plotgat in het midden. Je redacteur leest het en zegt: "Voor het woord 'kat' in paragraaf 3 had je 'hond' moeten schrijven." Maar je zit nog steeds vast in paragraaf 3 terwijl je probeert uit te zoeken waarom je 'kat' schreef in de eerste plaats. Het is verwarrend.
De TRD-Manier: Je schrijft een verhaal met een plotgat. Je redacteur neemt jouw concept, herstelt het plotgat en schrijft het hele middendeel herschreven zodat het verhaal perfect doorloopt. Daarna bestudeer jij de herschreven versie van de redacteur om de volgende keer beter te leren verhalen te schrijven.

Waarom Dit Beter Werkt

Lost de Oorzaak Aan: Door de "verkeerde afslag" te corrigeren voordat de student probeert te leren van de fout, is het advies van de leraar niet langer verwarrend. Het is één helder pad.
Ontdekt Nieuwe Paden: Zelfs als de student het antwoord de eerste keer goed heeft, kan de leraar hem een andere, kortere of slimmere manier laten zien om het probleem op te lossen. Dit leert de student om flexibeler te zijn.
Efficiëntie: Het paper ontdekte dat de "herschreven" paden vaak veel korter en duidelijker waren dan de oorspronkelijke, oeverloze pogingen van de student, wat het trainen daadwerkelijk sneller maakte.

De Resultaten

De onderzoekers testten dit op moeilijke wiskundige competities (zoals AIME en HMMT) en programmeeruitdagingen. Ze ontdekten dat:

De nieuwe methode (TRD) de oude methoden consequent versloeg.
Het vooral effectief was in het helpen van de AI bij de moeilijkste problemen waar het normaal gesproken op vastliep.
Het werkte, ongeacht of de leraar een aparte, grotere AI was of dezelfde AI die optrad als zijn eigen "slimme" leraar (met behulp van een techniek genaamd "privileged information").

Kortom, het paper betoogt dat om een AI effectief te onderwijzen, je niet alleen zijn fouten woord voor woord moet beoordelen terwijl hij ze maakt. In plaats daarvan moet je de leraar eerst het hele pad laten corrigeren, en de student vervolgens laten leren van de gecorrigeerde reis.

Technische Samenvatting: Trajectory-Refined Distillation (TRD)

1. Probleemstelling: Prefix-falen in On-Policy Distillatie

Het artikel identificeert een fundamentele structurele beperking in On-Policy Distillatie (OPD) en On-Policy Self-Distillatie (OPSD) voor Large Language Models (LLM's), die wordt aangeduid als prefix-falen (prefix failure).

In standaard OPD/OPSD genereert een studentmodel een rollout (traject) $y_o$ , en een teacher-model biedt dichte, per-token supervisie op basis van dat specifieke traject. De auteurs stellen dat wanneer het traject van de student een onjuist redeneerpad volgt (een "gefaalde prefix"), de supervisiesignaal van de teacher structureel wordt aangetast:

Bimodale Teacher-mengeling: Wanneer een prefix onjuist is, wordt de distributie van de teacher een mengeling van twee modi: één die de gefaalde prefix voortzet (voor sequentieconsistentie) en een andere die naar de juiste oplossing pivot (draait).
Gefragmenteerde Gradiënten: Standaard per-token verliesfuncties (KL-divergentie) evalueren supervisie langs het bevroren gefaalde traject. Gevolgelijk zijn de gradiënt-updates gefragmenteerd. De teacher beveelt herhaaldelijk een correctie-token aan (bijv. "Wacht" of "Eigenlijk") geconditioneerd op de foutieve context (de gefaalde prefix), in plaats van geconditioneerd op het ontrollende correctiepad.
Ineffectiviteit van Token-niveau Fixes: Bestaande remedies (bijv. loss clipping, top-K truncatie, herweging) opereren op het niveau van de token-loss. Ze proberen "informatieve" tokens te selecteren of "ruisende" tokens te onderdrukken, maar falen in het aanpakken van de kernoorzaak: de gefaalde prefix zelf blijft structureel intact. De supervisieparen die aan de student worden geleverd, wijken af van het ideale correctiepad, waardoor het model niet effectief leert om te herstellen van fouten.

2. Methodologie: Trajectory-Refined Distillation (TRD)

Om prefix-falen te mitigeren, stellen de auteurs Trajectory-Refined Distillation (TRD) voor, een methode die de interventie verschuift van het token-niveau naar het traject-niveau.

Kernmechanisme

In plaats van de ruwe student-rollout $y_o$ direct te gebruiken voor distillatie, introduceert TRD een verfijningsstap vóór de berekening van de loss:

Sampling: Het studentmodel $\pi_\theta$ genereert een ruwe on-policy rollout $y_o \sim \pi_\theta(\cdot | x)$ .
Verfijning (Refinement): Het teacher-model $\pi_T$ $π_{T}$ (of de student zelf in OPSD, geconditioneerd op geprivilegieerde informatie zoals de grondwaarheid $y^*$ $y^{*}$ ) genereert een verfijnd traject $y_r$ $y_{r}$ .
- In OPD: $y_r \sim \pi_T(\cdot | x, y_o)$ .
- In OPSD: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ .
- Cruciaal is dat de teacher wordt geprompt om de rollout van de student te herzien geleid door de referentieoplossing, waardoor de foutieve delen van de prefix effectief wordt "herschreven" terwijl de geldige redeneerpatronen van de student behouden blijven.
Distillatie: Het verfijnde traject $y_r$ wordt gebruikt als het supervisiedoel voor de standaard distributie-matching loss (Forward KL met volledige vocabulaire).

Theoretische Rechtvaardiging

TRD adresseert het probleem van "gefragmenteerde gradiënten" door ervoor te zorgen dat de supervisiecontext groeit langs het correctiepad in plaats van langs de gefaalde prefix.

Ideale Gradiënt-herstel: Door te superviseren langs $y_r$ , bevat de context voor elke token het gecorrigeerde prefix, waardoor de teacher een coherent, meerstaps correctiesignaal ( $g_{ideal}$ ) kan bieden in plaats van een repetitief, context-mismatched signaal ( $g_{frag}$ ).
Exploratie: Zelfs wanneer $y_o$ correct is, stelt TRD de student bloot aan alternatieve geldige afleidingen voorgesteld door de teacher, wat de distributie van correcte redeneerpaden die de student leert verbbreedt.

Algoritme

Het proces wordt samengevat in Algoritme 1:

Sample $y_o \sim \pi_\theta(\cdot | x)$ .
Genereer $y_r \sim \pi_T(\cdot | x, y_o)$ .
Update $\theta$ met behulp van de KL-divergentie tussen $\pi_\theta$ en de distributie van de teacher over $y_r$ .

3. Belangrijkste Bijdragen

Identificatie van Prefix-falen: Het artikel formaliseert "prefix-falen" als een structurele oorzaak van degradatie in on-policy distillatie, en legt uit waarom interventies op token-niveau falen bij het herstellen van redeneerfouten.
Correctie op Traject-niveau: TRD wordt voorgesteld als de eerste optimalisatie op traject-niveau die expliciet gericht is op het herstel van gefaalde prefixes, terwijl de on-policy support-restrictie wordt gerespecteerd zonder excessieve computationele overhead.
Verbetering in Dubbele Regimes: TRD verbetert de prestaties in twee regimes:
- Foutcorrectie: Het mitigeert prefix-falen door problematische prefixes te corrigeren vóór de distillatie.
- Exploratie: Het verbredt de redeneerdekking door studenten bloot te stellen aan alternatieve geldige afleidingen, zelfs wanneer het oorspronkelijke traject correct was.
Toepasbaarheid op Self-Distillatie: De methode is natuurlijk uitbreidbaar naar OPSD, waarbij de student fungeert als zijn eigen teacher geconditioneerd op geprivilegieerde informatie, waardoor een aparte teacher-model niet nodig is.

4. Experimentele Resultaten

De auteurs evalueerden TRD over vijf competitie-wiskunde benchmarks (AIME24/25, HMMT25, BeyondAIME, AMOBench) en code-generatie benchmarks (HumanEval+, MBPP+, LiveCodeBench) met behulp van Qwen3-modellen (1.7B, 4B, 8B).

Belangrijkste Bevindingen

Superieure Prestaties: TRD presteert consequent beter dan eerdere baselines (Forward/Reverse KL, met/zonder clipping/truncatie) in zowel OPD- als OPSD-settings.
Significante Winst op Moeilijke Taken: De verbeteringen zijn het meest uitgesproken op de moeilijkste benchmarks.
- AMOBench: TRD bereikte een relatieve verbetering van ~50% in Pass@16 voor Qwen3-8B onder de OPSD-setting vergeleken met het basismodel.
- AIME24/25: TRD toonde consistente winsten, met name voor kleinere modellen (bijv. +4.6% op AIME24 voor Qwen3-1.7B).
Stabiliteit: In tegen tegenstelling tot sommige baselines die presteren degraderen op bepaalde benchmarks (bijv. Reverse KL die prestaties verlaagt op AIME24), zakte TRD nooit onder de prestaties van het basismodel.
Efficiëntie: In de OPSD-setting waren de verfijnde trajecten $y_r$ aanzienlijk korter (mediane lengte gecomprimeerd met ~9x) dan de ruwe rollouts $y_o$ , wat de trainingstijd (wall-clock time) met ~60% verminderde ondanks de extra sampling-stap.
Code Generatie: TRD kwam overeen met of overtrof de baselines op HumanEval+ en MBPP+, hoewel de winst op LiveCodeBench beperkt was, wat suggereert dat de verfijningscapaciteit van de teacher begrensd kan zijn door de taakcomplexiteit.

5. Betekenis en Claims

Het artikel beweert dat TRD een noodzakelijke verschuiving vertegenwoordigt in de post-training strategie voor LLM's. De auteurs argumenteren dat:

Interventies op token-niveau zijn onvoldoende: De structurele aard van prefix-falen vereist correcties op traject-niveau.
On-policy support wordt behouden: TRD bereikt deze correcties zonder de on-policy restrictie te schenden (in tegen tegenstelling tot het gebruik van een vaste expert-policy), wat het compatibel maakt met standaard rekenefficiënte leerparadigma's.
Self-improvement is levensvatbaar: Het succes van TRD in de OPSD-setting demonstreert dat modellen effectief hun eigen redeneerpaden kunnen verfijnen met behulp van geprivilegieerde informatie (grondwaarheid) zonder dat een aparte, grotere teacher-model nodig is, wat een lichtgewicht pad naar zelfverbetering biedt.

Het werk suggereert dat toekomstige distillatiemethoden rekening moeten houden met de structurele integriteit van het redeneertraject, en niet alleen met de statistische uitlijning van individuele tokens.