Technische Samenvatting: Geometrische Zelf-distillatie voor Redeneer-generalisatie

1. Probleemstelling

On-policy distillatie (OPD) is opgekomen als een praktische post-training strategie voor grote taalmodellen (LLM's), waarbij dichte supervisie op token-niveau wordt geboden op trajecten die door het student-model zelf zijn gegenereerd. Een specifieke en effectieve variant, privileged-context zelf-distillatie (OPSD), conditioneert de "leraar" (hetzelfde model) op aanvullende informatie, zoals een volledige oplossingstrace of een hint, terwijl de student moet redeneren vanuit de oorspronkelijke probleemcontext alleen.

Hoewel deze opzet overvloedige supervisie biedt, identificeert het artikel een kritieke foutmodus: predictieve drift. Omdat de leraar beschikt over geprivilegieerde context, kan deze met hoge mate van zekerheid voortzettingen vastleggen die overduidelijk zijn gegeven de oplossing, maar die onzeker of onondersteund blijven vanuit het perspectief van de student zelf. Standaard distillatie-objectieven (bijv. de Kullback–Leibler divergentie) wegen de voorkeuren van de leraar op basis van de eigen waarschijnlijkheidsmassa van de leraar. Bijgevolg is het gradiëntsignaal sterk op staten waar de leraar zeer zeker is, maar de student een lage waarschijnlijkheid toekent (lage overlap). Over vele updates heen leiden deze sterke trekacties ertoe dat de student de zelfverzekerde voortzettingen van de leraar overneemt voordat de eigen informatie van de student deze kan rechtvaardigen. De accumulatie van dergelijke mismatchende updates leidt tot drift, waarbij het model onzekerheid onderdrukt en de out-of-distribution (OOD) redeneerprestaties verslechtert, wat vaak resulteert in "zelfverzekerde overeenstemming over foute antwoorden".

Bestaande benaderingen proberen dit te mitigeren door onbetrouwbare leraar-signalen te filteren of de divergentie-gewichten te hervormen (bijv. door gebruik te maken van Reverse KL of filtermechanismen). De auteurs stellen echter dat deze methoden er niet in slagen de kern van het probleem aan te pakken: hoe ver een vertrouwd signaal de student moet bewegen. Selectie alleen kan de magnitude van de update niet reguleren, en ongecontroleerde beweging accumuleert in drift.

2. Methodologie: GEOSD

De auteurs introduceren GEOSD (Geometric Self-Distillation), een objectief ontworpen om deze drift tegen te gaan door predictief gedrag te behandelen als beweging binnen een specifieke geometrische ruimte van next-token distributies. De methode werkt op twee complementaire manieren:

A. Lokale Druk-attenuatie (Hellinger Loss)

In plaats van standaard KL-divergentie gebruikt GEOSD een Hellinger-loss om de aantrekkingskracht van de leraar bij elke bezochte staat te moduleren.

• Mechanisme: De Hellinger-afstand is afgeleid van de geometrie van de waarschijnlijkheidssimplex onder een wortel-kwadraat embedding ($p \to \sqrt{p}$), waarbij distributies op een hypersfeer liggen. De gradiënt van de Hellinger-loss schaalt elke leraar-voorkeur door de geometrische gemiddelde van de student- en leraar-waarschijnlijkheden ($\sqrt{p_i q_i}$).
• Effect: Als de leraar een hoge waarschijnlijkheid toekent aan een token, maar de student een bijna nul-waarschijnlijkheid toekent, vervalt de overlap-term en wordt de gradiënt-trekkracht gedempt. Omgekeerd worden gedeelde voorkeuren sterk overgedragen. Dit zorgt ervoor dat supervisie geleidelijk wordt versterkt naarmate de student en de leraar in overeenstemming komen, in plaats van de student te dwingen de zelfverzekerde maar onondersteunde voortzettingen direct over te nemen.

B. Globale Driftcontrole (Fisher–Rao Proximale Term)

Hoewel overlap-weging de aantrekkingskracht bij individuele staten begrenst, voorkomt het niet dat kleine, redelijke updates uitmonden in een grote globale drift over de training heen.

• Mechanisme: GEOSD voegt een proximale term toe die de afstand tussen de huidige student-distributie en een recente checkpoint ($\theta_{ckpt}$) bestraft. Deze afstand wordt gemeten met de Fisher–Rao afstand (de booglengte op de hypersfeer), die de geaccumuleerde verplaatsing in de geometrie van predictieve distributies volgt.
• Effect: Dit fungeert als een "bewegende vertrouwensregio", die abrupte verplaatsingen dempt terwijl het constante beweging toestaat wanneer het leraar-signaal consistent is. In tegenstelling tot het verankeren aan het initiële model (wat adaptatie zou verbieden), staat het verankeren aan een recent checkpoint het model toe om te adapteren terwijl het wordt voorkomen dat het te ver afwijkt van zijn recente, stabiele gedrag.

C. Natural-Gradient Update

Zowel de Hellinger-loss als de Fisher–Rao proximale term zijn afstanden in dezelfde geometrie van next-token distributies. Om dit objectief te optimaliseren, gebruikt GEOSD een natural-gradient update.

• Mechanisme: In plaats van parameters in de Euclidische ruimte te updaten, wordt de update gepreconditioneerd door de Fisher Informatie Matrix (geëxtrapoleerd via K-FAC voor efficiëntie).
• Effect: Dit zorgt ervoor dat elke stap de kleinste verandering aan de voorspellingen van de student maakt die het objectief verbetert, in plaats van een stap van vaste grootte in de parameterruimte te maken. De proximale term fungeert effectief als een "checkpoint terugtrekking" in deze geometrie, waarbij het beleid terug naar het checkpoint wordt geleid op basis van de gedragsmatige salientie van de parameterrichtingen.

3. Belangrijkste Bijdragen

1. GEOSD Objectief: Een zelf-distillatie framework dat leraar-voorkeuren weegt op basis van leraar-student overlap (via Hellinger loss) en predictieve drift van een recent checkpoint bestraft (via Fisher–Rao afstand), geoptimaliseerd via natural gradients in een verenigde geometrie.
2. Empirische Verbetering: GEOSD behoudt de in-distribution (ID) winsten van zelf-distillatie terwijl het de OOD-redeneerprestaties aanzienlijk verbetert. Over drie modelfamilies (Qwen3, Olmo, DeepSeek-R1) en vijf schalen (1.7B tot 32B) verbetert het de gemiddelde OOD-nauwkeurigheid met 5.7–8.6 punten ten opzichte van het basemodel, terwijl standaard KL-objectieven de OOD-prestaties doen verslechteren.
3. Mechanistische Analyse: Het artikel karakteriseert waarom standaard matching faalt bij OOD: het concentreert de waarschijnlijkheidsmassa op de topkeuze van de leraar door alternatieven bij hoge-entropie staten "leeg te trekken", wat leidt tot zelfverzekerde overeenstemming over foute antwoorden. GEOSD behoudt de dominantie van de topkeuze terwijl alternatieve voortzettingen "bereikbaar" blijven, wat voortijdige ineenstorting voorkomt.

4. Experimentele Resultaten

De auteurs evalueerden GEOSD op wiskundige redeneerbenchmarks (DAPO-Math voor training; AIME 2024/2025, AMC 2023, MATH-500 voor OOD-testen).

• Prestaties: GEOSD behaalde de beste OOD-nauwkeurigheid over alle geteste benchmarks en modelfamilies. Bijvoorbeeld, op Qwen3-8B verbeterde GEOSD de OOD-nauwkeurigheid met +8.6 punten ten opzichte van het basemodel, vergeleken met een degradatie van -8.1 punt voor Forward KL.
• Robuustheid: De winsten hielden stand bij variërende hoeveelheden geprivilegieerde context (van korte prefixes tot volledige oplossingen) en over verschillende model-schalen (1.7B tot 32B).
• Ablatie: Het verwijderen van ofwel de Hellinger overlap-weging, de Fisher–Rao proximale term, of de natural-gradient update resulteerde in significante dalingen in OOD-prestaties, wat bevestigt dat alle drie de componenten noodzakelijk zijn.
• Vals Consensus: Analyse toonde aan dat standaardmethoden (zoals Forward KL) hoge percentages "vals consensus" vertonen (samples die sterk overeenstemmen over foute antwoorden) bij hoge-entropie staten. GEOSD verminderde dit percentage zelfs tot onder dat van het basemodel, waardoor informatieve zelfconsistentie behouden bleef.

5. Betekenis en Claims

Het artikel claimt dat de degradatie van OOD-redeneren in privileged-context zelf-distillatie niet louter het gevolg is van het gebruik van onbetrouwbare signalen, maar een consequentie is van hoe die signalen het model bewegen. Standaard objectieven laten de zelfverzekerdheid van de leraar de student voortijdig in een smalle modus dwingen door de waarschijnlijkheidsmassa van levensvatbare alternatieven leeg te trekken.

GEOSD demonstreert dat effectieve redeneer-distillatie niet alleen afhangt van welke leraar-signalen men moet vertrouwen, maar ook van hoe ver die signalen de student mogen bewegen. Door deze beweging geometrisch te reguleren — het dempen van de aantrekkingskracht waar de student onvoldoende ondersteuning heeft en het verankeren van het traject aan een recente stabiele staat — bereikt GEOSD een superieur evenwicht tussen in-distribution prestaties en out-of-distribution generalisatie. De auteurs concluderen dat het beheersen van predictieve drift essentieel is voor het opschalen van zelf-distillatie in redeneertaken.

Genoteerde Beperkingen: De auteurs erkennen dat hun bevindingen momenteel alleen zijn vastgesteld voor wiskundig redeneren. De toepasbaarheid van overlap-gewogen transfer en Fisher–Rao driftcontrole op andere domeinen (bijv. code, wetenschappelijk redeneren) waar hoge-entropie staten mogelijk andere structuren hebben, blijft ongetest. Daarnaast gaat de methode ervan uit dat het basisbeleid breed competent is; voor zeer zwakke basemodellen zou het mechanisme nuttige signalen die grote gedragsveranderingen vereisen, kunnen onderdrukken.