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

1. Il Problema

L'apprendimento per rinforzo (RL) ha dimostrato di migliorare significativamente le capacità di ragionamento dei modelli linguistici (LLM), specialmente in domini complessi come la matematica e la logica. Tuttavia, gli approcci esistenti (come RLVR - Reinforcement Learning with Verifiable Rewards) dipendono fortemente da reward verificabili basati su risposte a terra (ground truth). Questo limita l'applicabilità in scenari dove:

• I dati etichettati sono scarsi o costosi da ottenere.
• È necessario utilizzare grandi volumi di dati non etichettati per la distillazione della conoscenza.
• Si vuole democratizzare le capacità di ragionamento avanzate su modelli più piccoli ed efficienti, evitando la semplice imitazione superficiale tipica del Supervised Fine-Tuning (SFT).

La sfida principale è quindi creare un sistema di reward efficace per l'RL che non richieda ground truth, utilizzando dati non etichettati, mantenendo al contempo l'efficienza computazionale.

2. Metodologia

Gli autori propongono un framework di Knowledge Distillation basato su RL che trasforma il paradigma "LLM-as-a-Judge" da un valutatore statico a un generatore attivo di reward.

Configurazione del Problema

• Studente: Un modello linguistico più piccolo (es. Galactica-125M, CodeGen-350M) che genera traiettorie di ragionamento (Chain-of-Thought, CoT).
• Giudice: Un LLM più capace (es. Qwen3-8B) che valuta le risposte dello studente.
• Dati: Utilizzo di dati etichettati (per reward verificabili) e grandi volumi di dati non etichettati (per reward basati sul giudice).

Design del Reward

Il framework combina tre tipi di segnali di reward in un obiettivo di apprendimento convesso:

1. Reward Verificabile ($L^{VR}_{PPO}$): Basato su dati etichettati. Se la risposta estratta corrisponde alla ground truth, il reward è 1; se è numerica ma errata, 0.1; altrimenti 0.
2. Reward LLM-as-a-Judge ($L^{YoN}_{PPO}$): Questo è il contributo centrale. Il giudice LLM valuta la soluzione dello studente con un prompt che richiede una risposta binaria ("Sì" o "No").
   • Invece di generare una risposta completa, il giudice produce solo un token.
   • Il reward è calcolato direttamente dalle logits (probabilità) dei token "Sì" e "No": $s = \sigma(u) \cdot I\{\sigma(u) \ge a\}$, dove $u$ è la differenza logaritmica delle probabilità.
   • Questo approccio è efficiente (single-token output) e fornisce un segnale continuo e sfumato, non solo binario.
3. Reward da Modello di Reranking ($L^{Rerank}_{PPO}$): Utilizza un modello di reranking pre-addestrato per fornire una probabilità di correttezza come reward continuo.

Algoritmo di Ottimizzazione

• Viene utilizzato PPO (Proximal Policy Optimization) per la maggior parte degli esperimenti, con un warm-up SFT iniziale.
• Viene anche valutato GRPO (Group Relative Policy Optimization), che mostra prestazioni superiori ma con un costo computazionale maggiore (richiede la generazione di un gruppo di risposte per domanda).
• L'obiettivo finale è una combinazione pesata: $L_{overall} = \lambda L^{VR}_{PPO} + \mu L^{YoN}_{PPO} + \rho L^{Rerank}_{PPO}$.

3. Contributi Chiave

1. Distillazione senza etichette (Label-free): Dimostrano che è possibile addestrare modelli piccoli su grandi quantità di dati non etichettati utilizzando un LLM come giudice, eliminando la dipendenza esclusiva dalle risposte a terra.
2. Efficienza del Reward: L'uso di un output a singolo token ("Sì"/"No") per il giudice rende il calcolo del reward estremamente efficiente rispetto alla generazione di intere traiettorie di ragionamento.
3. Segnali di Addestramento Efficaci: Sostengono che i valutatori basati su LLM possono produrre segnali di addestramento robusti per il fine-tuning RL, superando i limiti della distillazione supervisionata tradizionale che spesso fallisce nel catturare i processi di ragionamento sottostanti.
4. Scalabilità: Il metodo funziona efficacemente sia su modelli piccoli (125M, 350M) che su modelli più grandi (6.7B), fornendo miglioramenti anche quando le prestazioni di base sono già elevate.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su dataset di ragionamento matematico (GSM8K, SVAMP, GSM-Plus, GSM-Symbolic).

• Modelli Piccoli (Galactica-125M):
  • Rispetto all'uso di soli reward verificabili (RLVR), l'aggiunta del reward del giudice su dati non etichettati ha portato a un miglioramento di 5-10 punti di accuratezza su GSM8K (da ~30% a ~40-41%).
  • Su dataset fuori dominio (Out-of-Distribution) come GSM-Symbolic e SVAMP, i miglioramenti sono stati ancora più significativi, dimostrando una maggiore robustezza e capacità di generalizzazione.
  • L'uso combinato di dati etichettati e non etichettati ha permesso di superare le prestazioni di un esperimento "oracle" che usava solo reward verificabili su dati non etichettati.
• Modelli Grandi (Galactica-6.7B):
  • Anche su modelli più grandi, dove le prestazioni di base sono già saturate, l'approccio proposto ha fornito un miglioramento di circa 1.5 punti (da 69.07% a 70.43% su GSM8K), confermando la scalabilità del metodo.
• Confronto PPO vs GRPO:
  • GRPO ha mostrato guadagni di prestazioni leggermente superiori rispetto a PPO, ma con un overhead computazionale più elevato.
• Stabilità: È stato notato un fenomeno di "reward hacking" (il modello genera risposte brevi o senza codice per ingannare il giudice) quando si usano solo dati non etichettati ($\lambda=0$). Questo è stato mitigato combinando i dati etichettati e applicando vincoli sulla lunghezza del codice generato.

5. Significato e Implicazioni

Questo lavoro è significativo perché:

• Democratizza l'RL: Rende possibile l'addestramento per rinforzo di modelli di ragionamento avanzati senza la necessità di costosi dataset etichettati o annotazioni umane.
• Efficienza Computazionale: Trasforma il "LLM-as-a-Judge", tradizionalmente usato per la valutazione post-hoc, in un componente attivo e leggero del ciclo di addestramento.
• Generalizzazione: Dimostra che i segnali di reward derivati da modelli LLM possono guidare l'apprendimento di capacità di ragionamento profonde, superando i limiti della semplice imitazione superficiale tipica dello SFT.
• Futuro della Ricerca: Apre la strada all'uso di dati non etichettati massicci per la distillazione di capacità cognitive complesse, suggerendo che la qualità del feedback (anche da modelli AI) è più critica della disponibilità di ground truth per certi compiti di ragionamento.

In sintesi, il paper propone un framework scalabile ed efficiente che colma il divario tra la necessità di dati non etichettati e la capacità di addestrare modelli di ragionamento robusti tramite RL, utilizzando LLM come giudici dinamici.