Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

Riepilogo Tecnico: Oltre la Normalizzazione: Ripensare la Funzione di Partizione come Programmatore di Difficoltà per RLVR

Enunciato del Problema
L'Apprendimento per Rinforzo con Ricompense Verificabili (RLVR) si è dimostrato efficace nel potenziare le capacità di ragionamento dei Modelli Linguistici su Grande Scala (LLM). Tuttavia, i metodi standard di massimizzazione della ricompensa come PPO e GRPO spesso portano a distribuzioni di output eccessivamente affilate, riducendo la diversità di generazione. Sebbene approcci recenti che utilizzano le GFlowNets affrontino questo problema addestrando politiche per corrispondere a una distribuzione target indotta dalla ricompensa, essi trattano tipicamente la funzione di partizione appresa ($Z$) esclusivamente come una costante di normalizzazione necessaria per la stabilità dell'addestramento. Questo documento sostiene che tale prospettiva sottoutilizza le informazioni codificate all'interno di $Z$, mancando un'opportunità per migliorare l'efficienza del campione senza incorrere in costi computazionali aggiuntivi.

Metodologia: PACED-RL
Gli autori propongono RL Guidato dalla Funzione di Partizione (PACED-RL), un framework di post-addestramento che reinterpreta la funzione di partizione della GFlowNet non solo come normalizzatore, ma come segnale per l'accuratezza online per prompt. La metodologia consta di tre componenti fondamentali:

1. Reinterpretazione Teorica della Funzione di Partizione:
   Il documento stabilisce un legame teorico tra la funzione di partizione ottimale $Z^*(x)$ e l'accuratezza online $p_{old}(x)$ di un prompt di domanda $x$ sotto la politica pre-aggiornamento. Modificando la perdita di Bilanciamento delle Traiettorie (TB) per ancorare la divergenza KL alla politica corrente ($\pi_{old}$) piuttosto che a un modello di riferimento fisso, gli autori derivano che $Z^*(x)$ codifica direttamente informazioni sull'accuratezza. Nello specifico, $p_{old}(x) \approx \beta \log Z^*(x)$, dove il termine di divergenza KL si è dimostrato empiricamente trascurabile durante le dinamiche di addestramento standard. Ciò permette alla funzione di partizione, già appresa per la normalizzazione, di fungere da stimatore di accuratezza online gratuito.

2. Selezione Adattiva dei Prompt:
   Sfruttando le stime di accuratezza derivate da $Z_\phi(x)$, PACED-RL implementa una strategia di selezione dei prompt consapevole della difficoltà. Invece di campionare i prompt in modo uniforme, il metodo seleziona avidamente i primi $m$ prompt le cui accuratezze stimate sono più vicine a un valore target $\tau$ (tipicamente 0.5). Questo concentra l'addestramento su campioni di "difficoltà intermedia", che lavori precedenti suggeriscono offrire i segnali di apprendimento più informativi, massimizzando così l'efficienza del campione.

3. Replay Prioritario basato sull'Errore di Stima dell'Accuratezza:
   Per migliorare ulteriormente l'efficienza e calibrare la funzione di partizione, il framework introduce un meccanismo di replay. Un buffer memorizza coppie prompt-output con ricompense corrette ($r=1$). All'interno di questo buffer, i campioni sono prioritizzati in base all'entità del loro errore di stima dell'accuratezza: $| \frac{N_{correct}}{N} - \hat{p}_{old}(x) |$. Ciò garantisce che il modello riveda i campioni in cui la previsione della funzione di partizione diverge maggiormente dalla realtà osservata, affinando le stime di accuratezza e la normalizzazione della distribuzione target.

Crucialmente, sia la selezione dei prompt che la prioritizzazione del replay riutilizzano le informazioni generate durante l'addestramento standard delle GFlowNets, ammortizzando il sovraccarico computazionale nel processo di ottimizzazione esistente.

Contributi Chiave

• Insight Teorico: Il documento dimostra che la funzione di partizione della GFlowNet per il post-addestramento degli LLM è intrinsecamente legata alle accuratezze online, permettendone l'uso diretto come stimatore di accuratezza senza meccanismi ausiliari.
• Framework PACED-RL: L'introduzione di un nuovo framework di addestramento che utilizza queste stime di accuratezza per la selezione adattiva dei prompt e il replay prioritario per errore, migliorando l'efficienza del campione pur preservando i benefici di diversità delle GFlowNets.
• Validazione Empirica: Valutazioni estese su benchmark di ragionamento matematico e generazione di codice mostrano che PACED-RL ottiene guadagni di performance significativi rispetto a baseline solide (GRPO, FlowRL) e metodi di campionamento adattivo (DS, LILO, MoPPS), senza il sovraccarico computazionale associato all'oversampling.

Risultati Sperimentali
Gli esperimenti sono stati condotti sulla generazione di codice (dataset DeepCoder, DeepSeek-R1-Distill-Qwen-1.5B) e sul ragionamento matematico (dataset DeepScaleR, Qwen2.5-Math-1.5B/7B).

• Guadagni di Performance: Sul benchmark AIME 24, PACED-RL ha migliorato le prestazioni medie pass@1 del 29.1% rispetto a GRPO e del 40.0% rispetto a FlowRL utilizzando il modello Qwen2.5-Math-7B.
• Efficienza del Campione: PACED-RL ha raggiunto le migliori prestazioni di GRPO e FlowRL in tempi significativamente inferiori (ad esempio, 0.42× il tempo di addestramento su LiveCodeBench). A differenza dei metodi che richiedono oversampling (DS, LILO), PACED-RL ha mantenuto un rapporto di rollout di 1.0× rispetto all'addestramento standard, evitando le penalità di latenza associate alla generazione di campioni extra per la stima dell'accuratezza.
• Diversità: Misurata tramite metriche pass@k, PACED-RL ha costantemente superato le baseline, mostrando miglioramenti fino al 14.2% rispetto a GRPO e al 9.1% rispetto a FlowRL, confermando che il metodo preserva la diversità di output e la capacità di esplorazione.
• Studi di Ablazione: Gli autori hanno verificato che omettere il termine di divergenza KL nello stimatore di accuratezza ha un impatto trascurabile sulle dinamiche di campionamento, risparmiando al contempo una significativa quantità di calcolo. Hanno inoltre dimostrato che il metodo sposta naturalmente l'addestramento verso domande più difficili man mano che il modello migliora, prevenendo il collasso su un sottoinsieme ristretto del dataset.

Significato e Affermazioni
Il documento afferma che PACED-RL offre un modo "principiato e pratico" per potenziare l'addestramento di matching delle distribuzioni per gli LLM. Ripensando la funzione di partizione come un segnale di accuratezza online, il metodo consente un addestramento più efficiente dal punto di vista del campione senza l'onere computazionale di stimatori di accuratezza ausiliari o una generazione eccessiva di rollout. Gli autori collocano questo lavoro come un passo verso un post-addestramento RLVR più efficiente, suggerendo che il ruolo della funzione di partizione si estende oltre la normalizzazione per agire come un programmatore di difficoltà dinamico. Si suggerisce che lavori futuri esplorino strategie al momento dell'inferenza (ad esempio, auto-consistenza adattiva) e l'estensione del framework a contesti di GFlowNet asincroni.