Reinforcement-aware Knowledge Distillation for LLM Reasoning

Riepilogo Tecnico: Distillazione della Conoscenza Consapevole del Reinforcement Learning per il Ragionamento di LLM (RLAD)

1. Definizione del Problema

I Large Language Models (LLM) hanno raggiunto capacità di ragionamento significative attraverso il post-training con Reinforcement Learning (RL), in particolare nei compiti di catena di pensiero (Chain-of-Thought, CoT) lunga. Tuttavia, gli elevati costi di inferenza di questi modelli grandi rendono necessaria la distillazione in modelli student più piccoli. I metodi esistenti di Knowledge Distillation (KD) affrontano limitazioni critiche quando applicati al post-training RL:

• Disallineamento della Distribuzione: I metodi KD standard (offline, off-policy o on-policy) spesso si affidano a tracce del docente fisse o a una regolarizzazione della divergenza KL statica. Nell'RL, la policy dello studente evolve continuamente, causando uno spostamento della distribuzione dei rollout. Gli obiettivi di distillazione statici diventano subottimali o instabili man mano che lo studente si allontana dalla distribuzione iniziale del docente.
• Interferenza degli Obiettivi: I metodi che integrano la KD nel ciclo RL (ad esempio, KDRL) aggiungono tipicamente un termine di divergenza KL tra lo studente e il docente come penalità di regolarizzazione. Questo crea un problema di ottimizzazione multi-obiettivo in cui l'obiettivo di massimizzazione della ricompensa compete con l'obiettivo di imitazione. Questa competizione richiede un attento tuning degli iperparametri (bilanciamento della perdita) e può portare a dinamiche di addestramento instabili, poiché la supervisione del docente potrebbe non allinearsi con le stime di vantaggio attuali dello studente.

2. Metodologia: Distillazione Consapevole dell'RL (RLAD)

Gli autori propongono RLAD, un framework progettato per eseguire un'imitazione selettiva durante il post-training RL. Il principio cardine è che lo studente debba essere guidato dal docente solo quando tale guida migliora l'attuale aggiornamento della policy sotto l'obiettivo RL.

Componente Chiave: Trust Region Ratio Distillation (TRRD)
Invece di utilizzare un regolarizzatore KL autonomo, RLAD introduce TRRD, che sostituisce il normale rapporto di campionamento per importanza in Group Relative Policy Optimization (GRPO) con una miscela docente-studente.

• Meccanismo: TRRD definisce un nuovo rapporto di verosimiglianza $r^{TRRD}_{i,t}$ per il $t$-esimo token della $i$-esima risposta:
  $$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}$$
  Qui, $\pi_{\theta_S}$ è l'attuale policy dello studente, $\pi_{\theta_S,old}$ è la precedente policy dello studente e $\pi_{\theta_T}$ è la policy del docente. $\alpha \in [0, 1]$ è un coefficiente di miscelazione.
• Interpretazione della Trust Region: Questo rapporto può essere visto come un aggiornamento in stile PPO dove la trust region è ancorata a una miscela della vecchia student policy e del docente ($r_{\pi_{mix}}$), piuttosto che solo alla vecchia student policy.
• Selezione Sensibile al Vantaggio (Advantage-Aware): L'aggiornamento è pesato dal vantaggio relativo al gruppo $\hat{A}_{i,t}$.
  • Vantaggio Positivo: Se il vantaggio è positivo, lo studente è incoraggiato ad aumentare la probabilità del token. Se anche il docente assegna un'alta probabilità, l'ancora della miscela è più alta, permettendo allo studente di aumentare la propria probabilità prima di colpire il limite di clipping (permettendo aggiornamenti più ampi). Se il docente assegna una bassa probabilità, l'ancora è più bassa, limitando l'entità dell'aggiornamento.
  • Vantaggio Negativo: Se il vantaggio è negativo, lo studente è incoraggiato a diminuire la probabilità. Se il docente assegna un'alta probabilità, il limite di clipping impedisce allo studente di sopprimere aggressivamente il token, preservando i comportamenti allineati al docente.
  • Vantaggio Vicino allo Zero: L'influenza del docente diventa trascurabile, prevenendo l'interferenza quando il segnale di ricompensa non è informativo.

Questo design bilancia implicitamente esplorazione, sfruttamento e imitazione assicurando che la guida del docente sia modulata dal segnale di ricompensa (vantaggio) e vincolata all'interno di una trust region.

3. Contributi Chiave

1. Framework RLAD: Un nuovo framework di distillazione per il post-training RL che integra la guida del docente tramite la Trust Region Ratio Distillation (TRRD), sostituendo la regolarizzazione KL statica con un obiettivo di rapporto di verosimiglianza pesato per vantaggio e con clipping.
2. Analisi Teorica: Un'analisi che dimostra come TRRD permetta un'imitazione selettiva condizionata dalla ricompensa. Preserva gli aggiornamenti vincolati dalla trust region e bilancia naturalmente il compromesso tra seguire il docente e ottimizzare le ricompense senza richiedere un delicato bilanciamento della perdita.
3. Validazione Empirica: Esperimenti estesi sul ragionamento logico (K&K Logistics) e sul ragionamento matematico a lungo contesto (AIME, AMC, MATH) che mostrano miglioramenti costanti rispetto ai baseline.

4. Risultati Sperimentali

Gli autori hanno valutato RLAD su compiti di ragionamento logico e ragionamento matematico complesso utilizzando modelli student con parametri da 0.6B a 8B e docenti fino a 32B.

• Ragionamento Logico (K&K Logistics):

  • Utilizzando un docente Qwen3-8B, RLAD ha migliorato significativamente l'accuratezza media per Qwen3-0.6B (da 0.76 a 0.94 a 8K di contesto) e Qwen3-1.7B (da 0.95 a 0.99 a 8K di contesto) rispetto a GRPO standard.
  • I guadagni sono stati più pronunciati sui subset più difficili (PPL7 e PPL8), dove RLAD ha mantenuto un chiaro margine rispetto a KDRL.
  • Le dinamiche di addestramento hanno mostrato una convergenza più rapida e ricompense finali più elevate rispetto a sia GRPO che KDRL.

• Ragionamento Matematico:

  • Modelli Base (stile R1-Zero): Su Qwen3-8B-Base, RLAD ha migliorato la performance media di +5.5 punti rispetto a GRPO su compiti a lungo contesto (30K).
  • Modelli Post-trained: RLAD ha superato la distillazione con KL esplicita (KDRL) e la distillazione offline.
  • Specifiche delle Metriche: RLAD ha mostrato guadagni maggiori su Pass@1 (Mean@32) rispetto a KDRL, che tendeva a mostrare guadagni relativi maggiori su Pass@32. Gli autori interpretano questo come il fatto che RLAD riesce con successo a guidare il miglioramento della policy allineata alla ricompensa (exploitation), mentre i guadagni di KDRL erano guidati maggiormente dalla pura imitazione di soluzioni del docente (diversità).
  • Stabilità: RLAD ha esibito dinamiche di addestramento più stabili con meno oscillazioni rispetto a KDRL, che soffriva di problemi di accoppiamento degli obiettivi.

• Robustezza ed Efficienza:

  • Docente Debole: Quando il docente era più debole dello studente, RLadi è rimasto vicino alle prestazioni di GRPO, mentre KDRL è degradato significativamente di più.
  • Iperparametri: RLAD si è rivelato ampiamente insensibile al coefficiente di miscelazione $\alpha$ (fissato a 0.5), a differenza di KDRL che è sensibile al peso $\lambda$ della KL.
  • Efficienza: L'integrazione di RLAD ha aggiunto solo circa il 12% di latenza di addestramento rispetto a GRPO standard, poiché il docente viene utilizzato solo per il calcolo della log-probabilità, non per la generazione dei percorsi (trajectories).

5. Significato e Assertions

Il paper sostiene che RLAD affronti il disallineamento fondamentale tra la supervisione statica del docente e gli aggiornamenti dinamici della policy RL. Incorporando la guida del docente nel meccanismo di aggiornamento della trust region, RLAD evita il "tuning fragile" richiesto dalla regolarizzazione KL multi-obiettivo.

Gli autori affermano che RLAD permette allo studente di muoversi verso comportamenti ad alto vantaggio all'interno di una distribuzione ancora, sfruttando opportunamente la guida del docente. Ciò si traduce in:

• Performance Superiore: Superamento costante di GRPO, KDRL e della distillazione offline, in particolare nei benchmark più impegnativi.
• Stabilità: Dinamiche di addestramento più stabili e migliore convergenza.
• Vero Ottimizzazione RL: A differenza dei metodi che si affidano pesantemente all'imitazione (che possono aumentare il Pass@32 ma non il Pass@1), RL-AD dà priorità al comportamento allineato alla ricompensa, portando a guadagni più forti nelle risposte corrette più probabili del modello (Pass@1).

Il paper conclude che, sebbene l'approccio dipenda da modelli docente open-weight con accesso ai logit ben calibrati, esso rappresenta un passo significativo nel rendere la distillazione compatibile con la natura dinamica del post-training RL.