Technische Samenvatting: Privileged Hidden Flow voor On-Policy Self-Distillation

Probleemstelling

On-policy self-distillation (OPSD) is een techniek voor het trainen van redeneermodellen waarbij een studentmodel zijn eigen rollouts samplet en wordt gesuperviseerd door een "privileged" (bevoorrechte) teacher. Deze teacher heeft dezelfde modelarchitectuur, maar is geconditioneerd op aanvullende informatie, zoals een geverifieerde referentie-oplossing, terwijl de student alleen de probleemprompt ziet.

Bestaande OPSD-doelstellingen richten zich uitsluitend op het matchen van de outputverdeling (de volgende-token-kansen) tussen de student en de privileged teacher. Hoewel dit de uiteindelijke voorspellingen afstemt, faalt het in het superviseren van de interne berekening die deze voorspellingen voortbrengt. Specifiek beïnvloedt de geprivilegieerde context de intermediaire verborgen toestanden (hidden states) en het traject van deze toestanden door de lagen van het netwerk. Huidige methoden beperken zich tot de uiteindelijke output-head, waardoor de "hidden flow"—hoe de interne representaties van het model stap voor stap evolueren tijdens de rollout—ongestuurd blijft. Het artikel stelt dat het distilleren van dit interne proces de redeneerprestaties kan verbeteren zonder de fundamentele on-policy trainingsopzet te wijzigen.

Methodologie: Privileged Hidden Flow (PHF)

De auteurs stellen Privileged Hidden Flow (PHF) voor, een methode die standaard OPSD aanvult door de beweging van verborgen toestanden te distilleren in plaats van de toestanden zelf.

Kernmechanisme

In plaats van de verborgen vector van de student $h^S_{\ell,t}$ te dwingen overeen te komen met de vector van de teacher $h^T_{\ell,t}$ op dezelfde tokenpositie (pointwise matching), matcht PHF de transities tussen tokenposities.

1. Definitie van Transitie: Voor een set geselecteerde gegenereerde posities $W = \{w_1, \dots, w_m\}$ is de lokale verborgen transitie gedefinieerd als het verschil tussen naburige posities: $\Delta h_{\ell,i} = h_{\ell, w_{i+1}} - h_{\ell, w_i}$.
2. Normalisatie: Deze transities worden genormaliseerd naar eenheidrichtingen om de focus te leggen op de geometrie van de beweging in plaats van de magnitude.
3. Verliescomponenten:
   • Direction Loss ($L_{dir}$): Stemt de richting van studenttransities af met de transities van de teacher op elke laag.
   • Geometry Loss ($L_{geo}$): Matcht de Gram-matrix van de transitierichtingen. Dit legt de relationele structuur van het traject vast (bijv. welke stappen parallel, orthogonaal of tegenovergesteld zijn) in plaats van absolute coördinaten.
   • Neighboring-Layer Relation ($L_{adj}$): Berekent de relatie tussen transitierichtingen in aangrenzende lagen ($\ell$ en $\ell+1$) om ervoor te zorgen dat de transformatie van beweging over de diepte consistent is met de teacher.

Structurele Invarianties

Een belangrijke theoretische bijdrage van PHF is de invariantie-eigenschap, die het probleem van "representational drift" aanpakt, waarbij verborgen coördinaten verschuiven tijdens de training:

• Offset Invariantie: De loss verandert niet als een constante offset aan het gehele traject van de student of de teacher wordt toegevoegd.
• Schaalinvariantie: De loss is invariant voor positieve scalaire vermenigvuldiging van transities (vóór normalisatie).
• Orthogonale Invariantie: De geometrie-term is invariant onder onafhankelijke orthogonale transformaties van de student- en teacher-transitiesets.

Deze eigenschappen maken het mogelijk om de vorm en richting van het redeneerproces te superviseren zonder dat de student en de teacher exact dezelfde punten in de hoog-dimensionale verborgen ruimte hoeven te bezetten.

Trainingsrecept

• Teacher Bron: De privileged teacher is een Exponential Moving Average (EMA) van de online student, geconditioneerd op de referentie-oplossing.
• Aggregatie: De methode aggregeert de transitie-loss over alle lagen in plaats van een subset te selecteren, en combineert de lokale transitie-loss met de naburige-laag-relatie.
• Doelstelling: De uiteindelijke loss is $L_{PHF} = L_{OPSD} + \alpha L_{Flow}$, waarbij de coëfficiënt $\alpha$ vast staat op 0.05 over alle schalen.
• Restricties: De methode gebruikt een vaste window van 128 gegenereerde posities om het geheugenbeheer te beheersen en introduceert geen reward-modellen, correctness routing of extra rollouts.

Belangrijkste Bijdragen

1. Privileged Transition-Transport Doelstelling: Introduceert een nieuwe loss die genormaliseerde token-naar-token verborgen transities en hun traject-geometrie matcht, waarmee verder wordt gegaan dan pointwise state matching.
2. Theoretische Karakterisering: Bewijst dat de transport-loss exacte algebraïsche invarianties bezit ten opzichte van traject-offsets en orthogonale transformaties, wat het onderscheidt van standaard verborgen-toestand matching.
3. All-Layer Aggregatie: Demonstreert dat het middelen van de transitie-loss over alle lagen, gecombineerd met een vaste naburige-laag-relatie, consistente verbeteringen oplevert zonder dat handmatige laagselectie nodig is.
4. Empirische Validatie: Valideert het vaste recept over drie model-schalen (Qwen3-1.7B, 4B, en 8B) onder het officiële OPSD wiskundig redeneerprotocol.

Experimentele Resultaten

De auteurs hebben PHF geëvalueerd op de AIME 2024, AIME 2025, en HMMT 2025 benchmarks met behulp van de Average@12 metriek (gemiddelde correctheid over 12 samples per probleem).

• Prestatiewinst: Onder een vast 100-stappen trainingsschema verbeterde PHF de geaggregeerde Average@12 ten opzichte van de gereproduceerde OPSD baseline:
  • Qwen3-1.7B: +2.2 punten (40.83 $\to$ 43.06)
  • Qwen3-4B: +1.5 punten (61.89 $\to$ 63.39)
  • Qwen3-8B: +1.7 punten (64.26 $\to$ 65.94)
• Consistentie: De winst was positief over alle schalen en bijna alle individuele benchmark-cellen.
• Ablatieonderzoeken:
  • Direction vs. Geometry: Noch de direction-only, noch de geometry-only varianten matchten het volledige PHF-recept over alle schalen, wat aangeeft dat beide componenten noodzakelijk zijn.
  • Laagselectie: Een "selected layers" controle (het matchen van een kleine, handmatig gekozen subset) presteerde slechter dan de all-layer aggregatie, wat de beslissing ondersteunt om over de diepte te middelen.
  • Pointwise Matching: Het artikel merkt op dat pointwise verborgen-toestand matching een competitieve boundary control is, maar minder robuust is tegen representational drift dan de transitie-gebaseerde aanpak.

Betekenis en Claims

Het artikel claimt dat PHF een bescheiden maar betekenisvolle verbetering biedt aan het OPSD-framework door geprivilegieerdeerde informatie structureler te gebruiken.

• Proces-niveau Supervisie: De primaire betekenis is het aantonen dat geprivilegieerde referenties een specifieke "verborgen proces" induceren (hoe het model door zijn eigen rollout beweegt) dat gedistilleerd kan worden. PHF vangt dit proces zonder dat de student de absolute verborgen coördinaten van de teacher hoeft te matchen.
• Deployment Compatibiliteit: In tegenstelling tot methoden die reward-modellen of externe verifiers vereisen tijdens de testfase, behoudt PHF de OPSD-deployment setting: het geëvalueerde model is de online student, en er is geen geprivilegieerde informatie nodig tijdens inferentie.
• Efficiëntie: De methode voegt slechts een scalaire loss-coëfficiënt en een verborgen-toestand opslagwindow toe, waardoor de computationele overhead van extra rollouts of complexe routing-mechanismen wordt vermeden.
• Omvang van Claims: De auteurs zijn expliciet over het feit dat hun bewijs beperkt is tot een vast recept, een specifieke modelfamilie (Qwen3) en wiskundige redeneertaken. Zij beweren niet dat hidden-flow matching elk aspect van redeneren (zoals bewijskwaliteit of kalibratie) verbetert of dat het een universele oplossing is, maar dat het een consistente, positieve geaggregeerde verschuiving biedt wanneer het wordt toegepast op on-policy self-distillation.