Self-Distilled Policy Gradient

Riepilogo Tecnico: Self-Distilled Policy Gradient (SDPG)

Definizione del Problema
Il Reinforcement Learning con Ricompense Verificabili (RLVR) è diventato uno standard per il post-training di Large Language Models (LLM) su compiti di ragionamento come la matematica e la generazione di codice. Algoritmi come il Group Relative Policy Optimization (GRPO) ottimizzano rispetto a ricompense basate su regole (es. risposte corrette/errate) senza richiedere annotazioni di preferenza umana. Tuttavia, l'RLVR affronta due limitazioni primarie:

1. Ricompense Sparse: Il segnale di ricompensa viene applicato a livello di sequenza, rendendo difficile l'attribuzione del credito a specifici passaggi di ragionamento (token), il che porta a un'esplorazione inefficiente.
2. Instabilità dell'Addestamento: Le prime fasi dell'addestramento soffrono spesso di vantaggi negativi, causando instabilità e una potenziale degradazione delle capacità linguistiche fondamentali.

Sebbene l'On-Policy Distillation (OPD) sia stata proposta per fornire una supervisione densa a livello di token, gli approcci tradizionali si affidano a modelli teacher esterni più grandi, che comportano elevati costi di memoria e introducono discrepanze nella distribuzione. Recentemente, i metodi di "self-distillation", in cui il modello funge sia da studente che da teacher utilizzando un contesto privilegiato (es. soluzioni ground-truth), affrontano il problema della memoria ma rischiano di rinforzare token localmente plausibili su traiettorie globalmente errate se non adeguatamente vincolati. Inoltre, la pura self-distillation può portare al collasso del modo (mode collapse) e a una riduzione dell'esplorazione.

Metodologia: Framework SDPG
Il documento propone il Self-Distilled Policy Gradient (SDPG), un framework che integra la self-distillation on-policy con intero vocabolario esatto con l'ottimizzazione della policy regolata tramite KL. L'obiettivo centrale combina tre componenti:

1. Policy Gradient basato sull'Esito ($L_{out}$):
   SDPG mantiene il segnale di ricompensa binario sparso proveniente da un verificatore. Utilizza un vantaggio relativo al gruppo ($A_{out}$) normalizzato all'interno di un gruppo di rollout per stabilizzare l'apprendimento, similmente a GRPO, ma formulato come una perdita surrogata diretta di tipo REINFORCE senza il clipping del rapporto di importanza tipico di PPO.

2. On-Policy Distillation Esatta con Intero Vocabolario ($L_{OPD}$):
   Inveve di campionare i token per approssimare il segnale del teacher, SDPG calcola la divergenza di Kullback-Leibler (KL) inversa esatta tra la distribuzione dello studente ($p_t$) e la distribuzione del teacher privilegiato ($q_t$) condizionata da un contesto aggiuntivo $c$ (es. soluzioni o percorsi di ragionamento).

• Identità del Gradiente: Gli autori derivano un risultato teorico chiave (Proposizione 3.1): il gradiente della perdita a intero vocabolario della reverse-KL rispetto alla policy dello studente è localmente identico a un aggiornamento di policy di tipo detached-sampling. L' "advantage" in questo aggiornamento è un log-rapporto centrato tra le probabilità del teacher e dello studente.
• Meccanismo di Gating: Per evitare che il teacher privilegiato rinforzi traiettorie errate, la perdita di distillazione è controllata da un gate basato sull'advantage dell'esito. La distillazione viene applicata solo quando il rollout è verificato come corretto ($A_{out} > 0$).

3. Regolarizzazione della Policy di Riferimento ($L_K$):
   Per evitare che la policy si discosti troppo dalle sue capacità iniziali o collassi, SDPG ancora la policy aggiornata a una policy di riferimento fissa ($\pi_{ref}$). Il documento impiega la divergenza Unnormalized KL (UKL) (forward o reverse) per questa regolarizzazione, che offre gradienti più eleganti e simmetrici rispetto alla standard KL normalizzata, specialmente quando si trattano misure non normalizzate.

Strategie di Stabilizzazione
Per gestire il rumore del segnale privilegiato e il compromesso tra esplorazione e sfruttamento, SDPG introduce due stabilizzatori:

• Positive Advantage Gating: Assicura che la distillazione avvenga solo su traiettorie approvate dal verificatore.
• Programma di Warmup-Decay ($\beta(k)$): Il coefficiente di distillazione $\beta$ viene aumentato gradualmente (warmup) per permettere al modello di trovare traiettorie corrette prima di fare affidamento sulla distillazione, e poi decresce (decay) verso la fine dell'addestramento. Ciò evita che il segnale privilegiato vincoli eccessivamente il modello una volta che ha interiorizzato la conoscenza necessaria, incoraggiando l'esplorazione nelle fasi finali.

Contributi Chiave

1. Derivazione Teorica: Il documento deriva una forma esatta di policy-gradient locale per la on-policy distillation a intero vocabolario della reverse-KL, dimostrando che è equivalente a un aggiornamento di detached-sampling con un advantage basato su log-rapporto centrato.
2. Obiettivo Innovativo: Propone un obiettivo unificato che combina ricompense binarie basate sull'esito, distillazione privilegiata a intero vocabolario con gating e regolarizzazione KL non normalizzata rispetto a un riferimento fisso.
3. Tecniche di Stabilizzazione: L'introduzione del positive-advantage gating e di un programma di warmup-decay per il coefficiente di distillazione per bilanciare la supervisione densa con l'esplorazione.

Risultati Sperimentali
Gli autori hanno valutato SDPG su benchmark di ragionamento matematico (AIME2024, AIME2025, AMC23) utilizzando Qwen3-4B.

• Performance: Entrambe le varianti di SDPG (utilizzando Unnormalized Forward KL e Unnormalized Reverse KL) hanno superato il baseline GRPO e il baseline di self-distillation RLSD in tutti i benchmark.
• Stabilità: Le varianti SDPG hanno raggiunto plateau ad alta ricompensa significativamente più velocemente di GRPO. Notevolmente, SDPG ha mantenuto un'entropia dell'attore più elevata durante l'addestramento rispetto a RLSD, che ha mostrato un collasso dell'entropia (un segno di collasso del modo) intorno allo step 250.
• Efficienza: I metodi SDPG hanno prodotto lunghezze di risposta sufficienti per il ragionamento multi-step, ma più concise rispetto agli output verbali di GRPO.

Significato
Il documento sostiene che SDPG riesce a colmare il divario tra l'RLVR a ricompensa sparsa e la self-distillation a supervisione densa. Trattando la self-distillation come un aggiornamento di policy gradient e controllando rigorosamente il segnale privilegiato attraverso il gating e la programmazione, SDPG ottiene una stabilità e una performance superiori. Dimostra che la supervisione densa a livello di token, derivata dal contesto privilegiato del modello stesso, può guidare efficacementamente i modelli di ragionamento senza l'overhead di memoria di teacher esterni o l'instabilità della pura self-distillation.