Riepilogo Tecnico: Ottimizzazione Adattiva dei Rollout tra Epoch per il Post-Training di RL (CERO)

1. Formulazione del Problema

Il post-training di Large Language Model (LLM) tramite Reinforcement Learning (RL), in particolare utilizzando algoritmi come GRPO, si basa tipicamente sul campionamento di un numero fisso di rollout per ogni prompt. Questa strategia di allocazione uniforme ignora l'eterogeneità dei segnali di addestramento forniti da diversi prompt. Alcuni prompt sono banali (generando successo o fallimento costanti con bassa varianza), mentre altri sono troppo difficili o già padroneggiati, offrendo rendimenti marginali decrescenti. Di conseguenza, una parte significativa del budget globale di rollout viene sprecata su prompt con limitato valore di apprendimento.

Il documento formula questo scenario come un problema di allocazione delle risorse online cross-epoch sotto un budget globale fisso $B$. A differenza di metodi precedenti che eseguono la riallocazione all'interno del batch o il filtraggio post-generazione, questo lavoro affronta l'allocazione dei rollout sia tra i prompt che tra le epoch di addestramento. L'obiettivo è massimizzare l'utilità cumulativa dei campioni raccolti, dove l'utilità è soggetta a rendimenti decrescenti a livello di prompt. Il problema è matematicamente definito come la massimizzazione della somma delle funzioni di utilità specifiche per prompt $U_i(\sum_k N_{i,k})$ soggetta al vincolo di budget globale $\sum_{i,k} N_{i,k} \le B$, dove $N_{i,k}$ è il numero di rollout assegnati al prompt $i$ all'epoch $k$.

2. Metodologia: CERO

Gli autori propongono CERO (Cross-Epoch Adaptive Rollout Optimization), un algoritmo bayesiano primal-dual online progettato per allocare adattivamente i rollout in base alla stima dell'informatività del prompt.

2.1 Informatività e Design dell'Utilità

• Modellazione della Posterior Beta: Per ogni prompt $x_i$, CERO mantiene una distribuzione posteriore Beta sulla sua latente probabilità di successo $p_i$, aggiornata in base alla cronologia dei rollout $H_i^k$.
• Punteggio di Informatività ($q_{k,i}$): Invece di utilizzare una stima puntuale della probabilità di successo, CERO stima la varianza di Bernoulli attesa: $q_{k,i} = E[p_i(1-p_i) | H_i^k]$. Questo punteggio è massimizzato quando $p_i \approx 0.5$ (difficoltà intermedia) e svanisce per prompt facili ($p_i \approx 1$) o difficili ($p_i \approx 0$), identificando efficacemente i prompt che offrono i segnali di apprendimento più incerti e quindi più informativi.
• Funzione di Utilità: Una funzione di utilità concava e saturante è definita come $U_i(n) = 1 - \exp(-\eta q_i n)$, dove $n$ è il numero cumulativo di rollout. Ciò cattura i rendimenti decrescenti: i primi rollout su un prompt informativo sono altamente preziosi, ma il valore marginale diminuisce man mano che vengono raccolti più campioni.

2.2 Ottimizzazione tramite Riformulazione Fenchel-Dual

L'accoppiamento temporale della funzione di utilità (che dipende dall'allocazione cumulativa) rende difficile l'ottimizzazione online diretta. CERO affronta questo problema tramite una riformulazione Fenchel-dual:

• Il problema viene trasformato in un obiettivo duale che coinvolge le variabili duali a livello di prompt $\theta_i$ e un moltiplicatore di budget globale $\mu$.
• Questa riformulazione decompone il problema temporalmente accoppiato in sottoproblemi di allocazione per ogni epoch coordinati dalle variabili duali.

2.3 Algoritmo Primal-Dual Online

CERO opera come uno strato "drop-in" sopra le pipeline di RL esistenti (ad es., GRPO):

1. Decisione di Allocazione: Ad ogni epoch, per ogni prompt, l'algoritmo risolve un sottoproblema lineare. Un prompt riceve il massimo numero di rollout consentiti ($N_{max}$) se il suo valore marginale stimato (variabile duale $\theta_i$) supera il prezzo attuale del budget globale ($\mu$); altrimenti, riceve zero.
2. Aggiornamento del Feedback: Dopo la generazione delle risposte e l'osservazione dei reward, la posteriore Beta per il prompt viene aggiornata e il punteggio di informatività $q_{k,i}$ viene rinfrescato.
3. Aggiornamenti delle Variabili Duali:
   • A livello di prompt ($\theta$): Aggiornato tramite gradiente discendente online proiettato per bilanciare i rollout allocati rispetto all'utilità implicata dalla corrente stima di informatività.
   • A livello di budget ($\mu$): Aggiornato in base alla discrepanza tra il consumo effettivo di rollout e il target di tasso di budget. Se il consumo supera il tasso, $\mu$ aumenta, rendendo l'allocazione futura più selettiva; se è inferiore, $\mu$ diminuisce.

3. Contributi Chiave

1. Formulazione del Problema: Il documento è il primo a formulare l'allocazione dei rollout nel post-training di RL per LLM come un problema di ottimizzazione online a budget fisso cross-epoch con rendimenti decrescenti specifici per prompt, accoppiando esplicitamente le decisioni attraverso l'orizzonte di addestramento tramite un budget globale condiviso.
2. Design dell'Algoritmo: Introduce CERO, un algoritmo bayesiano primal-dual che utilizza stime posteriori del valore del rollout a livello di prompt per guidare il budgeting. Utilizza la varianza di Bernoulli attesa della posteriore come proxy per la apprendibilità e impiega una riformulazione Fenchel-dual per gestire la non-separabilità temporale.
3. Validazione Teorica ed Empirica:
   • Teoria: Sotto utilità di prompt fisse, gli autori dimostrano un limite di regret $O(\sqrt{K})$ rispetto a un benchmark di allocazione offline, dove $K$ è il numero di epoch.
   • Esperimenti: Estesi esperimenti su compiti di ragionamento matematico (AIME, AMC) su quattro LLM open-weight (con parametri da 1.5B a 7B) dimostrano che CERO supera costantemente il vanilla GRPO.

4. Risultati Sperimentali

Gli autori hanno valutato CERO sul dataset DAPO-Math-17K utilizzando DeepSeek-R1-Distill-1.5B, Qwen3-4B (Base e Instruct) e Qwen2.5-Math-7B.

• Performance: CERO ha ottenuto miglioramenti costanti rispetto a vanilla GRPO su tutti i modelli e i benchmark. I guadagni medi sono variati da +1.53 a +6.04 punti a seconda del modello.
• Analisi dell'Efficienza: Il "Effective Prompt Ratio" (la frazione di prompt in un batch che contiene sia risposte di successo che di fallimento, che generano gradienti non nulli) è stato significativamente più alto per CERO rispetto a GRPO.
  • Nel baseline GRPO, l'effective prompt ratio è diminuito drasticamente con il progredire dell'addestramento, indicando che il modello aveva padroneggiato molti prompt, portando a gruppi di rollout degeneri.
  • CERO ha mantenuto un effective prompt ratio più elevato spostando adattivamente il budget dai prompt padroneggiati verso quelli con maggiore incertezza.
• Overhead: Il meccanismo di aggiornamento dell'allocazione introduce un overhead computazionale trascurabile (meno di un secondo per aggiornamento).

5. Significato e Rivendicazioni

Il documento sostiene che l'allocazione adattiva del budget dei rollout sia una direzione fondamentale ed efficace per migliorare l'efficienza campionaria nel post-training di RL. Allontanandosi dall'allocazione uniforme, CERO assicura che il budget di generazione fisso sia speso su prompt che forniscono i segnali di addestramento più informativi.

Gli autori posizionano CERO come uno strato di raccolta dati adattivo "drop-in" che modifica quali prompt ricevono i rollout e quanti, senza alterare l'obiettivo di policy-gradient sottostante. Il lavoro evidenzia che l'inefficienza degli attuali post-training di RL non deriva solo dalla regola di aggiornamento della policy, ma dalla sottottimale strategia di raccolta dati.

Limitazioni note dagli autori:

• La garanzia teorica di regret assume utilità di prompt fisse, mentre l'implementazione utilizza stime di plugin dinamiche per epoch della posteriore. Pertanto, l'analisi non cattura completamente la non-stazionarietà indotta dagli aggiornamenti della policy e dai refresh della posteriore.
• La validazione empirica è attualmente limitata ai compiti di ragionamento matematico; la validazione su altri domini e modelli di reward è suggerita come lavoro futuro.