DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay

Il paper propone DyJR, un framework di regolarizzazione che preserva la diversità nell'apprendimento per rinforzo dei modelli linguistici attraverso un buffer dinamico e una regolarizzazione basata sulla divergenza Jensen-Shannon, superando l'inefficienza dei metodi on-policy e prevenendo il collasso dei modi senza compromettere l'efficienza computazionale.

L'essenziale

🧠 Il Problema: L'Intelligenza Artificiale che "Dimentica" di Sognare

Immagina di avere un genio (un modello di Intelligenza Artificiale) che sta imparando a risolvere problemi di matematica complessi o a scrivere codice.
Per imparare, questo genio prova migliaia di soluzioni. Alcune funzionano, altre no.

Il metodo attuale (chiamato GRPO) funziona così:

1. Il genio prova una soluzione.
2. Se è giusta, viene lodato e impara da quella specifica soluzione.
3. Il problema: Dopo aver imparato, il genio butta via tutto il resto. Se ha provato 100 soluzioni diverse e solo 1 era perfetta, dimentica le altre 99.

La conseguenza? Il genio diventa un "monomane". Si fissa su quella unica strada che ha funzionato una volta e smette di esplorare altre possibilità. È come se un cuoco, dopo aver fatto un ottimo risotto, decidesse di cucinare solo risotto per il resto della vita, dimenticando come si fa la pasta o la pizza. In termini tecnici, si chiama "collasso delle modalità": l'AI perde la sua creatività e la sua capacità di trovare soluzioni alternative.

💡 La Soluzione: DyJR (Il "Diario di Bordo" Intelligente)

Gli autori propongono DyJR (Dynamic Jensen-Shannon Replay). Immagina DyJR non come un archivio polveroso, ma come un diario di bordo dinamico e intelligente.

Ecco come funziona, diviso in due idee semplici:

1. Il Diario che si Aggiorna da Solo (Il Buffer Dinamico)

I metodi vecchi provavano a salvare tutte le soluzioni passate, occupando un'enorme quantità di memoria (come riempire un magazzino infinito).
DyJR fa diversamente:

• La regola dell'età: Sa che le idee di oggi sono più fresche e utili di quelle di due settimane fa. Quindi, tiene solo le soluzioni recenti (quelle prodotte negli ultimi passi di allenamento).
• La regola del "momento giusto": All'inizio dell'allenamento, quando il genio sta ancora "sognando" e provando cose a caso, DyJR salva tante soluzioni diverse per catturare la massima creatività. Man mano che il genio diventa esperto e si stabilizza, il diario si restringe, tenendo solo le informazioni più recenti e pertinenti.
• Metafora: È come avere una lavagna. Invece di scrivere tutto su un muro infinito, cancelli la parte vecchia per fare spazio a quella nuova, ma ti assicuri di non cancellare le idee brillanti appena nate.

2. Non Copiare, ma "Bilanciare" (La Divergenza JS)

I metodi precedenti dicevano al genio: "Guarda questa soluzione perfetta del passato, copiala esattamente!". Questo portava il genio a imitare ciecamente, perdendo la sua individualità.
DyJR dice: "Guarda tutte le soluzioni che hai provato di recente. Non devi copiarne una sola, ma assicurati di non allontanarti troppo dalla varietà di quelle che hai già scoperto."

• Metafora: Immagina di essere un allenatore di calcio.
  • Il metodo vecchio dice: "Hai segnato un gol con un tiro potente? D'ora in poi tira solo così forte!" (Risultato: perdi la tecnica, il tiro diventa prevedibile).
  • DyJR dice: "Hai segnato con un tiro potente, ma anche con un tiro a giro e un rigore. Non dimenticare come si fanno tutti e tre. Mantieni la tua capacità di scegliere il tiro giusto in base alla situazione."

Questo "bilanciamento" si chiama Divergenza Jensen-Shannon. È una formula matematica che agisce come un freno di sicurezza: impedisce al modello di diventare troppo rigido e di perdere la sua diversità, senza però bloccare il suo apprendimento.

🚀 I Risultati: Perché è Geniale?

Hanno testato questo metodo su due campi difficili:

1. Matematica: Risolvere problemi di livello olimpico.
2. SQL: Tradurre domande in linguaggio per database (come chiedere a un computer di estrarre dati).

Cosa è successo?

• Migliore Accuratezza: Il modello ha risolto più problemi rispetto ai metodi precedenti (fino al 4-5% in più, che nel mondo AI è un'enorme differenza).
• Più Creatività: Analizzando le probabilità, hanno visto che il modello DyJR non si fissa su una sola risposta (Rank-1), ma mantiene aperte diverse opzioni (Rank-2, Rank-3), proprio come un umano che pensa a più soluzioni prima di agire.
• Efficienza: Non serve un supercomputer enorme. DyJR usa meno memoria perché non salva tutto, ma solo ciò che serve nel momento giusto.

🎯 In Sintesi

DyJR è come insegnare a un bambino a risolvere problemi non dicendogli "Memorizza la risposta giusta", ma dicendogli "Ricordati di tutte le strade diverse che hai provato per arrivare alla soluzione, così non ti blocchi se la strada principale si chiude".

Ridefinisce il modo in cui l'AI impara: non si tratta più solo di accumulare dati corretti, ma di preservare la diversità delle idee per diventare più robusti, creativi e intelligenti.

Sommario tecnico

1. Il Problema: Inefficienza e Collasso delle Modalità nel RLVR

L'apprendimento per rinforzo (RL) ha rivoluzionato le capacità di ragionamento dei Large Language Models (LLM), specialmente attraverso algoritmi on-policy come GRPO (Group Relative Policy Optimization) che utilizzano ricompense verificabili (RLVR). Tuttavia, questi metodi presentano due limiti fondamentali:

• Inefficienza nel campionamento: Gli algoritmi on-policy scartano i dati di roll-out (le traiettorie generate) dopo un singolo aggiornamento, sprecando risorse computazionali costose.
• Collasso delle modalità (Mode Collapse): I metodi esistenti di Experience Replay (come RLEP o Ex-GRPO) tentano di riutilizzare i dati storici trattandoli come campioni positivi aggiuntivi per aggiornamenti diretti del gradiente. Questo approccio, basato sull'ottimizzazione della verosimiglianza, tende a far sovrapposare il modello a percorsi di soluzione specifici, portando a un adattamento eccessivo (overfitting) e a una rapida perdita della capacità esplorativa. Il modello diventa troppo sicuro di un singolo percorso, riducendo la diversità delle risposte.

2. Metodologia: DyJR (Dynamic Jensen-Shannon Replay)

Gli autori propongono DyJR, un framework che ridefinisce lo scopo del Replay: non più solo rafforzare l'accuratezza, ma preservare la diversità delle traiettorie di successo. DyJR introduce due innovazioni chiave:

A. Costruzione Dinamica dei Dati (Buffer Temporale)

Invece di archiviare storicamente tutti i dati (che consuma molta memoria e include dati obsoleti), DyJR utilizza una strategia adattiva:

• Buffer Dinamico Time-Sensitive: Il buffer mantiene solo i campioni "perfetti" (ricompensa = 1) generati entro una finestra temporale recente (definita da un parametro Max Age $M$).
• Protocollo FIFO: I dati più vecchi vengono eliminati per mantenere il buffer sincronizzato con l'evoluzione attuale del modello.
• Selezione Adattiva: Vengono privilegiati i campioni ad alta confidenza (dove più percorsi generati sono corretti) per i task facili, ma il criterio si allenta per i task difficili per evitare la carenza di dati.
• Fase di Riscaldamento: Durante le prime fasi dell'addestramento (quando l'entropia crolla rapidamente), la capacità del buffer viene temporaneamente aumentata per catturare pattern ad alta entropia.

B. Regularizzazione tramite Divergenza Jensen-Shannon (JS)

DyJR abbandona l'aggiornamento diretto del gradiente sui dati di replay a favore di un vincolo di distribuzione:

• Obiettivo: Invece di massimizzare la probabilità dei dati storici, il modello viene regolarizzato per rimanere vicino a una distribuzione di riferimento dinamica (la miscela delle policy storiche nel buffer).
• Uso della JS Divergenza: Viene utilizzata la Divergenza Jensen-Shannon ($D_{JS}$) come termine di regolarizzazione. A differenza della KL Divergenza in avanti (che tende a coprire tutte le modalità ma può causare instabilità) o della KL inversa (che tende a collassare su una singola modalità), la JS è simmetrica e limitata.
• Funzione: Agisce come un "ancoraggio" flessibile che impedisce al modello di allontanarsi eccessivamente dai percorsi di successo diversificati senza alterare aggressivamente la direzione di ottimizzazione.

La funzione di perdita totale è:
$$\mathcal{L}_{total}(\theta) = \mathcal{L}_{GRPO}(\theta) + \alpha_{JS} \cdot \mathcal{L}_{JS}(\theta)$$

3. Contributi Chiave

1. Cambio di Paradigma: Spostamento dall'ottimizzazione dell'accuratezza tramite replay alla regularizzazione della diversità.
2. Strategia di Costruzione Dinamica: Un buffer non uniforme che si espande e si contrae in base alla fase di addestramento e alla vicinanza temporale dei dati, massimizzando l'efficienza della memoria.
3. Analisi Granulare: Dimostrazione empirica che la diversità nei primi stadi dell'addestramento è più preziosa delle traiettorie ad alta accuratezza degli stadi successivi.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su due domini principali: Ragionamento Matematico (dataset come AIME, HMMT, Minerva) e Text-to-SQL (BIRD, Spider).

• Prestazioni Superiori: DyJR supera significativamente GRPO e altri metodi basati su replay (RLEP, Ex-GRPO, DPH-RL).
  • Nel ragionamento matematico, DyJR ha raggiunto una media del 34.1% di accuratezza, contro il 29.8% di GRPO (+4.3%).
  • Nel Text-to-SQL, ha migliorato l'accuratezza Pass@1 del +3.3% su BIRD e del +5.0% su Spider rispetto a GRPO.
• Efficienza: Nonostante l'uso di un buffer, l'overhead di memoria GPU è trascurabile (meno di 1GB) rispetto ai metodi che archiviano intere storie di traiettorie. La velocità di addestramento è comparabile a quella di GRPO originale.
• Analisi della Diversità:
  • Entropia: Mentre GRPO mostra un crollo rapido dell'entropia e una probabilità del token Rank-1 che supera il 90%, DyJR mantiene un'entropia più alta e una distribuzione più bilanciata tra Rank-1, Rank-2 e Rank-3.
  • Scalabilità Pass@k: DyJR scala meglio all'aumentare del budget di campionamento ($k$), dimostrando una maggiore capacità di esplorazione rispetto ai baselines.
• Ablazione: L'uso della JS Divergenza è risultato superiore alla KL Divergenza in avanti in questo contesto, e un Max Age ($M$) piccolo (es. 8) ha funzionato meglio di finestre temporali più ampie, confermando l'importanza della "freschezza" dei dati.

5. Significato e Impatto

Il lavoro di DyJR è significativo perché risolve il compromesso tra efficienza del campione e diversità nel RL per LLM.

• Scalabilità: Dimostra che è possibile riutilizzare i dati storici senza i costi proibitivi o i rischi di collasso delle modalità associati ai metodi tradizionali.
• Robustezza: Il metodo è efficace sia su task di ragionamento profondo (matematica) che su task di generazione strutturata (SQL), indicando una generalizzazione trasversale.
• Insight Teorico: Fornisce la prova che, nel RLVR, il valore principale dei dati storici risiede nella loro capacità di mantenere la diversità esplorativa nelle fasi iniziali, piuttosto che nel semplice rafforzamento delle risposte corrette finali.

In sintesi, DyJR offre un approccio pratico ed efficiente per migliorare le capacità di ragionamento degli LLM, prevenendo la stagnazione in ottimi locali e mantenendo un comportamento esplorativo sano durante tutto il processo di addestramento.