Sintesi Tecnica: Dropout-GRPO per il Ragionamento Latente Continuo

1. Definizione del Problema

Il documento affronta un'incompatibilità strutturale fondamentale tra l'Ottimizzazione della Politica Relativa al Gruppo (GRPO) e i modelli di ragionamento latente continuo (specificamente l'architettura COCONUT).

• Il Requisito di GRPO: GRPO si basa sulla diversità di $K$ rollout all'interno di un gruppo per calcolare un vantaggio con media di gruppo non nullo ($A^{(k)} = r^{(k)} - \mu_r$). Se tutti i rollout producono traiettorie identiche, l'vantaggio collassa a zero, bloccando l'ottimizzazione.
• La Sfida del Ragionamento Latente: A differenza dei modelli standard di Chain-of-Thought (CoT) che derivano la diversità dal campionamento stocastico dei token, i modelli di ragionamento latente continuo (come COCONUT) operano tramite ricorrenza deterministica dello stato latente. Dato un prompt e parametri fissi, molteplici rollout generano traiettorie latenti e risposte identiche.
• La Conseguenza: L'applicazione di GRPO standard a questi modelli risulta in una varianza nulla ($\sigma_r = 0$), causando la scomparsa del gradiente della politica e impedendo la generazione di qualsiasi segnale di apprendimento.

2. Metodologia: Dropout-GRPO

Gli autori propongono Dropout-GRPO, un metodo che inietta la necessaria stocasticità nel processo di ragionamento latente attraverso un dropout strutturato, mantenendo al contempo le garanzie teoriche dell'inferenza variazionale.

Meccanismo Centrale: Variational Dropout a Maschera Condivisa

Invece del normale campionamento a livello di token, il metodo introduce la stocasticità a livello di parametro:

1. Generazione della Maschera: Per ogni rollout $k$, viene estratta una singola maschera di Bernoulli $\xi^{(k)}$.
2. Consistenza: Questa maschera rimane costante attraverso tutti i $T$ passi di ricorrenza latente all'interno di quel particolare rollout.
3. Perturbazione: La maschera induce una perturbazione strutturata $\tilde{\theta}(\xi)$ ai pesi della rete. La ricorrenza latente diventa $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$.
4. Ripetizione della Maschera (Mask Replay): Fondamentalmente, il seed del RNG per la maschera viene memorizzato durante la fase di rollout. Durante l'aggiornamento della politica, la stessa maschera viene rigenerata e applicata. Ciò assicura che la traiettoria valutata per il reward sia bit-identica alla traiettoria utilizzata per il calcolo del gradiente quando $\theta = \theta_{old}$.

Interpretazione Teorica

• Media del Modello Bayesiano: Seguendo Gal e Ghahramani, la maschera condivisa tratta ogni rollout come un singolo campione posteriore da una distribuzione strutturata sui parametri. La politica marginale $\bar{\pi}_\theta$ è interpretata come una media bayesiana del modello.
• Unbiasedness del Gradiente: Gli autori dimostrano che il gradiente surrogato utilizzando l'vantaggio con sola media ($A^{(k)} = r^{(k)} - \mu_r$) è uno stimatore non distorto (unbiased) del gradiente del reward atteso della politica marginale, scalato da $(1 - 1/K)$.
• Riduzione della Varianza: Riutilizzando la stessa maschera (Common Random Numbers) durante l'aggiornamento rispetto a quella usata durante il rollout, il metodo riduce la varianza rispetto al campionamento di una nuova maschera al momento dell'aggiornamento.

Raffinamenti Implementativi

Per stabilizzare l'addestramento, vengono introdotti quattro raffinamenti specifici:

1. Huberized $k_3$ KL: Sostituisce il classico clamping nella stima della divergenza KL con un'estensione di tipo Huber per evitare che i gradienti svaniscano per i token con grandi rapporti logaritmici, evitando "token fuggiti" che derivano senza regolarizzazione.
2. Annealing della Regione di Fiducia (Trust-Region): Il coefficiente KL $\beta$ viene attenuato proporzionalmente al tasso di apprendimento $\eta_t$ per evitare che la politica venga bloccata in un equilibrio dove la trazione KL annulla il gradiente della politica.
3. Filtro di Accuratezza a Livello di Gruppo: I gruppi con accuratezza empirica $\mu_r$ al di fuori di una finestra configurata (es. $[0.20, 0.95]$) vengono scartati. Questo funge da curriculum emergente, concentrando il calcolo su prompt con varianza non banale.
4. Skip Coordinato DDP: Garantisce la coerenza dell'addestramento distribuito quando si scartano i gruppi.

3. Contributi Chiave

1. Nuovo Paradigma per il RL Latente: Il documento identifica la natura deterministica della ricorrenza latente come la causa del fallimento di GRPO e propone il dropout variazionale a maschera condivisa come una soluzione minima e teoricamente pulita. È il primo metodo ad applicare con successo il RL relativo al gruppo ai modelli di ragionamento latente continuo.
2. Giustificazione Teorica:
   • Dimostra la ben definita del Jacobiano latente $\partial h_T / \partial \theta$ sotto la standard regolarità dei Transformer.
   • Dimostra l'unbiasedness del gradiente surrogato per la politica marginale.
   • Dimostra che il replay della maschera fornisce una riduzione della varianza tramite Common Random Numbers.
3. Validazione Empirica: Applica con successo GRPO al modello COCONUT sul benchmark GSM8K, un regime in cui il GRPO deterministico fallisce completamente.
4. Implementazione di Riferimento: Rilascia il codice che realizza il replay della maschera bit-identico attraverso l'intera ricorrenza latente.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti utilizzando Qwen2.5-1.5B come modello base con una profondità latente di $T=6$.

• Miglioramento delle Prestazioni: Dropout-GRPO ha migliorato l'accuratezza GSM8K pass@1 da un baseline COCONUT del 27.29% al 29.01%.
• Recupero delle Capacità: Il metodo ha invertito una regressiona delle prestazioni di 1.33 punti sul dataset SVAMP causata dall'addestramento curricolare, recuperando le prestazioni al 44.00% (eguagliando la varianza del modello base).
• Studi di Ablazione:
  • GRPO Deterministico: È fallito completamente, producendo una traiettoria di addestramento piatta identica all'inizializzazione SFT a causa dello zero vantaggio.
  • REINFORCE + Dropout (EMA): È fallito a causa dell'alta varianza; il gradiente del singolo rollout era troppo rumoroso per stabilizzare l'inizializzazione SFT.
  • Dropout-GRPO: L'unica configurazione che ha prodotto un segnale di apprendimento vitale, confermando che i baseline relativi al gruppo sono essenziali per la riduzione della varianza in questo contesto.

5. Significato e Claim

Il documento posiziona questo lavoro come un approccio pratico e teoricamente fondato per il post-training di Large Language Models (LLM) con ragionamento latente.

• Fattibilità: Dimostra che il post-training tramite RL è fattibile per i modelli che ragionano direttamente negli stati nascosti, una capacità precedentemente bloccata dalla mancanza di diversità nei rollout.
• Meccanismo: Inquadra il dropout non come un regolarizzatore (che viene spesso disabilitato nei moderni LLM), ma come una fonte di stocasticità esogena strutturata necessaria per l'ottimizzazione della politica variazionale.
• Limitazioni: Gli autori sono modesti riguardo l'entità del miglioramento (un guadagno di 1.72 punti su GSM8K), attribuendolo alla piccola scala del backbone (1.5B) e alla scarsità di dataset di ragionamento conciso adatti. Notano che la politica implementata (inferenza senza dropout) differisce dalla politica marginale ottimizzata, introducendo un bias dipendente dalle costanti di Lipschitz.

In sintesi, Dropout-GRPO colma il divario tra il reinforcement learning relativo al gruppo e il ragionamento latente continuo, consentendo l'ottimizzazione di modelli che precedentemente non potevano utilizzare GRPO a causa della loro dinamica interna deterministica.