Reinforcement-aware Knowledge Distillation for LLM Reasoning

Technische Samenvatting: Reinforcement-aware Knowledge Distillation voor LLM Reasoning (RLAD)

1. Probleemstelling

Large Language Models (LLMs) hebben significante redeneercapaciteiten bereikt door middel van Reinforcement Learning (RL) post-training, met name in taken met lange chain-of-thought (CoT). De hoge inferentiekosten van deze grote modellen vereisen echter distillatie naar kleinere studentmodellen. Bestaande Knowledge Distillation (KD) methoden kampen met kritieke beperkingen wanneer ze worden toegepast op RL post-training:

• Distributie-mismatch: Standaard KD-methoden (offline, off-policy of on-policy) vertrouwen vaak op vaste teacher-traces of statische KL-divergentie regularisatie. In RL evolueert het beleid van de student continu, waardoor de rollout-distributie verschuift. Statische distillatie-doelen worden suboptimaal of instabiel naarmate de student zich verwijdert van de initiële distributie van de teacher.
• Objectieve Interferentie: Methoden die KD integreren in de RL-loop (bijv. KDRL) voegen doorgaans een KL-divergentie term toe tussen de student en de teacher als een regularisatie-straf. Dit creëert een multi-objective optimalisatieprobleem waarbij het doel van het maximaliseren van de beloning concurreert met het imitatie-doel. Deze competitie vereist zorgvuldige hyperparameter-tuning (loss balancing) en kan leiden tot instabiele trainingsdynamiek, aange omdat de supervisie van de teacher mogelijk niet overeenkomt met de huidige advantage-inschattingen van de student.

2. Methodologie: RL-Aware Distillation (RLAD)

De auteurs stellen RLAD voor, een framework dat is ontworpen om selectieve imitatie uit te voeren tijdens RL post-training. Het kernprincipe is dat de student alleen door de teacher geleid moet worden wanneer die begeleiding de huidige beleidsupdate onder het RL-doel verbetert.

Kerncomponent: Trust Region Ratio Distillation (TRRD)
In plaats van een zelfstandige KL-regularisator te gebruiken, introduceert RLAD TRRD, die de standaard importance sampling ratio in Group Relative Policy Optimization (GRPO) vervangt door een teacher-student mengeling.

• Mechanisme: TRRD definieert een nieuwe likelihood ratio $r^{TRRD}_{i,t}$ voor de $t$-de token van de $i$-de respons:
  $$r^{TRRD}_{i,t}(\theta_S) = \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_S,old}(y_t | x, y_{<t})} \right)^\alpha \left( \frac{\pi_{\theta_S}(y_t | x, y_{<t})}{\pi_{\theta_T}(y_t | x, y_{<t})} \right)^{1-\alpha}$$
  Hierbij is $\pi_{\theta_S}$ het huidige studentbeleid, $\pi_{\theta_S,old}$ het vorige studentbeleid, en $\pi_{\theta_T}$ het teacherbeleid. $\alpha \in [0, 1]$ is een mengcoëfficiënt.
• Interpretatie van de Trust Region: Deze ratio kan worden gezien als een PPO-stijl update waarbij de trust region wordt verankerd aan een mengeling van de oude student en de teacher ($r_{\pi_{mix}}$), in plaats van alleen de oude student.
• Advantage-bewuste Selectie: De update wordt gewogen door de group-relative advantage $\hat{A}_{i,t}$.
  • Positieve Advantage: Als de advantage positief is, wordt de student aangemoedigd de waarschijnlijkheid van de token te verhogen. Als de teacher ook een hoge waarschijnlijkheid toekent, is de mengeling-anker hoger, wat de student toestaat de waarschijnlijkheid meer te verhogen voordat de clipping bound wordt bereikt (wat grotere updates toestaat). Als de teacher een lage waarschijnlijkheid toekent, is de anker lager, wat de omvang van de update beperkt.
  • Negatieve Advantage: Als de advantage negatief is, wordt de student aangemoedigd de waarschijnlijkheid te verlagen. Als de teacher een hoge waarschijnlijkheid toekent, voorkomt de clipping bound dat de student de token agressief onderdrukt, waardoor teacher-gealigneerde gedragingen behouden blijven.
  • Bijna-nul Advantage: De invloed van de teacher wordt verwaarloosbaar, wat interferentie voorkomt wanneer het beloningssignaal niet informatief is.

Dit ontwerp balanceert impliciet exploratie, exploitatie en imitatie door ervoor te zorgen dat de begeleiding van de teacher wordt gemoduleerd door het beloningssignaal (advantage) en binnen een trust region wordt geconstrained.

3. Belangrijkste Bijdragen

1. RLAD Framework: Een nieuw distillatie-framework voor RL post-training dat teacher-begeleiding integreert via Trust Region Ratio Distillation (TRRD), waarbij de statische KL-regularisatie wordt vervangen door een geclipt, advantage-gewogen likelihood-ratio objectief.
2. Theoretische Analyse: Een analyse die aantoont dat TRRD reward-geconditioneerde, selectieve imitatie mogelijk maakt. Het behoudt trust-region-gebonden updates en balanceert natuurlijk de afweging tussen het volgen van de teacher en het optimaliseren voor beloningen zonder dat delicate loss balancing nodig is.
3. Empirische Validatie: Uitgebreide experimenten op logisch redeneren (K&K Logistics) en lang-context wiskundig redeneren (AIME, AMC, MATH) tonen consistente verbeteringen aan ten opzichte van baselines.

4. Experimentele Resultaten

De auteurs evalueerden RLAD op logisch redeneren en complexe wiskundige redeneertaken met studentmodellen variërend van 0.6B tot 8B parameters en teachers tot 32B.

• Logisch Redeneren (K&K Logistics):

  • Met een Qwen3-8B teacher verbeterde RLAD de gemiddelde accuratesse aanzienlijk voor Qwen3-0.6B (van 0.76 naar 0.94 bij 8K context) en Qwen3-1.7B (van 0.95 naar 0.99 bij 8K context) vergeleken met standaard GRPO.
  • De winst was het meest uitgesproken op de moeilijkste subsets (PPL7 en PPL8), waar RLAD een duidelijk voordeel behield ten opzichte van KDRL.
  • De trainingsdynamiek toonde snellere convergentie en hogere uiteindelijke beloningen vergeleken met zowel GRPO als KDRL.

• Wiskundig Redeneren:

  • Base Modellen (R1-Zero stijl): Op Qwen3-8B-Base verbeterde RLAD de gemiddelde prestaties met +5.5 punten over GRPO op lang-context (30K) taken.
  • Post-getrainde Modellen: RLAD presteerde beter dan expliciete-KL distillatie (KDRL) en offline distillatie.
  • Specifieke Metrieken: RLAD vertoonde grotere winsten op Pass@1 (Mean@32) vergeleken met KDRL, dat de neiging had om grotere relatieve winsten te laten zien op Pass@32. De auteurs interpreteren dit als het feit dat RLAD succesvol de reward-aligned beleidsverbetering (exploitatie) aanstuurt, terwijl de winsten van KDRL meer werden gedreven door pure imitatie van diverse teacher-oplossingen.
  • Stabiliteit: RLAD vertoonde stabielere trainingsdynamiek met minder oscillaties vergeleken met KDRL, dat leed onder problemen met objectieve koppeling.

• Robuustheid en Efficiëntie:

  • Zwakke Teacher: Wanneer de teacher zwakker was dan de student, bleef RLAD dicht bij de GRPO-prestaties, terwijl KDRL significanter degradeerde.
  • Hyperparameters: RLAD bleek grotendeels ongevoelig voor de mengcoëfficiënt $\alpha$ (vastgesteld op 0.5), in tegen tegenstelling tot KDRL dat gevoelig is voor de KL-gewicht $\lambda$.
  • Efficiëntie: Het integreren van RLAD voegde slechts ~12% extra trainingslatentie toe vergeleken met standaard GRPO, aangezien de teacher alleen wordt gebruikt voor de berekening van de log-waarschijnlijkheid, niet voor de generatie van trajecten.

5. Betekenis en Claims

Het artikel claimt dat RLAD de fundamentele mismatch aanpakt tussen statische teacher-supervisie en dynamische RL beleidsupdates. Door teacher-begeleiding in de trust-region update te verwerken, vermijdt RLAD de "breekbare afstemming" die vereist is voor multi-objective KL-regularisatie.

De auteurs stellen dat RLAD de student in staat stelt om naar gedragingen met een hoge advantage te bewegen binnen een anker-distributie en daarbij opportunistisch gebruik te maken van teacher-begeleiding. Dit resulteert in:

• Superieure Prestaties: Consistente uitmuntendheid ten opzichte van GRPO, KDRL en offline distillatie, met name op uitdagende benchmarks.
• Stabiliteit: Stabielere trainingsdynamiek en betere convergentie.
• Echte RL Optimalisatie: In tegen tegenstelling tot methoden die zwaar leunen op imitatie (wat Pass@32 kan verhogen maar niet Pass@1), geeft RLAD prioriteit aan reward-aligned gedrag, wat leidt tot sterkere winsten in de meest waarschijnlijke correcte antwoorden van het model (Pass@1).

Het artikel concludeert dat hoewel de aanpak vertrouwt op open-weight, goed gekalibreerde teacher-modellen voor toegang tot logits, het een belangrijke stap voorwaarts is om distillatie compatibel te maken met de dynamische aard van RL post-training.