Sintesi Tecnica: Apprendere dal proprio futuro (d-OPSD)

1. Definizione del Problema

Sebbene l'On-policy Self-distillation (OPSD) si sia dimostrata efficace per il post-training di Large Language Models autoregressivi (AR-LLM), la sua applicazione ai Diffusion Large Language Models (dLLM) rimane inesplorata. Gli esistenti framework OPSD sono intrinsecamente progettati per modelli AR, creando due discrepanze fondamentali quando applicati ai dLLM:

1. Costruzione del Teacher Subottimale: L'OPSD basata su AR inietta informazioni privilegiate (ad esempio, soluzioni di riferimento) come un prefisso al prompt, affidandosi al condizionamento da sinistra a destra $p(\text{suffix}|\text{prefix})$. Ciò entra in conflitto con i dLLM, che generano sequenze in modo non autoregressivo e supportano la generazione con ordine arbitrario, inclusa la capacità di condizionarsi sui suffissi $p(\text{prefix}|\text{suffix})$. Inoltre, le soluzioni di riferimento statiche non si allineano con la natura "on-policy" della self-distillation.
2. Livello di Supervisione Incompatibile: I modelli AR utilizzano una supervisione della divergenza a livello di token basata sulla predizione del token successivo. Al contrario, i dLLM operano tramite un processo di denoising iterativo in cui più token vengono predetti simultaneamente ad ogni step, e solo un sottoinsieme viene rivelato mentre altri rimangono mascherati. La supervisione a livello di token è quindi incompatibile con le transizioni di stato step-wise dei dLLM.

Di conseguenza, esiste la necessità di un framework OPSD su misura che sfrutti le uniche meccaniche generative dei dLLM per consentire un efficiente auto-miglioramento senza dipendere da modelli teacher esterni più forti.

2. Metodologia: d-OPSD

Gli autori propongono d-OPSD (diffusion On-Policy Self-distillation), il primo framework OPSD specificamente progettato per i dLLM. Il metodo opera su un singolo modello $p_\theta$ che funge sia da studente che da teacher, utilizzando due innovazioni centrali:

A. Costruzione del Self-Teacher tramite "Self-Future"

Inveve di appendere soluzioni di riferimento statiche come prefisso, d-OPSD sfrutta la capacità di condizionamento bidirezionale dei dLLM.

• Meccanismo: Lo studente campiona una traiettoria completa $Y = \{y_T, \dots, y_0\}$ da uno stato completamente mascherato a una risposta pulita $y_0$.
• Input del Teacher: Il teacher viene costruito rivelando parzialmente i token dalla propria risposta finale $y_0$ generata dallo studente come condizione di suffisso. Nello specifico, ad ogni step di denoising $t$, l'input del teacher $y_{\text{teacher},t}$ consiste nella sequenza rumorosa corrente con un sottoinsieme di indici $S_t$ (selezionati dalla risposta finale) rivelati.
• Concetto: Ciò consente allo studente di apprendere dalla propria "esperienza futura" (la risposta auto-generata) condizionata come suffisso, aderendo strettamente al paradigma on-policy in cui tutti i dati sono generati dallo studente stesso.

B. Supervisione della Divergenza a Livello di Step

Riconoscendo che i dLLM decodificano tramite predizione del passo successivo piuttosto che del token successivo, gli autori spostano l'obiettivo di supervisione dal livello di token al livello di step.

• Meccanismo: Ad ogni step di denoising $t$, il modello calcola le distribuzioni per tutti i token. La supervisione si concentra esclusivamente sui top-$k$ token più sicuri tra le posizioni attualmente mascherate (il sottoinsieme $K_t$ che determina la transizione di stato da $t$ a $t-1$).
• Obiettivo: La funzione di perdita minimizza la divergenza di Kullback-Leibler (KL) tra la distribuzione dello studente e la distribuzione del teacher su questo specifico sottoinsieme $K_t$ ad ogni step:
  $$L_t = \frac{1}{|K_t|} \sum_{i \in K_t} D_{KL}(P_{\text{student},t}^i \parallel P_{\text{teacher},t}^i)$$
  L'obiettivo totale aggrega questa perdita a livello di step attraverso l'intera traiettoria di $T$ step.

3. Contributi Chiave

1. Identificazione della Subottimalità: Il paper identifica che le formulazioni OPSD esistenti, centrate sugli AR, sono fondamentalmente disallineate con i dLLM (condizionamento tramite prefisso e supervisione a livello di token).
2. Nuovo Framework (d-OPSD): Gli autori introducono il primo framework OPSD per i dLLM, caratterizzato da:
   • Una costruzione del self-teacher basata sul suffisso che utilizza risposte auto-generate come informazioni privilegiate, consentendo l'apprendimento da "esperienze future di sé".
   • Supervisione della divergenza a livello di step che si allinea con le meccaniche di denoising iterativo dei modelli di diffusione.
3. Validazione Empirica: Esperimenti estesi dimostrano che d-OPSD consente a un singolo modello di guidare il proprio auto-miglioramento, superando i baseline di Reinforcement Learning con Verifiable Rewards (RLVR) e Supervised Fine-Tuning (SFT).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su quattro benchmark di ragionamento: GSM8K, MATH500 (ragionamento matematico) e Countdown, Sudoku (pianificazione). Il modello base utilizzato è stato LLaDA-8B-Instruct.

• Performance di Ragionamento: d-OPSD supera costantemente o eguaglia i baseline RLVR (ad esempio, diffu-GRPO) e SFT in tutti i task e le lunghezze delle sequenze. Ad esempio, su GSM8K (256 token), d-OPSD ha raggiunto un'accuratezza dell'81,0% rispetto al 79,8% di diffu-GRPO e al 78,8% di SFT.
• Efficienza di Campionamento: d-OPSD dimostra una efficienza di campionamento vastamente superiore, convergendo in circa il 10% degli step di ottimizzazione richiesti da RLVR. Ad esempio, su GSM8K, d-OPSD ha richiesto solo 425 step rispetto ai 7.700 di diffu-GRPO.
• Confronto con OPSD in stile AR: Uno studio di ablazione che sostituiva il teacher basato sul suffisso con un teacher basato sul prefisso in stile AR ha mostrato che la proposta costruzione del suffisso supera significativamente il controparte AR. L'analisi delle metriche "Overlap Top-K" ha rivelato che l'approccio in stile AR non riesce a introdurre nuova conoscenza (alto overlap), mentre l'approccio suffix di d-OPSD fornisce un intervallo appropriato di nuovi schemi di pensiero affinché lo studente possa apprendere.
• Risultati delle Ablazioni:
  • Obiettivo di Divergenza: La Reverse KL supera la Forward KL.
  • Rapporto di Ritenzione: Un rapporto di ritenzione ($\rho_{\text{teacher}}$) di 0,25 ha prodotto risultati ottimali, sebbene anche 0,10 abbia fornito guadagni significativi.
  • Selezione Top-k: Derivare il sottoinsieme top-$k$ dalla distribuzione del teacher ha prodotto prestazioni migliori rispetto alla distribuzione dello studente.
  • Stabilità: Il clipping pointwise della KL è stato trovato essenziale per prevenire il collasso della policy dopo il picco di performance, un modo di fallimento osservato sia in RLVR che in d-OPSD non clippato.

5. Significato e Rivendicazioni

Il paper sostiene di aprire una via promettente per il post-training dei dLLM, stabilendo che la self-distillation può essere efficacemente adattata ai modelli di diffusione. La significatività risiede in:

• Efficienza: Fornire un metodo che raggiunge prestazioni di ragionamento allo stato dell'arte con un costo computazionale drasticamente ridotto (efficienza di campionamento) rispetto a RLVR.
• Allineamento Teorico: Collegare con successo il divario tra apprendimento on-policy e meccaniche di diffusione riformulando le informazioni privilegiate come suffissi di "self-future" e la supervisione come divergenza a livello di step.
• Auto-Miglioramento: Dimostrare che un singolo modello può servire come proprio teacher senza richiedere un modello esterno più forte o dataset di riferimento statici, a patto che la traiettoria di generazione sia utilizzata correttamente.

Gli autori osservano che, sebbene efficace, l'approccio condivide la suscettibilità di RLVR al collasso della policy dopo aver raggiunto il picco di performance, suggerendo la necessità di futuri lavori sulle tecniche di stabilizzazione.