Technische Samenvatting: $\beta$-OPSD

1. Probleemstelling

On-policy self-distillation (OPSD) is een veelbelovende techniek geworden voor het verbeteren van de redeneercapaciteiten van grote taalmodellen (LLM's). In standaard OPSD leert een studentmodel van zijn eigen gesamplede trajecten, terwijl het tegelijkertijd supervisie ontvangt van een "geprivilegieerde docent" (dezelfde modelarchitectuur met toegang tot grondwaarheid-oplossingen of geverifieerde sporen) die alleen tijdens de training beschikbaar is.

De auteurs stellen echter vast dat vanilla OPSD in de praktijk vaak broos is en aanzienlijke engineeringinspanning vereist om betrouwbaar te functioneren. Het kernprobleem is dat standaard OPSD de docent behandelt als het unieke optimalisatiedoel, waarbij de regularisatiestuctuur die de student aan zijn initiële of referentiebeleid verankert, impliciet wordt vastgelegd. Deze rigide formulering mist een mechanisme om te controleren hoe agressief de student zich verwijdert van zijn referentiebeleid richting de docent, wat vaak leidt tot instabiliteit of suboptimale convergentie.

2. Methodologie

Het artikel stelt $\beta$-OPSD voor, een framework dat on-policy self-distillation herformuleert als een bredere familie van KL-geregulariseerde beleidsoptimalisatieproblemen. De methodologie verloopt in drie hoofdfasen: het afleiden van het theoretische doel, het omzetten naar een praktisch distillatie-doel, en het verfijnen van de credit assignment.

2.1. Theoretische Afleiding: $\beta$-OPSD als Beleidsoptimalisatie

De auteurs demonstreren dat vanilla OPSD precies het $\beta = 1$ lid is van een familie van KL-geregulariseerde beleidsoptimalisatie. Ze definiëren de doelfunctie $J_\beta(\theta)$ als:
$$J_\beta(\theta) = \mathbb{E}_{y \sim \pi_\theta} \left[ \log \frac{p_T(y|x,c)}{\pi_{\text{ref}}(y|x)} \right] - \beta D_{KL}(\pi_\theta(\cdot|x) \parallel \pi_{\text{ref}}(\cdot|x))$$
Waar:

• $\pi_\theta$: Het studentbeleid.
• $p_T$: Het geprivilegieerde docentbeleid (geconditioneerd op probleem $x$ en geprivilegieerde informatie $c$).
• $\pi_{\text{ref}}$: Een referentiebeleid (bijv. de initiële student).
• $\beta$: Een controleerbare regularisatieparameter.

Kerninzicht: Wanneer $\beta=1$, vallen de referentietermen weg, waardoor de standaard reverse KL-minimalisatie naar de docent wordt hersteld. Voor $\beta > 1$ wordt het doel conservatiever, waarbij afwijkingen van het referentiebeleid zwaarder worden bestraft.

2.2. Optimaal Beleid en Geometrische Interpolatie

De auteurs leiden het gesloten vorm optimal beleid $\pi^\star_\beta$ af voor dit doel:
$$\pi^\star_\beta(y|x,c) = \frac{\pi_{\text{ref}}(y|x)^{1-1/\beta} p_T(y|x,c)^{1/\beta}}{Z_\beta(x,c)}$$
Deze oplossing vertegenwoordigt een geometrische interpolatie tussen het referentiebeleid en de geprivilegieerde docent. De parameter $\beta$ (specifiek de inverse $1/\beta$) fungeert als een gewicht dat de positie langs dit pad bepaalt.

• $\beta = 1 \implies$ Het doel is de docent.
• $\beta \to \infty \implies$ Het doel nadert het referentiebeleid.

2.3. Praktische Training: Gescheduled Logit-interpolatie

Directe optimalisatie van het RL-doel is computationeel duur vanwege de onhandelbare normalisator $Z_\beta$ en gradiënten met een hoge variantie. In plaats daarvan zet $\beta$-OPSD de theoretische oplossing om in een efficiënt distillatie-doel.

• Lokale Approximatie: Omdat de sequentie-niveau normalisator onhandelbaar is, benadert de methode het optimale beleid bij elke decoderingsprefix $h_t$ door de logits van de referentie en de docent lineair te interpoleren:
  $$\tilde{p}_{\beta_k}(\cdot | h_t, c) = \text{softmax}\left( \left(1 - \frac{1}{\beta_k}\right) z_{\text{ref}} + \frac{1}{\beta_k} z_T \right)$$
• Curriculum Scheduling: In plaats van $\beta$ vast te leggen, gebruikt de methode een schema waarbij het docentgewicht $w_k = 1/\beta_k$ lineair toeneemt van $w_{\text{start}}$ naar $w_{\text{end}}$ over de trainingsstappen. Dit creëert een vloeiend curriculum, beginnend met doelen die dichter bij de student liggen en geleidelijk sterkere docentbegeleiding op te nemen.

2.4. Return-to-Go Credit Assignment

Standaard token-niveau OPSD updates zijn "myopic" (kortzichtig), waarbij ze elke token-loss onafhankelijk behandelen. Dit negeert hoe een vroege token-keuze de toekomstige prefixes en de daaropvolgende mismatch met het doel beïnvloedt.
Om dit aan te pakken, introduceert $\beta$-OPSD return-to-go (RTG) credit assignment. De loss voor een token $t$ wordt gewogen door de cumulatieve toekomstige mismatch (gedisconteerde return) langs het traject:
$$G_{t, \gamma} = \sum_{s=t}^{T} \gamma^{s-t} \rho_s$$
waarbij $\rho_s$ de log-ratio mismatch is op stap $s$. Dit zorgt ervoor dat vroege tokens krediet (of schuld) krijgen voor de distributionele mismatch die ze later in de sequentie induceren, waardoor de token-niveau updates worden uitgelijnd met het sequentie-niveau doel.

3. Belangrijkste Bijdragen

1. Theoretische Unificatie: Het artikel stelt vast dat vanilla OPSD het $\beta=1$ geval is van een bredere KL-geregulariseerde beleidsoptimalisatiefamilie, waardoor de impliciete regularisatiekeuze expliciet en controleerbaar wordt.
2. $\beta$-OPSD Framework: Het introduceert een gegeneraliseerd doel waarbij het optimale beleid een geometrische interpolatie is tussen de referentie en de docent, wat in de praktijk wordt gerealiseerd via gescheduled logit-interpolatie.
3. Efficiënt Trainingsalgoritme: De auteurs stellen een methode voor om de dure RL-afgeleide oplossing te benaderen met goedkope token-niveau distillatie, waardoor complexe RL-mechanismen zoals advantage estimation worden vermeden.
4. Sequentie-niveau Uitlijning: De integratie van return-to-go credit assignment corrigeert de myopische aard van standaard token-KL updates, en biedt een onbevooroordeelde estimator voor het sequentie-niveau doel.

4. Experimentele Resultaten

De methode werd geëvalueerd op wiskundige redeneerbenchmarks (AIME 2024, AIME 2025, HMMT 2025) met behulp van Qwen3-modellen variërend van 1.7B tot 8B parameters.

• Prestatiewinst: $\beta$-OPSD presteerde consequent beter dan vanilla OPSD, SFT en GRPO baselines.
  • Op Qwen3-1.7B verbeterde $\beta$-OPSD de gemiddelde prestatie (avg@12) met 5,74 procentpunten ten opzichte van vanilla OPSD, met specifieke winsten van 9,16 punten op AIME 2024.
  • Op grotere modellen (4B en 8B) behield het een competitieve of superieure prestatie, met een verbetering van respectievelijk 1,76 en 1,66 punten in het algemeen gemiddelde.
• Ablatie-studies:
  • Doelverdeling: Het vervangen van de vaste docent-target door de gescheduled logit-interpolant target leverde significante winsten op (bijv. +6,03 punten op AIME 2024), wat het voordeel van een vloeiende transitie bevestigt.
  • Credit Assignment: Het gebruik van return-to-go gewichten in plaats van lokale token-gradiënten verbeterde de prestaties (bijv. +5,55 punten op AIME 2025), wat de belangrijkheid van toekomst-bewuste updates valideert.
  • Referentieconstructie: De beste prestatie werd bereikt met een dynamische student (huidig model) en een vaste docent (geprivilegieerde docent), wat suggereert dat het volgen van de evoluerende studentdistributie terwijl een stabiel toezichthoudend anker behouden blijft, optimaal is.

5. Betekenis en Claims

Het artikel claimt een "geprincieerde route van self-distillatie naar beleidsoptimalisatie en terug" te bieden. De primaire betekenis ligt in:

• Stabiliteit: Door $\beta$ als een controleerbare parameter te behandteren, adresseert de methode de broosheid van vanilla OPSD, waardoor een vloeiendere optimalisatie mogelijk is die een afweging maakt tussen nabijheid tot het referentiebeleid en docentbegeleiding.
• Efficiëntie: Het bereikt de voordelen van beleidsoptimalisatie (sequentie-niveau doelen, KL-regularisatie) zonder de computationele kosten van volledige RL, waarbij de efficiëntie van standaard distillatie behouden blijft.
• Generaliseerbaarheid: De aanpak biedt een mechanisme om de distributionele kloof tussen student en docent te verkleinen, wat bijzonder gunstig is voor kleinere modellen waar deze kloof groter is.

De auteurs concluderen dat vloeiende doelpaden en traject-niveau credit assignment effectieve principes zijn voor on-policy distillatie, wat een robuust alternatief biedt voor directe docent-imitatie.