Sintesi Tecnica: Sul Bias di Posizione della Distillazione On-Policy

Problema

La Distillazione On-Policy (OPD) migliora l'efficienza dell'apprendimento del Reinforcement Learning (RL) per i Large Language Models (LLM) fornendo una supervisione densa a livello di token da un modello insegnante (teacher) a un modello studente (student) che genera i propri rollout. Tuttavia, gli obiettivi standard della OPD mediano uniformemente le perdite a livello di token (tipicamente tramite la divergenza KL), assumendo implicitamente un valore uguale per tutti i token in una sequenza.

Gli autori identificano un critico Bias di Posizione in questo approccio: man mano che i rollout dello studente si estendono, la traiettoria generata devia progressivamente dalla distribuzione del teacher. Di conseguenza, la qualità della supervisione del teacher degrada nelle posizioni successive. Le osservazioni empiriche rivelano che:

1. Asimmetria nel Valore della Supervisione: L'addestramento utilizzando solo il primo 30% dei token (prefix) produce prestazioni comparabili o superiori rispetto all'uso di tutti i token, mentre l'addestramento sull'ultimo 30% (suffix) fornisce benefici di apprendimento trascurabili.
2. Divergenza Persistente: Anche dopo la convergenza, la divergenza KL tra il teacher e lo studente rimane significativa, suggerendo che la OPD ottimizzi lo studente solo all'interno di una piccola regione locale dello spazio delle policy.
3. Inefficienza: Pesare uniformemente tutti i token spreca il limitato budget di aggiornamento del gradiente su token successivi dove lo studente è già derivato in regioni Out-of-Distribution (OOD) rispetto al teacher, rendendo il segnale di supervisione rumoroso o poco utile.

Metodologia: Distillazione On-Policy con Pesatura per Importanza (IW-OPD)

Per affrontare questo bias, gli autori propongono la Distillazione On-Policy con Pesatura per Importanza (IW-OPD), un metodo che regola dinamicamente i pesi dei token in base alla discrepanza accumulata tra le distribuzioni dello studente e del teacher.

Fondamento Teorico

Gli autori inquadrano la OPD come un problema di ottimizzazione vincolata (una proiezione locale) in cui la policy dello studente viene aggiornata all'interno di una regione di fiducia definita da un budget di divergenza KL $\rho$.

• Derivazione della Policy Ottimale: Risolvere la minimizzazione vincolata di $D_{KL}(q \| \pi_T)$ soggetta a $D_{KL}(q \| \pi_\theta) \leq \rho$ produce una policy ottimale in forma chiusa $q^\star$.
• Ripesatura del Rapporto di Verosimiglianza: La soluzione $q^\star$ ripesa la policy dello studente $\pi_\theta$ tramite il rapporto di verosimiglianza $(\pi_T / \pi_\theta)^\alpha$. Ciò implica che i token in cui la probabilità del teacher è significativamente più alta di quella dello studente (alta compatibilità) dovrebbero ricevere un peso maggiore, mentre i token in cui lo studente è derivato (basso rapporto) dovrebbero essere de-pesati.

Implementazione Pratica

L'uso diretto del rapporto di verosimiglianza grezzo è numericamente instabile a causa dell'accumulo esponenziale di probabilità su sequenze lunghe. IW-OPD introduce una discrepanza cumulativa del prefix non firmata stabile come proxy per il peso di importanza:

1. Discrepanza Accumulata: Invece del rapporto grezzo, il metodo somma gli advantage assoluti a livello di token ($|A^{OPD}_k|$) dall'inizio della sequenza fino al token corrente $t$. Ciò cattura la "deriva" totale dal percorso preferito dal teacher.
2. Normalizzazione: La discrepanza accumulata è standardizzata all'interno di ogni campione in un intervallo $[0, 1]$, garantendo una scalatura stabile.
3. Interpolazione: Per evitare la completa rimozione della supervisione per i token derivati, il peso finale è un'interpolazione tra il segnale OPD standard (peso 1) e il peso di importanza derivato dalla discrepanza.
   • Formula: $\tilde{r}^{IW-OPD}_t = 1 + \gamma \cdot (1 - \frac{\sum_{k<t} |A^{OPD}_k|}{\sum_{k<T} |A^{OPD}_k|})$.
   • Qui, $\gamma$ controlla il contributo della pesatura per importanza.

Questo approccio aumenta naturalmente il peso dei primi token (dove il prefix è compatibile con il teacher) e riduce il peso dei token successivi (dove la deriva accumulata è elevata), riallocando efficacementamente il budget del gradiente alle parti più informative della traiettoria.

Contributi Chiave

1. Identificazione del Bias di Posizione: Il paper dimostra empiricamente e spiega teoricamente che la OPD soffre di un bias di posizione in cui i primi token forniscono una supervisione significativamente più preziosa rispetto ai token successivi a causa della deriva distributiva.
2. Prospettiva di Ottimizzazione Vincolata: Gli autori forniscono una comprensione fondata della OPD come una proiezione locale vincolata, derivando che la strategia di aggiornamento ottimale richiede intrinsecamente una ripesatura basata sul rapporto.
3. Algoritmo IW-OPD: Propongono IW-OPD, un algoritmo pratico e stabile che stima l'importanza del token utilizzando la discrepanza cumulativa del prefix non firmata. Non richiede valutazioni aggiuntive del teacher oltre alla standard OPD.
4. Validazione Empirica: Estesi esperimenti mostrano che IW-OPD converge più velocemente e raggiunge prestazioni finali superiori rispetto alla standard OPD in varie scale di modelli e configurazioni teacher-student.

Risultati Sperimentali

Gli autori hanno valutato IW-OPD utilizzando modelli Qwen3 (0.6B, 1.7B, 4B, 30B e 235B) come teacher e student su benchmark di ragionamento matematico (AIME-2024, AIME-2025, HMMT-2025) e generazione di codice (HumanEval+, MBPP+).

• Guadagni di Performance: IW-OPD supera costantemente la standard OPD.
  • Nel setting dello studente Qwen3-1.7B con un teacher Qwen3-30B, IW-OPD ha migliorato l'accuratezza di AIME-2025 di +6.9 punti al passo 10 e di +1.7 punti alla convergenza.
  • In tutte le scale di studenti (da 0.6B a 30B) e setting di teacher, IW-OFP ha migliorato il punteggio di performance medio.
• Efficienza Campionaria: IW-OPD converge significativamente più velocemente. In diversi casi, il checkpoint di IW-OPD al passo 10 ha eguagliato o superato la performance finale della standard OPD.
• Effetti di Scaling:
  • Teacher più forti: IW-OPD rende i teacher più forti più efficienti dal punto di vista campionario. Mentre la standard OPD fatica inizialmente con gap più ampi tra teacher e studente (richiedendo più passi per beneficiare del teacher più forte), IW-OPD sfrutta immediatamente il segnale del teacher più forte.
  • Studenti più piccoli: I guadagni relativi di IW-OPD aumentano man mano che il modello studente diventa più piccolo (ad esempio, +9.5% di guadagno relativo per lo studente 0.6B rispetto a +1.8% per il 4B), indicando che è particolarmente efficace nella distillazione cross-scale dove la deriva è più severa.

Significato e Affermazioni

Il paper sostiene che una supervisione on-policy efficace deve tenere conto del contesto della traiettoria, non solo del disaccordo a livello di token. Riconoscendo che la supervisione del teacher non è uniformemente affidabile lungo una traiettoria generata dallo studente, IW-OPD rialloca il finito budget del gradiente ai prefix "compatibili con il teacher".

Gli autori concludono che questa semplice modifica — de-pesare i token in cui lo studente è derivato dalla distribuzione del teacher — migliora significativamente sia l'efficienza campionaria che la performance finale della distillazione on-policy. Questa intuizione suggerisce che i futuri metodi on-policy dovrebbero considerare la storia cumulativa del rollout (il prefix) quando assegnano crediti o pesi di supervisione, piuttosto che trattare tutti i token in una traiettoria come indipendenti e ugualmente preziosi.