ReCo: Reweighting GRPO Against Distributional Concentration

Problema

Il Group Relative Policy Optimization (GRPO) è emerso come un metodo di apprendimento per rinforzo standard per il post-training dei modelli linguistici, particolarmente in domini con ricompense verificabili come la matematica e la programmazione. Sebbene il GRPO migliori l'efficienza di campionamento per piccoli budget di campionamento (basso $k$ in Pass@$k$), osservazioni recenti indicano che esso non riesce a superare, e spesso sottoperforma, il modello base quando il budget di campionamento è elevato (alto $k$). Ciò suggerisce una riduzione della copertura dei percorsi di ragionamento. Il documento identifica che il GRPO concentra la massa di probabilità sulle risposte che il modello base genera già con alta probabilità, riducendo così la diversità delle traiettorie di risoluzione dei problemi e limitando la capacità del modello di scoprire soluzioni corrette che risiedono in regioni a bassa probabilità della distribuzione di base.

Metodologia

Gli autori propongono ReCo (Reweighting GRPO Against Distributional Concentration), un metodo progettato per mitigare questa concentrazione modificando l'aggiornamento del GRPO a due livelli:

1. Ripesatura a livello di risposta (Response-Level Reweighting):
   Nel GRPO standard, le risposte ad alta probabilità vengono campionate più frequentemente all'interno di un gruppo di rollout. Poiché il vantaggio del gruppo è mediato equamente tra le risposte campionate, queste traiettorie frequenti accumulano una massa di gradiente aggregata maggiore semplicemente a causa della loro occorrenza ripetuta. ReCo affronta questo problema normalizzando il contributo del gradiente di ogni risposta rispetto al suo conteggio di occorrenza atteso all'interno del gruppo.

   • Il peso è definito come l'inverso del conteggio atteso: $w_i = \frac{1}{G \cdot \bar{\pi}_{\theta_{old}}(o_i | q)}$, dove $\bar{\pi}$ è la verosimiglianza normalizzata per lunghezza.
   • Questo impedisce alle risposte frequenti di dominare il gradiente a livello di gruppo, garantendo che l'aggiornamento non rafforzi in modo sproporzionato le traiettorie che sono già probabili sotto la vecchia policy.

2. Ripesatura a livello di token (Token-Level Reweighting):
   Il GRPO standard utilizza un rapporto di importanza ($\pi_\theta / \pi_{\theta_{old}}$) per scalare i gradienti a livello di token. Ciò crea un ciclo di feedback in cui i token che diventano più probabili sotto la policy corrente ricevono aggiornamenti più grandi, rafforzando ulteriormente i punti decisionali saturati e riducendo le traiettorie alternative.

   • ReCo sostituisce lo standard rapporto di importanza con un rapporto basato sulla varianza: $r_{var} = \frac{p_\theta(1 - p_\theta)}{p_{old}(1 - p_{old})}$.
   • Il termine $p(1-p)$ rappresenta la varianza di Bernoulli. Questo rapporto assegna scale di aggiornamento maggiori ai punti decisionali non saturati (dove $p$ non è né vicino a 0 né a 1, il che significa che le scelte di token alternative rimangono plausibili) e riduce il peso delle scelte saturate (dove $p \approx 1$).
   • Questo meccanismo incoraggia gli aggiornamenti nei punti decisionali in cui la policy non è ancora collassata su una singola scelta, preservando così la diversità del ragionamento.

Contributi Chiave

• Analisi della concentrazione del GRPO: Il documento traccia la riduzione del Pass@$k$ per grandi valori di $k$ a due meccanismi specifici nell'aggiornamento del GRPO: la dominanza delle risposte frequenti a livello di gruppo e l'amplificazione dei token saturati a livello di token.
• Algoritmo ReCo: L'introduzione di uno schema di ripesatura a due livelli che normalizza i contributi delle risposte rispetto all'occorrenza attesa e sostituisce il rapporto di importanza del token con una metrica basata sulla varianza.
• Validazione Empirica: Valutazione estesa su diverse famiglie di modelli (Qwen2.5-Math, Llama-3.1) e scale (1.5B, 7B, 8B) su cinque benchmark di ragionamento matematico (AIME 2024/2025, AMC 2023, OlympiadBench, MATH 500).

Risultati

• Performance Pass@$k$: ReCo supera costantemente il GRPO nel Pass@$k$ per valori elevati di $k$ (ad esempio, Pass@64), indicando una migliore copertura delle traiettorie di ragionamento corrette. Per valori piccoli di $k$, ReCo rimane paragonabile al GRPO, mantenendo l'efficienza di campionamento.
• Recupero del Modello Base: Sul backbone Llama-3.1-8B-Instruct, dove si è osservato che il GRPO standard scendeva al di sotto delle prestazioni del modello base nel Pass@$k$, ReCo riduce significativamente questo degrado e persino supera il modello base su determinati benchmark (AIME 2024/2025).
• Metriche di Diversità: L'analisi della dinamica di addestramento mostra che mentre il GRPO perde rapidamente diversità (diminuendo il Distinct-2, aumentando il Self-BLEU), ReCo mantiene la diversità del rollout vicina a quella del modello base durante tutto l'addestramento.
• Varianza del Token: ReCo mantiene una varianza di Bernoulli a livello di token sostanzialmente più alta nei segmenti iniziali e medi delle risposte rispetto al GRPO, indicando che i percorsi di ragionamento rimangono non saturati per un tempo maggiore.
• Generalizzabilità: Il metodo migliora anche il Pass@$k$ nei benchmark di ragionamento sul codice (HumanEval, HumanEval+, LiveCodeBench), suggerendo che il problema della concentrazione non è limitato al ragionamento matematico.
• Compatibilità: ReCo può essere combinato con altre varianti di GRPO (ad esempio, DAPO, Regola 80/20) per ottenere ulteriori miglioramenti.

Significato

Il documento afferma che ReCo affronta i meccanismi interni del GRPO che portano alla concentrazione distributiva senza richiedere obiettivi di diversità esterni o regolarizzatori. Ripesando l'aggiornamento stesso in base alla probabilità di occorrenza e alla saturazione dei punti decisionali, ReCo preserva un insieme più ampio di percorsi di ragionamento corretti. Ciò consente al modello di mantenere i vantaggi dell'efficienza di campionamento dell'apprendimento per rinforzo espandendo al contempo la sua copertura dello spazio delle soluzioni, mitigando efficacementamente il "collasso della diversità" osservato nell'addestramento standard del GRPO. Il lavoro suggerisce che i limiti del RLVR nell'espandere la capacità di ragionamento sono in parte dovuti alla dinamica di ottimizzazione dell'aggiornamento della policy stessa, che può essere corretta attraverso una ripesatura mirata.