Hybrid Policy Distillation for LLMs

Titolo: Hybrid Policy Distillation (HPD) per Large Language Models (LLM)

1. Il Problema

La compressione dei Large Language Models (LLM) tramite Knowledge Distillation (KD) è fondamentale per il loro dispiegamento pratico, ma le metodologie esistenti presentano sfide significative legate a tre assi di progettazione intrecciati:

1. Direzione della divergenza: La scelta tra Forward KL (FKLD) e Reverse KL (RKLD). Il FKLD favorisce la copertura delle modalità (mode coverage) ma tende a produrre previsioni eccessivamente "smussate" (over-smoothed), mentre il RKLD favorisce la ricerca delle modalità principali (mode-seeking) ma può diventare instabile quando il divario tra studente e insegnante è ampio.
2. Strategia di ottimizzazione: L'uso di funzioni di perdita (loss) basate sulla divergenza rispetto a segnali di ricompensa (reward).
3. Regime dei dati: L'uso di dati off-policy (da un dataset fisso o dall'insegnante) rispetto a dati on-policy (campionati dalla politica corrente dello studente).
   • I metodi off-policy soffrono di un disallineamento tra comportamento di addestramento e inferenza.
   • I metodi on-policy (come OPD) offrono un allineamento migliore ma comportano un elevato costo computazionale dovuto al campionamento di sequenze complete (rollouts).

L'obiettivo è sviluppare un metodo che bilanci copertura e ricerca delle modalità, mantenga l'efficienza computazionale dei dati off-policy e garantisca stabilità durante l'ottimizzazione.

2. Metodologia: Hybrid Policy Distillation (HPD)

Gli autori propongono una visione unificata della distillazione come un obiettivo di verosimiglianza logaritmica pesata (reweighted log-likelihood) a livello di token.

Concetti Chiave:

• Unificazione degli obiettivi: Sia la SFT (Supervised Fine-Tuning) che la distillazione tramite KL possono essere visti come minimizzazioni di una funzione obiettivo pesata:
  $$L(\theta) = \min_{\theta} -E_{(s_t, a_t) \sim D_\pi} [w(a_t | s_t) \log q_\theta(a_t | s_t)]$$
  dove $w(a_t | s_t)$ è un termine di pesatura che cattura la discrepanza locale tra studente e insegnante.

Il meccanismo HPD:
HPD ibrida i vantaggi del FKLD e del RKLD combinando dati off-policy e un campionamento on-policy leggero (senza rollouts completi). Il processo avviene a livello di token:

1. Stima del Gap (Negative K1 Estimator): Per ogni token, viene calcolato un estimatore del gap di KL inverso ($k_1$) tra la distribuzione dello studente $q_\theta$ e quella dell'insegnante $p$.
   $$k_1 = q_\theta(a^*_t | s_t) [\log p(a^*_t | s_t) - \log q_\theta(a^*_t | s_t)]$$
2. Gestione dei Token Esperti (Off-policy):
   • Se $k_1 > 0$ (lo studente sottostima il token esperto), si applica un peso positivo derivante dal FKLD per rafforzare il token.
   • Se $k_1 \le 0$ (lo studente sovrastima il token), il peso FKLD viene mascherato per evitare gradienti conflittuali, mantenendo solo il segnale di correzione.
3. Gestione dei Token Campionati (Lightweight On-policy):
   • Lo studente campiona un token non esperto $a_t \sim q_\theta$.
   • Se il token campionato è sovrastimato ($k'_1 < 0$), il suo peso viene impostato a un valore negativo. Questo segnale negativo sopprime la probabilità del token non esperto e ridistribuisce la massa di probabilità verso i token esperti.
   • Se il token esperto è sottostimato mentre un token non esperto è sovrastimato, il peso del token esperto viene raddoppiato per accelerare l'allineamento.

Vantaggi: HPD mantiene l'efficienza della supervisione "one-hot" (tipica della SFT) ma incorpora la densità informativa della distillazione KL, gestendo dinamicamente la direzione della divergenza in base allo stato di apprendimento del modello.

3. Contributi Chiave

1. Visione Unificata: Presentano una riformulazione della distillazione come obiettivo di verosimiglianza pesata, collegando SFT, FKLD e RKLD in un unico quadro teorico.
2. Algoritmo HPD: Propongono un metodo che integra ibridamente FKLD e RKLD su dati off-policy e campionamento on-policy leggero, senza richiedere rollouts completi costosi.
3. Stabilità ed Efficienza: HPD offre una maggiore stabilità nell'ottimizzazione e una migliore efficienza computazionale rispetto ai pipeline standard SFT + OPD.
4. Inizializzazione Robusta: HPD funge da eccellente inizializzazione per fasi successive di ottimizzazione (come OPD o DPO), portando a miglioramenti delle prestazioni finali.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre domini: ragionamento matematico (lunga generazione), dialogo e codice (generazione breve).

• Ragionamento Matematico (Off-policy):
  • HPD supera costantemente i baselines (SFT, SeqKD, RKLD, JSD) su modelli Qwen2.5 e LLaMA3.
  • Esempio significativo: Su Qwen2.5 (7B $\to$ 3B), HPD migliora le prestazioni medie del 41.0% rispetto alla SFT standard. Su LLaMA3 (8B $\to$ 3B), il miglioramento è del 77.9%.
  • HPD mantiene un'entropia stabile durante l'addestramento, evitando il collasso dell'entropia tipico della SFT.
• Personalizzazione e Dialogo (Off-policy):
  • Su benchmark come MT-Bench e AlpacaEval, HPD mostra la migliore capacità di preservare la coerenza del dialogo multi-turno e le capacità di allineamento.
• Codice (Off-policy):
  • HPD ottiene le migliori prestazioni medie su HumanEval e MBPP, dimostrando una maggiore robustezza e stabilità rispetto ad altri metodi di distillazione.
• Dati On-policy:
  • Quando HPD viene utilizzato come inizializzazione per l'OPD (On-Policy Distillation), supera le pipeline che iniziano con SFT, confermando che HPD fornisce una distribuzione dello studente più allineata e stabile.
• Impatto su DPO:
  • L'uso di HPD come fase di inizializzazione prima del DPO (Direct Preference Optimization) porta a guadagni di prestazioni superiori rispetto all'inizializzazione con SFT o RKLD.

5. Significato e Impatto

Questo lavoro è significativo perché risolve il compromesso storico tra l'efficienza computazionale della distillazione off-policy e la qualità dell'allineamento della distillazione on-policy.

• Praticità: HPD non richiede costosi rollouts completi, rendendolo scalabile per l'addestramento di modelli su larga scala.
• Flessibilità: Funziona efficacemente su diversi domini (matematica, codice, dialogo) e scale di modelli.
• Fondazione per il Futuro: Dimostra che una distillazione ibrida ben progettata può fornire una base solida per tecniche di allineamento avanzate (RLHF, DPO), permettendo di costruire modelli piccoli ma potenti con prestazioni vicine a quelle dei modelli insegnanti molto più grandi.

In sintesi, HPD rappresenta un avanzamento metodologico che unifica teoricamente e praticamente le diverse strategie di distillazione, offrendo una soluzione robusta ed efficiente per la compressione degli LLM.