Technische Samenvatting: OISD – On-Policy Internal Self-Distillation van Taalmodellen

1. Probleemstelling

Recente versterkingsleer (RL) post-training benaderingen voor Large Language Models (LLMs), met name die welke Group Relative Policy Optimization (GRPO) met verifieerbare beloningen gebruiken, hebben de wiskundige redenering aanzienlijk verbeterd. Deze methoden zijn echter voornamelijk gericht op de einduitvoer. Ze optimaliseren het uiteindelijke activeringsbeleid met behulp van schaarse, uitkomstgerelateerde beloningen (bijvoorbeeld de juistheid van het eindantwoord), terwijl ze grotendeels voorbijgaan aan de voorspellende signalen, redeneerprocessen en interne representaties die in de tussenliggende lagen zijn gecodeerd.

Hoewel Kennisdistillatie (KD) een mechanisme biedt om dichte token-level supervisie te verlenen, lijden bestaande paradigma's aan specifieke beperkingen:

• Off-Policy Distillatie: Creëert een distributiemismatch omdat de student wordt getraind op trajecten gegenereerd door een vaste leraar, en niet op zijn eigen generaties.
• On-Policy Distillatie: Vermindert de mismatch, maar vertrouwt doorgaans op externe leraar-modellen.
• On-Policy Self-Distillatie: Recentere methoden verwijderen externe leraren, maar zijn afhankelijk van extern ingevoerde bevoorrechte informatie (bijvoorbeeld bevoorrechte prompts, grond-waarheid oplossingen of aanvullende feedback) om het model te superviseren.

Er bestaat een gat in het benutten van de eigen interne berekeningen van het model als bron van supervisie, zonder dat externe bevoorrechte data of meerdere activeringsbeleid vereist zijn. Bovendien wijst eerdere research erop dat redeneerrepresentaties zich systematisch ontwikkelen over de lagen heen (van brede kandidaatvoorspellingen in vroege lagen tot geconvergeerde beslissingen in diepe lagen), wat suggereert dat tussenliggende lagen door diepere lagen kunnen worden geleid om de coherentie van het redeneren te verbeteren.

2. Methodologie: OISD Framework

De auteurs stellen OISD (On-Policy Internal Self-Distillation) voor, een framework dat de eindlaag van het model behandelt als een "losgekoppelde interne leraar" en een geselecteerde tussenliggende laag als een "auxiliaire interne student". Dit vindt plaats binnen één enkele forward pass tijdens het rollout- en GRPO-optimalisatieproces.

Kernarchitectuur

• Enkel Activeringsbeleid: De eindlaag ($L$) blijft het enige beleid voor het genereren van rollouts en het optimaliseren van de GRPO-doelstelling.
• Interne Leraar-Student: Tijdens dezelfde forward pass wordt de verborgen toestand van een geselecteerde tussenliggende laag ($l < L$) geprojecteerd naar de tokenruimte met behulp van de unembedding-matrix (volgend op het logit lens framework). Dit creëert een voorlopige "student"-distributie die is afgeleid van dezelfde context als de "leraar" (eindlaag).
• Geen Externe Bevoorrechte Informatie: De supervisiesignalen zijn volledig afgeleid van de eigen interne toestanden van het model op dezelfde on-policy rollout, zonder dat externe grond-waarheid oplossingen of bevoorrechte prompts nodig zijn.

Aligneringsdoelstellingen

OISD maakt gebruik van twee complementaire interne aligneringsmechanismen om voorspellende signalen van de eindlaag naar de tussenliggende laag over te dragen:

1. Logit-Alignering ("Hoe te Denken"):

   • Doel: Overdracht van hoogwaardige redeneergedragingen en voorspellende overtuigingen.
   • Mechanisme: De Jensen–Shannon (JS)-divergentie wordt berekend tussen de met temperatuur geschaalde token-distributies van de tussenliggende laag ($p^l_t$) en de eindlaag ($q^L_t$).
   • Gewicht: Het verlies wordt gewogen met de getekende GRPO-voordeel ($A_t$). Tokens met een hoog voordeel trekken de tussenliggende laag naar de distributie van de eindlaag, terwijl tokens met een laag voordeel deze wegduwen om het versterken van slechte voortzettingen te voorkomen.
   • Formule: $\mathcal{L}^{(l)}_{\text{think}} = \mathbb{E}_t [ A_t \cdot \text{JS}(p^{l,\tau}_t, \text{sg}(q^{L,\tau}_t)) ]$

2. Attention-Alignering ("Waar Naar Kijken"):

   • Doel: Handhaving van consistente attentionpatronen en evidence-routing-gedrag.
   • Mechanisme: De JS-divergentie wordt berekend tussen de attentiongewichten van de tussenliggende laag en de eindlaag over een gesamplede set causale sleutels (combinerend lokale vensters en globale gestraddle posities om het geheugen te beheren).
   • Gewicht: Net als bij logit-alignering wordt het verlies gewogen met het getekende voordeel. Positieve voordelen moedigen de tussenliggende laag aan om naar dezelfde bewijslast te kijken als de eindlaag; negatieve voordelen verminderen de overeenstemming met patronen met lage beloning.
   • Formule: $\mathcal{L}^{(l)}_{\text{attn}} = \mathbb{E}_t [ A_t \cdot \frac{1}{H} \sum_{h=1}^H \text{JS}(\tilde{a}^l_{t,h}, \text{sg}(\tilde{a}^L_{t,h})) ]$

Trainingsdoelstelling

Het totale verlies combineert het standaard GRPO-verlies met de twee interne aligneringsverliezen:
$$\mathcal{L}_{\text{OISD}} = \mathcal{L}_{\text{GRPO}} + \lambda_{\text{think}} \mathcal{L}_{\text{think}} + \lambda_{\text{attn}} \mathcal{L}_{\text{attn}}$$
waarbij $\lambda_{\text{think}}$ en $\lambda_{\text{attn}}$ de sterkte van de distillatie regelen.

3. Belangrijkste Bijdragen

1. Nieuw Paradigma: Introductie van On-Policy Internal Self-Distillation, een redeneer-RL paradigma dat de eigen interne berekeningen van het model benut voor supervisie terwijl één enkel activeringsbeleid wordt gehandhaafd.
2. OISD Framework: Een specifieke implementatie die de eindlaag gebruikt als een losgekoppelde leraar en een tussenliggende laag als student, waarbij on-policy voorspellende signalen worden overgedragen zonder externe bevoorrechte informatie.
3. Dual Aligneringsmechanismen: Ontwikkeling van twee complementaire doelstellingen:
   • Logit-Alignering: Superviseert het redeneerproces ("hoe te denken").
   • Attention-Alignering: Superviseert de evidence-routing ("waar naar kijken").
4. Empirische Validatie: Aantonen van aanzienlijke en consistente verbeteringen ten opzichte van sterke redeneer-RL baselines (inclusief GRPO en BuPO) over vier wiskundige redeneerbenchmarks.

4. Experimentele Resultaten

De auteurs hebben OISD geëvalueerd op Qwen3-4B, Qwen3-8B en OctoThinker modellen over vier benchmarks: AMC23, MATH500, AIME24 en AIME25.

• Prestatiewinst:
  • Op Qwen3-4B behaalde OISD een gemiddelde score van 64,75, wat een prestatie is van 9,67 punten boven GRPO (55,08) en 6,24 punten boven BuPO (58,51).
  • De verbeteringen waren bijzonder significant op AIME-benchmarks (bijvoorbeeld +14,31 punten boven GRPO op AIME24).
  • Op Qwen3-8B behaalde OISD het beste overall gemiddelde van 66,72, presterend beter dan GRPO (64,23) en BuPO (66,36).
• Pass@K Analyse: OISD behaalde consistent hogere Pass@K-rates over alle $K$-waarden in vergelijking met GRPO, wat wijst op een verbetering in de algehele kwaliteit van de gesamplede oplossingsset, en niet alleen in de gemiddelde nauwkeurigheid.
• Ablatiestudies:
  • Zowel Logit- als Attention-alignering droegen positief bij. Logit-alignering toonde de grootste winst op AIME24, terwijl Attention-alignering gemiddeld iets beter presteerde.
  • Het combineren van beide leverde de beste resultaten op, wat hun complementaire aard bevestigt.
• Vergelijking met OPSD: OISD (64,75) presteerde aanzienlijk beter dan OPSD (56,63), een state-of-the-art on-policy self-distillatiemethode die vertrouwt op bevoorrechte grond-waarheid oplossingen. Dit onderstreept de effectiviteit van de puur interne supervisie van OISD.
• Kwalitatieve Analyse:
  • Trainingsdynamiek: OISD toonde een gestage verbetering in beloning en een initiële toename in entropie (wat wijst op bredere exploratie) voordat het stabiliseerde.
  • Attention Visualisatie: Tussenliggende lagen in OISD aligneerden geleidelijk met de attention van de eindlaag naar kritieke bewijstokens (bijvoorbeeld eindpunten, supremum) tijdens het trainen.
  • Logit-Lens Visualisatie: OISD vertoonde een eerdere alignering tussen tussenliggende en eindvoorspellingen op redeneringskritieke tokens in vergelijking met GRPO, wat vaak gefragmenteerde tussenliggende voorspellingen produceerde.

5. Betekenis en Claims

Het paper claimt dat OISD het onderzoek naar on-policy internal self-distillation inleidt. De primaire betekenis hiervan ligt in:

• Eliminatie van Externe Afhankelijkheid: Het biedt dichte interne supervisie zonder dat externe leraren, bevoorrechte prompts of grond-waarheid oplossingen tijdens het trainen vereist zijn.
• Beleidconsistentie: Het behoudt één enkel activeringsbeleid voor zowel rollout-generatie als optimalisatie, waardoor distributiemismatch-problemen van off-policy methoden worden vermeden.
• Versterkt Redeneren: Door tussenliggende representaties te aligneren met de redeneer- en attentionpatronen van de eindlaag, moedigt OISD het model aan om sterkere, meer globaal consistente tussenliggende representaties te leren.

De auteurs merken beperkingen op, waaronder de focus op wiskundig redeneren met verifieerbare antwoorden en het ontbreken van een diepgaande studie naar optimale laagselectie. Zij suggereren dat toekomstig werk zich moet richten op generalisatie naar complexere redeneertaken en multi-layer interne supervisie.