Reinforcement Learning-based Knowledge Distillation with LLM-as-a-Judge

Probleemstelling

Reinforcement Learning (RL) heeft bewezen de redeneercapaciteiten van Large Language Models (LLMs) aanzienlijk te verbeteren, maar bestaande methoden (zoals RL met verifieerbare beloningen, RLVR) zijn sterk afhankelijk van ground truth-labels (bijv. het juiste antwoord op een wiskundige vraag). Dit beperkt de toepasbaarheid in scenario's waar gelabelde data schaars, duur of niet beschikbaar is.

Daarnaast is er een kritieke behoefte aan het "democratiseren" van geavanceerde redeneervermogens door deze te distilleren naar kleinere, efficiëntere modellen. Traditionele kennisdistillatie via Supervised Fine-Tuning (SFT) faalt vaak omdat het slechts oppervlakkige token-nabootsing bevordert en geen onderliggende redeneerprocessen overneemt, wat leidt tot slechte generalisatie.

De kernvraag is: Hoe kunnen we kleinere modellen trainen met RL om redeneercapaciteiten te leren van grotere modellen, zonder afhankelijk te zijn van ground truth-labels voor de trainingsdata?

Methodologie

Het paper stelt een nieuw RL-framework voor dat LLM-as-a-Judge gebruikt als een dynamische, label-vrije beloningsbron. Het framework combineert verifieerbare beloningen (voor gelabelde data) met beoordelingen van een sterke "rechter"-LLM (voor ongelabelde data).

1. Architectuur en Training

• Student Model: Een kleiner taalmodel (bijv. 125M of 350M parameters) dat redeneertrajecten (Chain-of-Thought, CoT) genereert.
• Judge Model: Een krachtig LLM (in dit geval Qwen3-8B) dat de output van de student evalueert.
• Warm-up: Het studentmodel wordt eerst getraind via SFT op gelabelde data om een basisredeneervermogen te ontwikkelen.
• RL Fine-tuning: Het model wordt verder getraind met PPO (Proximal Policy Optimization). De beloningssignaal is een combinatie van drie componenten:
  1. Verifieerbare Beloning ($L^{VR}_{PPO}$): Voor gelabelde data. Als de gegenereerde code (Program-aided Language Models, PoT) correct compileert en het juiste antwoord geeft, is de beloning 1. Anders 0 of 0.1.
  2. LLM-as-a-Judge Beloning ($L^{YoN}_{PPO}$): Voor ongelabelde data. De judge LLM krijgt de vraag en het antwoord van de student en moet alleen "Yes" of "No" antwoorden. In plaats van een volledige generatie te wachten, wordt de beloning berekend op basis van de logits (kansverdeling) voor de tokens "Yes" en "No".
     • De beloning wordt berekend als een continue schaalwaarde: $s = \sigma(u) \cdot I_{\{s \ge a\}}$, waarbij $u$ het verschil is in log-probabiliteit tussen "Yes" en "No". Dit maakt de berekening zeer efficiënt (één token stap).
  3. Reranking Model Beloning ($L^{Rerank}_{PPO}$): Een vooraf getraind reranking-model dat als tweede judge fungeert om de diversiteit van de signalen te vergroten.

De totale loss functie is een convexe combinatie:
$$L_{overall} = \lambda \cdot L^{VR}_{PPO} + \mu \cdot L^{YoN}_{PPO} + \rho \cdot L^{Rerank}_{PPO}$$
waarbij $\lambda + \mu + \rho = 1$.

2. Beloningsefficiëntie

Een cruciaal innovatief aspect is dat de judge LLM niet hoeft te genereren of te redeneren om een score te geven; het gebruikt alleen de waarschijnlijkheid van het eerste token ("Yes" of "No"). Dit elimineert de hoge rekentijd die normaal nodig is voor het genereren van volledige beoordelingen.

Belangrijkste Bijdragen

1. Label-vrije RL Distillatie: Het transformeert de "LLM-as-a-Judge" van een statische evaluatiemethode naar een actieve, real-time beloningsgenerator voor RL, waardoor training op enorme hoeveelheden ongelabelde data mogelijk wordt.
2. Efficiënte Beloningsberekening: Door gebruik te maken van single-token output (Yes/No) en logits, wordt de overhead van het gebruik van een groot LLM als judge geminimaliseerd.
3. Superieure Prestaties op Kleine Modellen: Het framework toont aan dat kleine modellen (125M/350M) significant kunnen profiteren van ongelabelde data wanneer ze worden geleid door een sterke judge, zelfs zonder ground truth voor die specifieke data.
4. Robuustheid: Het framework verbetert niet alleen de prestaties op de trainingsdomeinen, maar ook op out-of-distribution benchmarks, wat wijst op een beter geïnternaliseerd redeneervermogen.

Resultaten

De experimenten zijn uitgevoerd op wiskundige redeneerbenchmarks (GSM8K, SVAMP, GSM-Plus, GSM-Symbolic) met modellen van 125M, 350M en 6.7B parameters.

• Kleine Modellen (Galactica-125M & CodeGen-350M):
  • Alleen RL met verifieerbare beloningen (op gelabelde data) gaf een verbetering van ongeveer 5-6 punten ten opzichte van SFT.
  • Door de toevoeging van ongelabelde data met de LLM-as-a-Judge beloning, steeg de nauwkeurigheid op GSM8K van ~30% naar 40.71%.
  • Dit is een aanzienlijke verbetering ten opzichte van alleen RLVR (30.48%) en SFT (25.25%).
  • Op out-of-distribution datasets (zoals GSM-Symbolic en SVAMP) waren de winsten nog groter, wat aantoont dat het model robuuster is geworden.
• Groot Model (Galactica-6.7B):
  • Zelfs voor grotere modellen, waar de prestaties vaak verzadigd zijn, leverde de combinatie van beloningen een verbetering op van 69.07% naar 70.43% op GSM8K (en tot 71.95% met meerderheidsstemming).
• Ablatie Studies:
  • Het gebruik van alleen ongelabelde data met de "Noisy" judge-beloning presteerde beter dan het gebruik van beperkte gelabelde data met verifieerbare beloningen.
  • De combinatie van alle drie de loss-termen ($\lambda=0.5, \mu=0.25, \rho=0.25$) gaf de beste resultaten.
  • GRPO vs. PPO: GRPO (Group Relative Policy Optimization) leverde iets betere resultaten dan PPO, maar PPO werd gekozen voor de hoofdstudie vanwege de lagere rekentijd (GRPO vereist het genereren van meerdere antwoorden per vraag).

Betekenis en Conclusie

Dit paper demonstreert dat LLM-based evaluators effectieve trainingsignalen kunnen produceren voor RL-finetuning, zelfs zonder ground truth. Dit opent de deur voor het trainen van krachtige, kleine redeneermodellen op schaalbare hoeveelheden ongelabelde data, wat de afhankelijkheid van dure menselijke annotatie of verifieerbare antwoorden vermindert.

De methode lost het probleem op van "oppervlakkige distillatie" door het studentmodel actief te laten exploreren in de oplossingsruimte via RL, geleid door de "latent" redeneercapaciteiten van een groter model. Dit is een belangrijke stap richting het democratiseren van geavanceerde AI-redenering voor resource-beperkte omgevingen.