Sintesi Tecnica: Reinforcement Learning da Feedback Ricco con Distribuitional DAgger

Formulazione del Problema

Il paper affronta i limiti del paradigma dominante nell'addestramento dei modelli di ragionamento: il Reinforcement Learning da Ricompense Verificabili (RLVR). Nell'RLVR, i modelli generano molteplici soluzioni candidate e a ciascuna viene assegnata una singola ricompensa binaria (corretto/errato) basata sulla risposta finale. Questo approccio soffre di due gravi difetti:

1. Scarsa Assegnazione del Credito (Credit Assignment): La ricompensa terminale viene trasmessa uniformemente a tutti i token di una risposta, fallendo nel distinguere quali specifici passaggi di ragionamento abbiano contribuito al successo o al fallimento.
2. Inapplicabilità a Domini Non Verificabili: L'RLVR si basa sulla verifica automatica delle risposte finali, rendendolo inadatto a compiti in cui le risposte di verità fondamentale (ground-truth) non sono disponibili o sono difficili da controllare.

Sebbene molti contesti offrano un "feedback ricco" (ad esempio, tracce di esecuzione, log di errore, correzioni di esperti o auto-valutazioni), i metodi attuali che tentano di sfruttare ciò tramite l'auto-distillazione on-policy (ad esempio, SDPO, OPSD) affrontano ostacoli teorici e pratici. Gli approcci esistenti tipicamente minimizzano le $f$-divergenze (come la KL inversa o la Jensen-Shannon) tra una politica studentesca e una politica docente condizionata dal feedback. Gli autori sostengono che questi metodi manchino di garanzie per il miglioramento monotono della politica e non riescano a eseguire un'assegnazione del credito a livello di sequenza, portando spesso a una convergenza subottimale.

Metodologia: DistIL

Gli autori propongono DistIL (Distributional Imitation Learning), una variante distribuionale del classico algoritmo DAgger adattata per il reinforcement learning con feedback ricco.

Obiettivo Centrale

DistIL ottimizza una perdita di cross-entropia forward tra la politica studentesca ($\pi_\theta$) e la politica docente condizionata dal feedback ($\pi_T$) sugli stati visitati dallo studente:
$$\mathcal{L}_{\text{DistIL}}(\theta) := \mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot|x)} \left[ \sum_{t=1}^H H_\times(\pi_T(\cdot | s_t), \pi_\theta(\cdot | s_t)) \right]$$
dove $H_\times(p, q) = -\mathbb{E}_{y \sim p}[\log q(y)]$ è la cross-entropia forward. Fondamentalmente, il docente $\pi_T$ è condizionato dal feedback ricco $f$ (ad esempio, log di esecuzione o verità fondamentale) ma rimane fisso durante l'aggiornamento dello studente.

Gradiente e Assegnazione del Credito

A differenza dei precedenti metodi di auto-distillazione che approssimano i gradienti utilizzando solo discrepanze locali per token (effettivamente congelando la distribuzione degli stati), DistIL utilizza gradienti completi (full gradients). Il gradiente si decompone in due termini:

1. Imitazione Locale: Un termine pesato dal docente che aggiorna lo studente per corrispondere alla distribuzione del docente nello stato corrente.
2. Assegnazione del Credito Consapevole del Futuro (Future-Aware Credit Assignment): Un termine che propaga il disaccordo della cross-entropia osservato nei passi temporali futuri alle decisioni precedenti. Ciò assicura che le scelte dei token iniziali che influenzano la traiettoria verso uno stato con un alto disaccordo docente-studente ricevano il giusto credito o penalità.

L'algoritmo supporta docenti black-box (modelli esterni o esperti umani) poiché non richiede l'accesso alle log-probabilità del docente, ma solo campioni o la capacità di stimare l'obiettivo.

Contributi Chiave

1. Limiti Teorici dei Metodi Esistenti

Il paper fornisce prove formali dimostrando che gli obiettivi di auto-distillazione on-policy basati su $f$-divergenze (inclusa la KL inversa e la Jensen-Shannon) soffrono di due problemi fondamentali:

• Assenza di Garanzia di Miglioramento Monotono: Anche se il docente ha una ricompensa attesa strettamente superiore rispetto allo studente, minimizzare una $f$-divergenza non garantisce che un passo di gradiente naturale aumenti la ricompensa attesa dello studente. La direzione dell'aggiornamento può involontariamente aumentare la probabilità di azioni subottimali a causa del peso dato dai rapporti logaritmici.
• Convergenza Subottimale tramite Approssimazione Locale: Ignorando il termine di "credito futuro" (l'impatto delle azioni correnti sulle distribuzioni degli stati futuri), le approssimazioni locali del gradiente possono convergere a politiche strettamente subottimali. Gli autori dimostrano che, in semplici MDP a due step, gli aggiornamenti locali possono fallire nell'assegnare il credito alle decisioni precoci che sono cruciali per raggiungere stati ad alta ricompensa.

2. Garanzie Teoriche di DistIL

Si dimostra che DistIL possiede proprietà che allineano l'ottimizzazione con la massimizzazione della ricompensa:

• Miglioramento Monotono della Politica: Sotto le standard assunzioni di realizzabilità locale, l'obiettivo della cross-entropia forward garantisce che muoversi verso un docente migliore aumenti la ricompensa attesa dello studente.
• Limiti di Regret (Regret Bounds): L'algoritmo ammette un limite di regret sublineare, con tassi di convergenza dipendenti dai coefficienti di concentrabilità e dalla varianza del docente.
• Ottimizzazione Pass@N: L'obiettivo è dimostrato essere un massimizzatore di un limite inferiore della verosimiglianza pesata dal docente del successo (teacher-weighted likelihood of success). Ciò fornisce una spiegazione fondata del perché DistIL migliori le metriche Pass@N (la probabilità di generare almeno una risposta corretta in $N$ campioni) per tutti i valori di $N$.

3. Prestazioni Empiriche

DistIL è stato valutato in tre domini:

• Ragionamento Scientifico: Su SciKnowEval (Chimica, Fisica, Biologia, Scienza dei Materiali), DistIL ha superato SDPO, GRPO (on-policy e off-policy) e ha ottenuto punteggi Best@16 e Maj@16 più elevati. Notevolmente, DistIL ha mostrato un addestramento stabile, mentre SDPO ha mostrato oscillazioni significative e un crollo delle prestazioni in certi domini.
• Coding: Su LiveCodeBench, utilizzando le tracce di esecuzione come feedback, DistIL ha superato SDPO e GRPO sia in termini di correttezza esatta che di accuratezza parziale.
• Ragionamento Matematico Difficile: Su problemi impegnativi (AIME, HMMT) dove il modello base ha un tasso di successo pari a zero, i metodi RLVR (GRPO) non sono riusciti a migliorare a causa di segnali di vantaggio nulli. DistIL ha sfruttato con successo le soluzioni di verità fondamentale per migliorare le prestazioni, superando i baseline di auto-distillazione (SDPO, OPSD) e il SFT.

Significato e Rivendicazioni

Il paper sostiene che DistIL colmi il divario tra l'apprendimento per imitazione e il reinforcement learning fornendo un framework teoricamente fondato per apprendere da feedback ricchi. La sua importanza risiede in:

1. Correzione dell'Obiettivo: Allontanarsi dalla minimizzazione della $f$-divergenza (che può essere disallineata rispetto alla ricompensa) verso la cross-entropia forward, che garantisce il miglioramento monotono quando il docente è migliore.
2. Abilitazione dell'Apprendimento a Livello di Sequenza: Utilizzando i gradienti completi, DistIL risolve il problema dell'assegnazione del credito inerente al ragionamento a lungo termine, assicurando che le decisioni precoci siano valutate in base alle loro conseguenze a valle.
3. Robustezza Pratica: Il metodo funziona con docenti black-box e diversi tipi di feedback (dai log di correttezza binaria ai log di esecuzione), rendendolo applicabile a una gamma più ampia di compiti di ragionamento rispetto allo standard RLVR.

Gli autori concludono che la combinazione di un obiettivo allineato alla ricompensa e un'assegnazione del credito consapevole del futuro è essenziale per distillare efficacemente il feedback ricco nei modelli di ragionamento, portando a guadagni empirici costanti rispetto ai baseline allo stato dell'arte.