Sintesi Tecnica: OPD+: Ripensare il Design dell'Advantage per la Distillazione On-Policy

1. Definizione del Problema

La distillazione on-policy (OPD) è una tecnica utilizzata per trasferire le capacità da modelli linguistici teacher capaci a modelli student base, ottimizzando un segnale di supervisione denso a livello di token derivato dai rollout generati dallo studente. Mentre le implementazioni esistenti di OPD (ad es. Tinker) spesso utilizzano la divergenza di Kullback-Leibler (KL) come obiettivo primario, esse impiegano una specifica scelta di design: interrompere il gradiente (stop-gradient) sul termine di reward (la differenza del log-rapporto tra la verosimiglianza dello studente e quella del teacher).

Questo design è motivato dal desiderio di stabilità numerica, tracciando paralleli con le reti target nei Deep Q-Networks (DQN). Tuttavia, gli autori identificano un critico difetto teorico in questo approccio:

1. Stime del Gradiente Distorte: Per le $f$-divergenze generali, la funzione di reward dipende dai parametri della policy dello studente $\theta$. Interrompere il gradiente su questo termine di reward ignora la dipendenza diretta del reward da $\theta$, portando a stime del gradiente matematicamente errate.
2. Degradazione delle Prestazioni: Questa distorsione causa un collasso catastrofico dell'addestramento per funzioni di divergenza diverse dalla Reverse KL (ad es. Forward KL e Jensen-Shannon Divergence), rendendole impossibili da addestrare nelle normali impostazioni OPD. Anche per la Reverse KL, sebbene la distorsione teoricamente svanisca grazie all'identità dello score, i risultati empirici suggeriscono che il design stop-gradient renda comunque subottimale le prestazioni.

2. Metodologia: Il Framework OPD+

Il paper propone OPD+, un framework di ottimizzazione corretto che rivaluta il design dell'advantage per la distillazione on-policy basata sulla $f$-divergenza.

2.1 Fondamento Teorico

Gli autori formulano l'obiettivo OPD come la minimizzazione della $f$-divergenza tra la policy dello studente $p_\theta$ e la policy del teacher $q$. La funzione obiettivo è definita come:
$$J_f(\theta) = \mathbb{E}_{x \sim X, y \sim p_\theta(\cdot|x)} \left[ \sum_{n=1}^{L_y} f(u_{\theta,n}) \right]$$
dove $u_{\theta,n} = q(y_n | y_{<n}, x) / p_\theta(y_n | y_{<n}, x)$ è il rapporto di densità a livello di token, e $f$ è una funzione convessa con $f(1)=0$.

2.2 La Correzione

Il paper deriva il vero gradiente di questo obiettivo. La OPD standard utilizza uno stimatore del gradiente della policy con "stop-gradient" dove il reward $r = -f(u_{\theta,n})$ è trattato come fisso. Gli autori dimostrano che il vero gradiente contiene due termini:

1. Il termine standard della funzione di score: $\mathbb{E}[r \nabla_\theta \log p_\theta]$.
2. Un termine di gradiente del reward diretto: $\mathbb{E}[\nabla_\theta r]$.

Per una $f$-divergenza generica, il termine di gradiente del reward diretto non svanisce. Nello specifico, $\nabla_\theta r = u_{\theta,n} f'(u_{\theta,n}) \nabla_\theta \log p_\theta$.
Di conseguenza, la corretta funzione di advantage non è semplicemente $-f(u)$, ma un termine pesato:
$$w_f(u) = -f(u) + u f'(u)$$

2.3 Caso Speciale: Reverse KL

Per la Reverse KL, $f(u) = -\log u$. Il termine di correzione diventa:
$$w(u) = -(-\log u) + u(-1/u) = \log u + 1$$
Il termine $+1$ aggiuntivo è un baseline costante. Poiché $\mathbb{E}[\nabla_\theta \log p_\theta] = 0$ (identità dello score), questo termine teoricamente svanisce in attesa. Tuttavia, gli autori sostengono che in pratica, questo baseline costante agisce come un meccanismo di riduzione della varianza, migliorando la velocità di convergenza e la stabilità.

2.4 Implementazione

OPD+ richiede solo una modifica di una singola riga di codice nelle implementazioni OPD esistenti: sostituire il reward della divergenza negativa grezza $-f(u)$ con il peso corretto $w_f(u)$. Questo framework supporta varie $f$-divergenze, incluse Forward KL, Reverse KL e Jensen-Shannon Divergence (JSD).

3. Contributi Chiave

1. Framework di Ottimizzazione Generico: Il paper fornisce un framework unificato per la distillazione on-policy basata su $f$-divergenze generali, derivando il gradiente e le funzioni di perdita corretti che tengono conto della dipendenza del reward dai parametri della policy.
2. Dimostrazione Teorica della Distorsione: Dimostra matematicamente che l'operazione standard di stop-gradient porta a stime del gradiente distorte per le $f$-divergenze generali, invalidando la pratica comune di trattare il reward come fisso.
3. Algoritmo OPD+: Gli autori introducono OPD+, l'algoritmo corretto che elimina il collasso dell'addestramento osservato con Forward KL e JSD e migliora le prestazioni anche per la Reverse KL.
4. Validazione Empirica: Il lavoro mette in discussione la convinzione prevalente che la Reverse KL sia la metrica di distanza superiore o l'unica praticabile per la distillazione on-policy, dimostrando che altre divergenze possono superare la Reverse KL quando ottimizzate correttamente.

4. Risultati Sperimentali

Gli autori hanno valutato OPD+ sulla famiglia di modelli QWEN3 (1.7B e 8B parametri) in due domini distinti: Matematica Competitiva (DeepMath, AIME 24/25, HMMT 25) e Uso di Strumenti Agente (Agentic Tool Use).

4.1 Salvataggio delle Divergenze "Non Addestrabili"

Sotto il design standard stop-gradient (OPD), la Forward KL e la JSD hanno mostrato un fallimento catastrofico, scendendo allo 0% di accuratezza nei compiti di tool-use e fallendo nel raggiungere il baseline SFT-inizializzato nei compiti di matematica.

• Correzione OPD+: Applicando l'advantage corretto $w(u)$, OPD+ ha eliminato completamente questo modo di fallimento.
• Prestazioni JSD: JSD+ non solo si è stabilizzata, ma ha superato il baseline standard della Reverse KL. Nel benchmark impegnativo AIME 25, JSD+ ha raggiunto il 50.42% di accuratezza rispetto al 48.54% della Reverse KL. Nel compito di Tool Use, JSD+ ha raggiunto il 60.29%, superando il baseline della Reverse KL del 59.97%.

4.2 Miglioramento della Reverse KL

Anche nel dominio in cui la Reverse KL è tradizionalmente dominante, OPD+ ha prodotto miglioramenti empirici netti:

• AIME 24: Reverse KL+ è migliorata dal 64.27% al 64.69%.
• Tool Use: Reverse KL+ è migliorata dal 59.97% al 61.44%.
• Convergenza: La dinamica dell'addestramento (Figura 1) mostra che le varianti OPD+ accelerano la convergenza precoce e prevengono il plateau o la degradazione delle prestazioni osservati nei baseline standard.

4.3 Auto-distillazione (OPSD)

Il metodo è stato validato anche in un contesto di Auto-distillazione On-Policy (OPSD), dove il modello funge sia da studente che da teacher (con il teacher condizionato su soluzioni ground-truth privilegiate). OPD+ ha costantemente superato le versioni vanilla stop-gradient in tutti i tipi di divergenza, con la Reverse KL corretta che mostrava le migliori prestazioni tra le divergenze.

5. Significato e Rivendicazioni

Il paper sostiene che l'inferiore prestazione storica delle altre $f$-divergenze (Forward KL, JSD) nella distillazione on-policy non è un limite intrinseco delle divergenze stesse, ma piuttosto un effetto collaterale di un coefficiente di gradiente mal specificato causato dal design stop-gradient.

Correggendo la funzione di advantage per renderla matematicamente fedele all'obiettivo della $f$-divergenza, OPD+:

1. Espande lo Spazio di Design: Permette l'uso di una varietà più ampia di distanze statistiche per la distillazione, andando oltre lo standard del settore della Reverse KL.
2. Migliora le Prestazioni: Stabilisce nuovi picchi di prestazione state-of-the-art sui benchmark di ragionamento e tool-use.
3. Mette in Discussione le Convenzioni: Mette in discussione la comune convinzione che la Reverse KL sia intrinsecamente superiore per via della sua stabilità, suggerendo che tale stabilità fosse un artefatto del metodo di ottimizzazione piuttosto che della scelta della divergenza.

Gli autori concludono che il loro lavoro fornisce un algoritmo matematicamente più corretto ed empiricamente più efficace, capace potenzialmente di innescare ulteriore ricerca nel design delle funzioni di advantage per i metodi di distillazione on-policy. Riconoscono i limiti riguardanti l'analisi della scalabilità del modello e il confronto tra la verosimiglianza a singolo token e la distillazione dei logit, indicandoli come aree per il lavoro futuro.