Sintesi Tecnica: Sulle stime del vantaggio per le politiche Max@K

Definizione del Problema

L'apprendimento per rinforzo con ricompense verificabili (RLVR) è un paradigma dominante per il post-training di grandi modelli linguistici (LLM) su compiti di ragionamento. Tuttavia, questi contesti si affidano tipicamente a ricompense sparse basate sull'esito (ad esempio, un segnale binario corretto/errato), il che rende difficile l'esplorazione in ampi spazi combinatori. Sebbene metodi come la regolarizzazione dell'entropia e i bonus basati sul conteggio esistano, essi introducono obiettivi ausiliari che richiedono una calibrazione attenta.

Un approccio complementare consiste nell'ottimizzare direttamente gli obiettivi di inferenza come il pass@K (la probabilità che almeno una delle $K$ completazioni campionate sia corretta) e la sua generalizzazione continua, il max@K (il valore atteso del massimo della ricompensa tra $K$ campioni). Diversi lavori recenti hanno proposto stimatori di gradiente di politica (PG) non distorti per questi obiettivi. Tuttavia, le relazioni tra questi stimatori rimangono poco chiare. Nello specifico, gli stimatori esistenti utilizzano spesso segnali, baseline e normalizzazioni differenti. Un problema critico identificato in questo lavoro è che i principali stimatori (ad esempio, quelli basati sul Miglioramento Atteso o Expected Improvement) sono non distorti come stimatori PG, ma producono vantaggi non centrati. Un vantaggio non centrato (con una media non nulla) corrisponde a una baseline sottostimata, il che può indirizzare l'ottimizzazione verso azioni subottimali e aumentare la varianza del gradiente.

Metodologia

Il documento propone MaxPO (Max@K Policy Optimization), un metodo che introduce una baseline Leave-Two-Out (L2O) per correggere la stima del vantaggio per gli obiettivi max@k.

1. Analisi degli stimatori esistenti

Gli autori rivisitano lo stimatore del Miglioramento Atteso (EI) utilizzato in lavori precedenti (ad esempio, Walder e Karkhanis [54]). Dimostrano che, sebbene lo stimatore basato su EI sia uno stimatore PG non distorto, il segnale di vantaggio risultante $s(a)$ è strettamente non negativo (a causa dell'operazione ReLU nella definizione di EI). Di conseguenza, $E[s(a)] > 0$, il che significa che il vantaggio non è centrato. Questa mancanza di centratura lascia spazio alla riduzione della varianza.

2. La baseline Leave-Two-Out (L2O)

Per ottenere un vantaggio esattamente centrato preservando l'imparzialità (unbiasedness) del PG, gli autori introducono la baseline L2O.

• La Sfida: Una standard baseline Leave-One-Out (L1O) fallisce nel contesto max@K perché la ricompensa di un campione $i$ può influenzare implicitamente il calcolo del vantaggio di altri campioni $j$ (poiché $i$ agisce come un potenziale comparatore nell'operazione max). Questa dipendenza viola la condizione di indipendenza richiesta per la stima PG non distorta.
• La Soluzione: La baseline L2O $b_{-i}^{L2O}$ è calcolata mediando i punteggi di Miglioramento Atteso di tutti gli altri campioni $j$, dove ogni punteggio è calcolato su un sottoinsieme di comparatori che esclude sia il campione target $i$ sia il campione $j$ stesso.
• Proprietà:
  • Indipendenza: Escludendo $i$ dall'insieme dei comparatori per ogni $j$, la baseline $b_{-i}^{L2O}$ è indipendente dalla ricompensa $r_i$, garantendo che lo stimatore PG rimanga non distorto.
  • Centratura: Gli autori dimostrano che il vantaggio medio realizzato per batch è esattamente zero per ogni batch campionato: $\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$.
  • Efficienza: Sebbene un calcolo ingenuo sarebbe $O(B^3)$, gli autori derivano un'implementazione vettorizzata utilizzando formulazioni basate su ReLU che calcola la baseline L2O in tempo $O(B^2)$, rendendola praticabile per l'addestramento di LLM accelerato da GPU.

3. Vantaggio Canonico a Batch Finito

Il documento deriva un vantaggio canonico a batch finito per max@K, definito come $u_i - v_i$:

• $u_i$: Il valore atteso condizionale max@K di un gruppo contenente la risposta $i$.
• $v_i$: Il valore atteso max@K "leave-one-out" calcolato da gruppi che escludono $i$.
  Gli autori mostrano che lo stimatore MaxPO (EI + L2O) recupera esattamente questo segnale canonico ($s_i - b_{-i}^{L2O} = u_i - v_i$). Ciò fornisce una visione unificata degli stimatori esistenti, classificandoli come "non centrati" (ad esempio, solo EI/PKPO) e "(normalizzati) canonici" (ad esempio, MaxPO e certe varianti analitiche di pass@K).

4. Integrazione nell'addestramento di LLM

MaxPO è integrato in algoritmi di RL basati su gruppi (come GRPO). Per ogni prompt, viene campionato un gruppo di $G$ output. La baseline L2O viene calcolata all'interno di ogni gruppo per generare vantaggi a livello di sequenza, che vengono poi propagati ai loss a livello di token.

Contributi Chiave

1. Identificazione dei Vantaggi Non Centrati: Il documento identifica che i principali stimatori PG per max@K (specificamente quelli basati su EI) sono non distorti come stimatori di gradiente, ma non producono vantaggi centrati, portando a proprietà di varianza subottimali.
2. Baseline Leave-Two-Out (L2O): Un nuovo design di baseline che preserva l'imparzialità del PG nel setting non lineare di max@K, garantendo al contempo che il vantaggio realizzato per batch abbia media zero.
3. Unificazione Canonica: La derivazione del vantaggio canonico a batch finito ($u_i - v_i$) e la prova che MaxPO recupera questo segnale, offrendo un quadro unificato per confrontare e categorizzare gli esistenti stimatori pass@K/max@K.
4. Implementazione Efficiente: Un algoritmo vettorizzato $O(B^2)$ per calcolare la baseline L2O, rendendo possibile l'applicazione pratica nel post-training di LLM.

Risultati Sperimentali

Gli autori validano il metodo attraverso esperimenti controllati su bandit/maze e compiti di ragionamento con LLM.

• Riduzione della Varianza: Negli ambienti bandit, la baseline L2O riduce significativamente la varianza del gradiente rispetto ai grezzi stimatori EI, con un effetto che cresce all'aumentare della dimensione dello spazio delle azioni. Nell'addestramento di LLM (Llama-3.2-3B-Instruct), il metodo ha ridotto la proxy della varianza del gradiente del 77,4% rispetto a PKPO (il baseline EI-only).
• Guadagni di Performance: Su cinque benchmark di ragionamento matematico (AIME24, AIME25, AMC23, MATH500, Minerva), MaxPO ha superato i forti baseline inclusi PKPO, GRPO ed Entropy-Adv.
  • Su Qwen2.5-Math-7B, MaxPO ha migliorato il pass@256 medio del compito del 5,2% rispetto a PKPO.
  • Su Llama-3.2-3B-Instruct, ha migliorato il pass@256 del 2,4% rispetto a PKPO.
• Scaling del Calcolo di Inferenza: I guadagni di performance di MaxPO sono stati più pronunciati ai livelli di calcolo di inferenza più elevati (k maggiore in pass@k), suggerendo che il metodo aiuta la policy a utilizzare meglio i campioni aggiuntivi in fase di test.

Significato e Rivendicazioni

Il documento sostiene che MaxPO fornisce una base principata per progettare e confrontare stimatori di gradiente di politica per obiettivi di inferenza come max@K. Identificando la forma del vantaggio canonico e correggendo il bias degli stimatori esistenti, il lavoro offre un metodo per migliorare l'efficienza e la stabilità dei modelli di ragionamento in fase di post-training.

Gli autori notano con modestia che la loro analisi si concentra sulla stima on-policy a batch finito sotto campionamento i.i.d. Riconoscono limitazioni riguardanti l'addestramento off-policy, il campionamento adattivo e i contesti con generazioni correlate come direzioni per il lavoro futuro. Inoltre, evidenziano che, sebbene il metodo migliori l'ottimizzazione, il miglioramento dell'ottimizzazione di obiettivi di ragionamento con campionamento ripetuto comporta impatti più ampi, potenzialmente rafforzando modelli che potrebbero essere utilizzati impropriamente per assistenza ingannevole o risoluzione automatizzata di problemi.