Technische Samenvatting: Relay-OPD (De estafettestok doorgeven)

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

On-policy distillatie (OPD) verbetert de kennisoverdracht van sterke leraarmodellen naar zwakkere studentmodellen door token-niveau supervisie te funderen in de eigen gegenereerde trajecten van de student. Deze aanpak lijdt echter aan een kritiek gebrek dat bekend staat als prefix-falen. In taken met lange ketens van redeneringen (long-chain reasoning), als een studentmodel zich vroeg in de generatie vastlegt op een onjuiste redeneerrichting, bouwen alle daaropvolgende tokens voort op deze afwijking. Dit resulteert in lange, misgestuurde voortzettingen die:

1. Onbetrouwbare en potentieel schadelijke supervisie van de leraar uitlokken.
2. Substantieel trainingscompute verspillen aan trajecten die niet meer gecorrigeerd kunnen worden.

Bestaande oplossingen hebben structurele beperkingen:

• Vaste lengte-afkap (bijv. FastOPD, ESR): Snijdt rollouts af op rigide posities, ongeacht waar de werkelijke redeneerfout plaatsvindt.
• Offline herschrijven (bijv. TRD): Repareert trajecten pas na voltooiing, wat vaak zichtbare artefacten achterlaat.
• Token-niveau mixing (bijv. SKD): Schakelt over op basis van generieke distributieve discrepantie in plaats van expliciete signalen van redeneerfalen.

Er is een gebrek aan een mechanisme dat online werkt, fouten corrigeert terwijl ze ontstaan, en ingrijpt op basis van de specifieke redeneerstatus zelf.

2. Methodologie: Relay On-Policy Distillatie (Relay-OPD)

De auteurs stellen Relay-OPD voor, een methode die studentgeneratie afwisselt met korte, gerichte overnames door de leraar op gedetecteerde foutpunten.

2.1 Kerninzicht: Asymmetrie tussen Leraar en Student

De methode berust op een geobserveerde asymmetrie bij gefaalde prefixes:

• De Student: Neigt om langs de huidige (foutieve) richting door te gaan.
• De Leraar: Neigt te stoppen, te heroverwegen en de redenering te heroriënteren (vaak met "reflectie"-tokens zoals Maar, Wacht, Echter).

Deze divergentie dient als een label-vrije handoff-trigger. Wanneer het meest waarschijnlijke volgende token van de leraar een reflectie-token is, maar de top-$K$ support set van de student geen dergelijke tokens bevat, wordt een handoff getriggerd.

2.2 Constructie van het Relay-traject

Wanneer een trigger afgaat, neemt de leraar kort de controle over om een leraar-segment (teacher leg) te produceren, waarna de student de generatie hervat.

• Handoff Trigger ($\phi(h)$): Gedefinieerd als $1[a_T(h) \in \mathcal{R}] \cdot 1[K_S(h) \cap \mathcal{R} = \emptyset]$, waarbij $\mathcal{R}$ de verzameling reflectie-tokens is, $a_T$ het top-token van de leraar is, en $K_S$ de top-$K$ set van de student is.
• Relay Budget ($M, L$): Om te voorkomen dat het traject te ver afwijkt van het beleid van de student, begrenst een budget de interventie:
  • $M$: Maximaal aantal leraar-overnames.
  • $L$: Aantal extra paragrafen gegenereerd door de leraar na het reflectie-token.
• Proces: De student genereert autoregressief. Bij detectie van een trigger (en indien het budget nog niet is uitgeput), genereert de leraar $L$ paragrafen (beginnend met het reflectie-token). De student hervat vervolgens de generatie vanaf de gecorrigeerde prefix.

2.3 Optimalisatie-doelstelling

Relay-OPD optimaliseert het studentbeleid op het resulterende relay-traject $z$. In tegenstelling tot standaard OPD, die de originele tokens van de student gebruikt, gebruikt Relay-OPD de daadwerkelijk gegenereerde tokens $z_t$ uit het relay-traject (inclusclusief de door de leraar gegenereerde tokens).
De doelstelling gebruikt een reverse-KL-stijl single-sample advantage:
$$A_t^{Relay} = \log \pi_T(z_t | h_t^z) - \log \bar{\pi}_\theta(z_t | h_t^z)$$
Dit wordt geoptimaliseerd via een clipped PPO-stijl loss, waardoor de student selectief de corrigerende signalen van de leraar kan absorberen in plaats van volledig de distributie van de leraar te fitten.

2.4 Efficiënte Implementatie

Om de overhead van het schakelen tussen twee generatie-pipelines te vermijden, verenigt Relay-OPD het proces binnen een enkele speculative decoding engine:

• Student fungeert als het draft-model.
• Leraar fungeert als het target-model.
• De engine schakelt dynamisch tussen het target-beleid ($\pi_{tgt}$) op basis van de huidige staat (Student leg versus Teacher leg).
• Teacher legs voeren standaard speculative decoding uit tegen de leraar, terwijl student legs alle drafts accepteren. Dit zorgt ervoor dat de distributieve identiteit van het relay-traject exact behouden blijft.

3. Belangrijkste Bijdragen

1. Identificatie van Asymmetrie: Het paper identificeert dat de divergentie tussen leraar en student op gefaalde prefixes een betrouwbaar, label-vrij signaal is voor interventie. Experimenten tonen aan dat het corrigeren van prefix-falen slechts vroege, lokale interventie vereist, waarbij de leraar-tokens slechts 0,35% van de totale gegenereerde tokens beslaan.
2. Relay-OPD Framework: Een nieuwe methode die deze asymmetrie omzet in een budgetgestuurd relay-rollout, waarbij student- en leraargeneratie worden verenigd in een enkele speculative decoding engine.
3. Empirische Superieuriteit: Aangetoond over acht wiskundige redeneer-benchmarks en twee student-schalen (0.6B en 1.7B).

4. Experimentele Resultaten

De auteurs evalueerden Relay-OPD met behulp van Qwen3-4B-Instruct-2507 als leraar en Qwen3-0.6B/1.7B-Non-Thinking als studenten op acht benchmarks (AIME 2024/2025/2026, MATH500, AMC 2023, OlympiadBench, HMMT Feb/Nov).

• Prestaties:
  • 1.7B Student: Behaalde de beste of op één na beste resultaten op alle benchmarks.
    • Gemiddelde nauwkeurigheid: 46,96% (tegenover 41,23% voor standaard OPD, een verbetering van +5,73%).
    • Presteerde beter dan de sterkste baseline, FastOPD, met +1,49%.
  • 0.6B Student: Consistente winst, met een verbetering van de gemiddelde nauwkeurigheid met +3,01% ten opzichte van standaard OPD.
• Efficiëntie:
  • Trainingslengte: Verminderde de gemiddelde trainings-trajectlengte met >50% (van 4.658 naar 2.296 tokens voor 1.7B).
  • Convergentie: Bereikte het beste checkpoint in minder stappen (Stap 35 versus Stap 55 voor OPD).
  • Inference: Verminderde de gemiddelde responslengte met 14–28% vergeleken met FastOPD terwijl de nauwkeurigheid verbeterde.
• Ablatie-studies:
  • Noodzaak van de Teacher Leg: Het simpelweg stoppen bij de trigger (Trigger-stop) leverde 43,48% nauwkeurigheid op; het toevoegen van de leraar-leg ($L=3$) verhoogde dit naar 46,25%, wat de waarde van gecorrigeerde context bewijst.
  • Doelstelling: De reverse-KL-stijl doelstelling op gegenereerde tokens presteerde beter dan zowel de "Student Draft Token" doelstellingen als de "Teacher FKL" doelstellingen, wat het voordeel bevestigt van selectieve absorptie van corrigerende signalen.

5. Betekenis en Claims

Het paper claimt dat prefix-falen een fundamentele bottleneck is in on-policy distillatie die bestaande truncatie- of herschrijfmethoden niet dynamisch kunnen aanpakken. Relay-OPD adresseert dit door:

• Online te handelen: Fouten te corrigeren terwijl ze ontstaan, in plaats van achteraf (post-hoc).
• Interventie te minimaliseren: Een "relay"-aanpak te gebruiken waarbij de leraar slechts kort ingrijpt om te heroriënteren, waardoor het traject dicht bij het beleid van de student blijft.
• Efficiëntie: De rekenkosten van training aanzienlijk te verlagen door trajecten te verkorten en verspilde berekening op onherstelbare fouten te vermijden.

De auteurs concluderen dat Relay-OPD een robuust, label-vrij mechanisme biedt voor sterke-naar-zwakke distillatie in redeneertaken, waarbij state-of-the-art resultaten worden behaald terwijl de trainingskosten drastisch worden verlaagd.