Blockwise Advantage Estimation for Multi-Objective RL with Verifiable Rewards

Titolo: Blockwise Advantage Estimation for Multi-Objective RL with Verifiable Rewards

1. Il Problema: Interferenza degli Obiettivi e Assegnazione del Credito

Il paper affronta un limite fondamentale degli algoritmi di Reinforcement Learning (RL) attualmente popolari per il post-training dei Large Language Models (LLM), come il Group Relative Policy Optimization (GRPO).

In GRPO, l'algoritmo assegna un singolo valore di "advantage" (vantaggio) scalare a tutti i token di una generazione completa. Questo approccio presenta due problemi critici quando il compito è multi-obiettivo o strutturato (es. una soluzione matematica seguita da un'analisi della propria fiducia):

• Interferenza degli obiettivi: Se si combinano più segnali di reward (es. correttezza della soluzione + calibrazione della fiducia) in un unico scalare, i segnali possono contrastarsi, portando a un degrado delle prestazioni.
• Errata assegnazione del credito (Credit Assignment): Un reward ottenuto alla fine di una sequenza viene applicato uniformemente a tutti i token precedenti, anche se alcuni token (come quelli di una riflessione finale) non hanno influenzato direttamente quel reward.
• Problema del Baseline Condizionale: Per i blocchi di testo successivi, il reward dipende dal prefisso generato precedentemente. Calcolare un baseline accurato richiederebbe costosi "rollout" (nuove generazioni) da stati intermedi, rendendo l'addestramento computazionalmente proibitivo.

2. Metodologia: Blockwise Advantage Estimation (BAE)

Gli autori propongono una famiglia di metodi compatibili con GRPO che decompongono la generazione in blocchi contigui $X_1, X_2, \dots, X_K$, dove ogni blocco è associato a un obiettivo specifico $r_k$.

I pilastri della metodologia sono:

• Decomposizione in Blocchi: Invece di un unico vantaggio per l'intera sequenza, il metodo calcola vantaggi proxy separati $\hat{A}_k$ per ogni blocco e applica il vantaggio $\hat{A}_k$ solo ai token appartenenti a quel blocco.
• Outcome-Conditioned Baseline (OCB): Per risolvere il problema del baseline senza costi computazionali aggiuntivi, gli autori introducono l'OCB. Invece di usare la media del gruppo (che ignora che i prefissi sono diversi), l'OCB stratifica i campioni del gruppo in base a un esito intermedio discreto $o$ (ad esempio, se la soluzione precedente è corretta o errata). Il baseline per un blocco viene quindi calcolato come la media dei reward all'interno di quel sottogruppo specifico.
• Ottimizzazione Modulare: Il metodo permette di ottimizzare obiettivi diversi (es. correttezza tramite verifier e calibrazione tramite Brier Score) in modo indipendente, evitando la necessità di progettare funzioni di reward scalari complesse e "hand-tuned".

3. Contributi Chiave

1. Framework BAE: Un metodo critic-free (senza necessità di un modello critic separato) e compatibile con GRPO per l'RL multi-obiettivo su generazioni strutturate.
2. Efficienza Computazionale: L'introduzione dell'OCB permette di approssimare i valori degli stati intermedi utilizzando solo le statistiche del gruppo già esistente, eliminando la necessità di rollout Monte Carlo costosi.
3. Soluzione al Reward Hacking: Dimostrano che la decomposizione in blocchi previene comportamenti degeneri (come il modello che risponde "non lo so" per minimizzare l'errore di calibrazione) che si verificano quando si usa un reward combinato standard.

4. Risultati Sperimentali

Il metodo è stato testato su modelli Qwen (da 3B a 7B) in compiti di ragionamento matematico con stima della fiducia (self-confidence).

• Accuratezza e Calibrazione: L'approccio BAE con OCB è competitivo con lo stato dell'arte (RLCR) in termini di accuratezza matematica e supera o eguaglia le prestazioni in termini di Expected Calibration Error (ECE) e Brier Score.
• Test-Time Scaling (TTS): Il metodo preserva e potenzia i benefici del TTS. Poiché la fiducia del modello è ben calibrata, è possibile utilizzare strategie di selezione (come la Confidence-weighted majority) per migliorare l'accuratezza finale durante l'inferenza.
• Robustezza OOD: Il metodo mostra una forte resilienza anche su dataset fuori distribuzione (Out-of-Distribution) come GSM8K e AIME.
• Applicabilità: È stato dimostrato efficace anche in scenari di "Two-Attempt Refinement", dove il modello genera una prima soluzione e poi una seconda dopo una riflessione.

5. Significatività e Conclusioni

Il lavoro è significativo perché offre una "ricetta modulare" per l'addestramento di agenti e modelli di ragionamento complessi. Invece di cercare la combinazione perfetta di pesi per un unico reward scalare, i ricercatori possono ottimizzare segmenti specifici della generazione in modo isolato ma statisticamente corretto. Questo approccio è fondamentale per lo sviluppo di modelli capaci di gestire task a lungo raggio, agenti autonomi e processi di ragionamento multi-step dove la precisione di ogni fase è cruciale per il successo finale.