Riepilogo Tecnico: Disaccoppiamento di KL e Traiettorie nella Distillazione di LLM

Enunciato del Problema

La distillazione della conoscenza è diventata una pietra angolare del post-addestramento dei Modelli Linguistici di Grandi Dimensioni (LLM), in particolare per i compiti di ragionamento, operando spesso in parallelo con l'Apprendimento per Rinforzo (RL). Tuttavia, lo spazio di progettazione della distillazione rimane scarsamente compreso. Le pipeline attuali seguono tipicamente uno dei due paradigmi: distillazione off-policy (addestramento su tracce generate dal docente, spesso tramite Supervised Fine-Tuning) o distillazione on-policy (OPD) (addestramento su rollouts generati dallo studente con feedback del docente).

Il paper sostiene che questi paradigmi prevalenti accoppiano implicitamente due scelte di progettazione ortogonali:

1. Fonte del Prefisso: Se la supervisione viene applicata a prefissi generati dal docente (off-policy) o dallo studente (on-policy).
2. Direzione KL a livello di token: Se l'obiettivo minimizza la KL forward ($KL(p_{teacher} \parallel q_{student})$) o la KL reverse ($KL(q_{student} \parallel p_{teacher})$).

La pratica convenzionale accoppia i prefissi del docente con la KL forward e i prefissi dello studente con la KL reverse. Gli autori sostengono che questo accoppiamento non è intrinseco a livello di token e che disaccoppiare questi assi rivela uno spazio di progettazione più ampio con quattro obiettivi validi, ciascuno corrispondente a regimi di apprendimento distinti.

Metodologia

Quadro Teorico: Disaccoppiamento degli Assi

Gli autori decompongono la divergenza KL a livello di sequenza utilizzando la regola della catena autoregressiva. Dimostrano che:

• La KL Forward ($KL(p_T \parallel q_S)$) si accoppia naturalmente con i prefissi del docente (off-policy) perché l'aspettativa è calcolata sulla distribuzione delle traiettorie del docente.
• La KL Reverse ($KL(q_S \parallel p_T)$) si accoppia naturalmente con i prefissi dello studente (on-policy) perché l'aspettativa è calcolata sulla distribuzione delle traiettorie dello studente.

Trattando questi come assi indipendenti, gli autori definiscono quattro obiettivi:

1. Prefisso del Docente + KL Forward: Corrisponde a SFT off-policy (distillazione standard).
2. Prefisso dello Studente + KL Forward: Corrisponde a SFT on-policy in stile DAgger.
3. Prefisso del Docente + KL Reverse: Corrisponde a un obiettivo di distillazione in stile RL offline.
4. Prefisso dello Studente + KL Reverse: Corrisponde a Distillazione On-Policy (OPD), vista come RL on-policy con ricompensa densa.

Il paper stabilisce identità a livello di gradiente che mostrano che:

• La KL Forward produce un gradiente di cross-entropy in stile SFT con target soft del docente.
• La KL Reverse produce un gradiente di policy in stile RL con una ricompensa densa definita dal rapporto logaritmico docente-studente ($r = \log p_T - \log q_S$).

Configurazione Sperimentale

Gli autori conducono uno studio controllato sul ragionamento matematico utilizzando modelli Qwen3 (studente da 0.6B, docenti da 4B/8B). Valutano i quattro obiettivi in due regimi:

1. Distillazione Autonoma: Confronto delle prestazioni finali dei modelli addestrati esclusivamente con ciascun obiettivo.
2. Distillazione-poi-RL: Utilizzo dei checkpoint distillati come inizializzazioni per l'Group Relative Policy Optimization (GRPO).

Monitorano accuratezza, entropia predittiva e lunghezza della risposta su varie lunghezze di sequenza (128 vs 4096 token) e benchmark (MATH500, GSM8K, AIME24, AMC23).

Contributi Chiave

1. Identificazione di Tre Trade-off

L'analisi empirica rivela tre trade-off fondamentali che governano la progettazione della distillazione:

• Trade-off Accuratezza–Entropia (Direzione KL): La KL reverse raggiunge costantemente una maggiore accuratezza autonoma (Avg@k), ma spinge l'entropia predittiva verso il collasso, riducendo la diversità (Pass@k più basso) e rendendo il modello meno stabile per il successivo RL. La KL forward preserva l'entropia, supportando un miglioramento RL più stabile, sebbene con una minore accuratezza autonoma iniziale.
• Trade-off Qualità–Calcolo (Fonte del Prefisso): I prefissi dello studente (on-policy) producono una maggiore accuratezza a parità di passi di addestramento supervisionando stati che lo studente visita effettivamente. Tuttavia, i prefissi del docente (off-policy) sono più efficienti dal punto di vista computazionale a parità di FLOP perché permettono di riutilizzare i logit del docente in cache e le traiettorie offline, evitando il costo della generazione online dello studente.
• Trade-off Accuratezza–Stabilità (Lunghezza dell'Addestramento): La distillazione su sequenze lunghe (4096 token) migliora l'accuratezza del ragionamento, ma, in particolare sotto KL reverse, causa un grave collasso dell'entropia e "inflazione della lunghezza" (risposte che raggiungono il limite di generazione). L'addestramento su sequenze corte è più stabile ma meno accurato.

2. Metodi Proposti per Bilanciare i Trade-off

Motivati da questi risultati, gli autori propongono due metodi pratici:

• Miscelazione KL: Una combinazione ponderata di KL forward e reverse a ogni prefisso. Gli autori scoprono che nella distillazione su sequenze lunghe, una miscela con prevalenza forward è cruciale. Mantiene i benefici di accuratezza della KL reverse prevenendo al contempo il collasso dell'entropia e l'inflazione della lunghezza, stabilizzando la dinamica di addestramento.
• Curriculum di Lunghezza a Cancellazione dell'Entropia: Un programma di addestramento che inizia con orizzonti brevi e aumenta la lunghezza della sequenza solo mentre l'entropia predittiva rimane sopra una soglia di stabilità. Questo approccio evita il collasso dell'entropia associato all'addestramento fisso su orizzonti lunghi.

Risultati

Distillazione Autonoma

• La KL Reverse generalmente supera la KL Forward nell'accuratezza autonoma (Avg@k), con guadagni medi di +2.45 punti.
• I Prefissi dello Studente superano i Prefissi del Docente nell'accuratezza, specialmente in contesti di sequenze lunghe (guadagno di +3.55 Avg@k).
• Collasso dell'Entropia: La KL reverse pura in contesti di sequenze lunghe spinge l'entropia vicino allo zero e causa l'inflazione delle lunghezze delle risposte fino al limite massimo (es. 8192 token), degradando il Pass@k.

Distillazione-poi-RL (GRPO)

• Stabilità di Inizializzazione: I modelli inizializzati con KL Forward (anche se con minore accuratezza autonoma) mostrano un miglioramento più affidabile durante il GRPO perché mantengono un'entropia più elevata.
• Insidia della KL Reverse: I modelli inizializzati con KL reverse spesso iniziano con una maggiore accuratezza ma degradano o raggiungono un plateau durante il RL a causa della bassa entropia e della ridotta capacità di esplorazione.
• Successo del Curriculum: Il curriculum di lunghezza a cancellazione dell'entropia migliora significativamente le prestazioni. Rispetto all'addestramento fisso a 4096 token, ha migliorato l'Avg@k di 3.6 punti e il Pass@k fino a 5.8 punti, riducendo al contempo la lunghezza media delle risposte di circa 3 volte (es. da ~8185 a ~2688 token).

Risultati della Miscelazione KL

In contesti di sequenze lunghe, una miscela con alto peso della KL forward (es. 0.8 Forward / 0.2 Reverse) ha preservato la maggior parte dell'accuratezza della KL reverse, stabilizzando l'entropia e prevenendo l'inflazione della lunghezza.

Significato e Affermazioni

Il paper afferma di fornire un quadro teorico unificato che collega ricette di addestramento disparate (SFT, DAgger, RL Offline, OPD) come istanze specifiche di due scelte di progettazione ortogonali. Disaccoppiando la fonte del prefisso e la direzione KL, gli autori sostengono che i praticanti possano andare oltre i paradigmi rigidi per progettare obiettivi di distillazione del ragionamento che bilanciano esplicitamente:

• Accuratezza (tramite KL reverse e sequenze lunghe),
• Diversità/Esplorazione (tramite KL forward e preservazione dell'entropia),
• Efficienza Computazionale (tramite prefissi del docente/caching), e
• Addestrabilità RL a valle (tramite entropia stabile).

I metodi proposti (miscelazione KL e curriculum a cancellazione dell'entropia) sono presentati come strumenti pratici per navigare i trade-off identificati, consentendo l'addestramento di modelli di ragionamento più piccoli che siano sia accurati che stabili per un'ulteriore ottimizzazione della policy. Gli autori sottolineano che il "migliore" obiettivo autonomo non è necessariamente la migliore inizializzazione per una pipeline RL, evidenziando la necessità di una visione olistica del processo di post-addestramento.