Future-KL Regularized GRPO: Process-Level Credit Assignment from $f$-Divergence Regularization

Riepilogo Tecnico: GRPO con Regularizzazione Future-KL

Enunciazione del Problema
L'Ottimizzazione della Politica Relativa di Gruppo (GRPO) è emersa come paradigma dominante per l'apprendimento per rinforzo (RL) senza critico nel post-addestramento dei Modelli Linguistici di Grande Dimensione (LLM), in particolare per il ragionamento matematico. La GRPO elimina la necessità di una funzione di valore appresa normalizzando le ricompense degli esiti all'interno di un gruppo di risposte campionate. Tuttavia, le implementazioni standard della regolarizzazione KL nella GRPO trattano tipicamente la penalità come una penalità locale a livello di token sul lato della perdita (ad esempio, utilizzando lo stimatore $k_3$).

Gli autori identificano un difetto fondamentale in questo approccio: ignora l'assegnazione del credito a livello di processo indotta dalla generazione autoregressiva. In una traiettoria autoregressiva, un token iniziale non solo sostiene il proprio costo di regolarizzazione locale, ma altera anche la distribuzione dei prefissi futuri, influenzando di conseguenza i costi KL di tutti i token successivi. I termini di perdita locali standard differenziano solo la penalità del token campionato, omettendo il "ritorno futuro di regolarizzazione causale" (future-regularization return-to-go).

Inoltre, il documento evidenzia un'incompatibilità strutturale tra la normalizzazione di gruppo della GRPO e le formulazioni standard KL-in-reward. A differenza degli obiettivi di ricompensa lineari (ad esempio, MaxEnt RL) dove ricompensa e regolarizzatore possono essere fusi, la normalizzazione di gruppo della GRPO induce un'utilità non lineare a livello di prompt (specificamente, una trasformazione arcsin per ricompense binarie). Questa non linearità significa che aggiungere penalità KL alle ricompense prima della normalizzazione distorce l'obiettivo implicito, collassando le informazioni KL a livello di token in un singolo scalare a livello di traiettoria e fallendo nel preservare il segnale di gradiente corretto.

Metodologia
Il documento deriva il gradiente della politica on-policy per obiettivi di tipo GRPO sotto regolarizzazione generale $f$-divergenza a livello di token. La chiave dell'insight teorico è che il gradiente si scompone in due termini distinti con proprietà di scaling diverse:

1. Il Termine di Ricompensa: Ereditato dall'utilità non lineare $\ell(p)$, questo termine porta un guadagno dipendente dal prompt $\ell'(p)$. Per ricompense binarie, questo recupera il vantaggio GRPO standardizzato.
2. Il Termine di Regolarizzatore: Questo termine è indipendente da $\ell'$ e si scompone in una penalità locale a livello di token più un ritorno futuro di regolarizzazione causale.

Sulla base di questa derivazione, gli autori propongono l'Ottimizzazione della Politica con Regularizzazione Future-KL (FRPO). L'algoritmo modifica l'aggiornamento standard della GRPO in modo disaccoppiato:

1. Calcolare il vantaggio dell'esito normalizzato per gruppo ($\hat{A}^{GRPO}$) strettamente dalle ricompense originali del verificatore, preservando l'utilità non lineare implicita.
2. Costruire un termine di correzione future-KL a livello di token calcolando la somma cumulativa inversa dei log-rapporti per-token ($\delta_{i,j} = \log \frac{\pi_\theta}{\pi_{ref}}$) dal token corrente fino alla fine della traiettoria.
3. Aggiungere questo termine future-KL al vantaggio dopo la costruzione del vantaggio, piuttosto che modellare la ricompensa prima della normalizzazione.

La regola di aggiornamento per un token $t$ nella risposta $i$ è:
$$\hat{A}^{FRPO}_{i,t} = \text{sg}\left( \hat{A}^{GRPO}_i - \beta \sum_{j=t}^{T_i} \delta_{i,j} \right)$$
dove $\text{sg}$ denota stop-gradient. Questo approccio non richiede un critico, né passaggi aggiuntivi del modello, e mantiene lo stesso sovraccarico computazionale della GRPO standard.

Contributi Chiave

1. Identificazione dell'Utilità Non Lineare: Il documento dimostra che la normalizzazione di gruppo della GRPO induce un'utilità non lineare ($\ell(p) = 2 \arcsin \sqrt{p}$ per ricompense binarie), rendendo ricompensa e KL non fusibili. Aggiungere KL alle ricompense prima della normalizzazione distorce strutturalmente l'obiettivo.
2. Derivazione del Gradiente On-Policy: Gli autori derivano il gradiente on-policy corretto per la regolarizzazione $f$-divergenza a livello di token, mostrando che ciò richiede un termine di ritorno futuro di regolarizzazione assente nelle implementazioni locali sul lato della perdita.
3. Algoritmo FRPO: Propongono FRPO, un aggiornamento semplice e disaccoppiato che preserva il vantaggio di esito della GRPO mentre inietta l'assegnazione del credito future-KL necessaria. Questo metodo separa il ranking a livello di traiettoria dalla regolarizzazione KL a livello di processo.
4. Interpretazione di Ottimalità: La struttura future-KL è mostrata allinearsi alle condizioni di ottimalità dell'RL regolarizzato per entropia, dove la politica ottimale soddisfa una forma di ricompensa di processo che coinvolge i costi KL futuri.

Risultati Sperimentali
Gli autori valutano FRPO su compiti di ragionamento matematico utilizzando modelli base Qwen3 (1.7B, 4B e 30B-A3B) addestrati sul dataset DAPO-Math-17k.

• Prestazioni: Sul modello 30B-A3B, FRPO migliora il pass@16 di quasi il 5% di guadagno assoluto rispetto alla baseline più forte (GRPO standard con KL sul lato della perdita) sul benchmark MATH500. Supera anche le baseline su AIME24, AIME25 e altri benchmark di livello olimpico.
• Stabilità e Deriva: FRPO mantiene un'entropia finale più alta e una deriva della politica inferiore (misurata tramite PPO-KL) rispetto alle baseline KL sul lato della perdita. Ciò suggerisce che il metodo elicita pattern di ragionamento a catena di pensiero (CoT) lunghi senza una deviazione eccessiva dalla politica di riferimento.
• Ablazione: Gli esperimenti confermano che il KL inverso (mode-seeking) performa meglio del KL diretto (mode-covering) per questo compito. Il documento nota anche che l'uso dello stimatore $k_1$ (log-rapporto) è preferito rispetto a $k_3$ per il termine future-KL, poiché la non negatività di $k_3$ crea un segnale di penalità unilaterale che ostacola l'assegnazione del credito.
• Modalità di Fallimento: Il documento osserva che gli approcci "KL-in-reward" (aggiungere KL alle ricompense prima della normalizzazione) portano al collasso del modello o al degrado delle prestazioni, validando l'affermazione teorica che tale fusione è strutturalmente errata sotto l'utilità non lineare della GRPO.

Significato
Il documento sostiene che il significato di questo lavoro risiede nella correzione di un errore di implementazione sottile ma critico in un framework RL ampiamente utilizzato. Spostando la prospettiva da "penalità di perdita locale" a "assegnazione del credito a livello di processo", gli autori mostrano che la GRPO può essere migliorata senza introdurre funzioni di valore complesse o costi computazionali aggiuntivi. Il metodo FRPO proposto fornisce una modifica teoricamente fondata e leggera che potenzia le capacità di ragionamento preservando al contempo le proprietà di stabilità e generalizzazione del modello base. Il lavoro sottolinea che in contesti autoregressivi, la regolarizzazione è intrinsecamente un segnale a livello di processo, e ignorare il costo futuro (cost-to-go) porta a un'assegnazione del credito subottimale.