Sintesi Tecnica: rePIRL: Learn PRM with Inverse RL for LLM Reasoning

Definizione del Problema

Nel ragionamento dei Large Language Models (LLM), i Process Reward Models (PRM) vengono utilizzati per fornire ricompense intermedie a livello di token, con l'obiettivo di migliorare l'efficienza dell'addestramento, ridurre la varianza e prevenire il reward hacking. Tuttavia, i metodi esistenti per apprendere PRM efficaci affrontano limitazioni significative:

• Metodi Offline: Gli approcci supervisionati (ad es., Uesato et al., 2022; Lightman et al., 2023) e i lavori di tipo DPO richiedono l'accesso ad annotazioni a livello di token o etichette di preferenza da traiettorie esperte, che sono costose da ottenere. I metodi basati su Monte Carlo Tree Search (MCTS) (ad es., Wang et al., 2024) richiedono l'esecuzione di una policy esperta per generare traiettorie, il che è computazionalmente costoso e impraticabile se la policy esperta non è disponibile.
• Metodi Online: Approcci come PRIME (Cui et al., 2025a) si basano su forti assunzioni riguardanti le funzioni di ricompensa e il modello di policy in addestramento. I metodi di ricompensa intrinseca (ad es., Zhao et al., 2025b) utilizzano l'entropia o la confidenza del modello come ricompense, ma soffrono di "collasso dell'entropia", dove le prestazioni degradano nelle fasi finali dell'addestramento a causa di un eccessivo comportamento di ricerca del modo (mode-seeking).

La sfida centrale è apprendere un PRM efficace con assunzioni minime: specificamente, senza richiedere annotazioni di ricompensa a livello di token, etichette di preferenza, accesso alla policy esperta o forti assunzioni sulla forma della funzione di ricompensa.

Metodologia: Framework rePIRL

Gli autori propongono rePIRL, un framework ispirato all'Inverse Reinforcement Learning (IRL) che apprende i PRM esclusivamente da traiettorie esperte. Il metodo modella il ragionamento multi-step degli LLM come un processo decisionale di Markov (MDP) a livello di token e impiega un processo di apprendimento duale che aggiorna la policy e il PRM in modo intercambiabile.

1. Fondamento Teorico

• Costruzione dell'Obiettivo: Seguendo il classico framework IRL, gli autori definiscono una variabile nascosta $o_t \in \{0, 1\}$ che indica se un'azione è campionata dalla policy esperta. Modellano la probabilità che un'azione sia campionata dall'esperto utilizzando una funzione basata sull'energia: $p(o_t=1|s_t, a_t) \propto \exp(r_\phi(s_t, a_t))$.
• Massimizzazione della Verosimiglianza: La funzione di ricompensa $r_\phi$ viene appresa massimizzando la verosimiglianza delle traiettorie esperte. Ciò porta a un obiettivo che coinvolge una funzione di partizione $z(\phi)$, che è intrattabile per gli LLM a causa dei vasti spazi di stato e azione.
• Trattabilità tramite Campionamento: Per risolvere l'intrattabilità della funzione di partizione, gli autori applicano un trucco del gradiente e il campionamento per importanza (importance sampling). Invece di campionare dalla policy ottimale (che è sconosciuta), campionano traiettorie $\bar{\tau}$ dalla policy corrente $\pi_\theta$. Il gradiente è approssimato utilizzando pesi di importanza $w = \exp(r_\phi(\bar{\tau})) / \pi_\theta(\bar{\tau})$.
• Ottimizzazione della Policy: Per prevenire il collasso del modello (policy deterministiche), la policy viene aggiornata utilizzando un obiettivo di Maximum Entropy RL (Haarnoja et al., 2018), che bilancia la massimizzazione della ricompensa con la regolarizzazione dell'entropia. Questo viene risolto utilizzando algoritmi come RLOO, PPO o GRPO.

2. Algoritmo di Apprendimento

Il ciclo di addestramento alterna l'aggiornamento del PRM e della policy:

1. Rollout: Generare rollout della policy e recuperare dimostrazioni esperte.
2. Calcolo della Ricompensa: Calcolare le ricompense di esito (outcome rewards) e le ricompense di processo medie.
3. Aggiornamento del PRM: Aggiornare il modello di ricompensa utilizzando una funzione di perdita che contrasta le traiettorie esperte contro i rollout della policy, pesate dai punteggi di campionamento per importanza. Crucialmente, gli autori utilizzano le ricompense di processo medie invece delle ricompense cumulative per mitigare il reward hacking.
4. Aggiornamento della Policy: Aggiornare il modello di policy utilizzando le ricompense di processo apprese combinate con le ricompense di esito (se disponibili), impiegando obiettivi di massima entropia.

3. Unificazione dei Metodi Esistenti

Un contributo teorico chiave è la dimostrazione che rePIRL può unificare vari metodi esistenti di apprendimento dei PRM (DPO, DQO, Math-Shepherd, PRIME) sotto un unico framework introducendo specifiche assunzioni aggiuntive. Questo posiziona rePIRL come il metodo più generale, richiedendo il minor numero di assunzioni.

Contributi Chiave

1. Framework Innovativo: Introduzione di rePIRL, un framework ispirato all'IRL per l'apprendimento dei PRM da traiettorie esperte senza annotazioni a livello di token, etichette di preferenza o accesso alla policy esperta.
2. Teoria Unificata: Un'analisi teorica che mostra come rePIRL unifichi i metodi state-of-the-art per l'apprendimento offline e online dei PRM (inclusi DPO, DQO, PRIME e approcci basati su MCTS), chiarendo le loro assunzioni sottostanti.
3. Algoritmo Scalabile: Sviluppo di un algoritmo di apprendimento scalabile progettato per gli LLM, che utilizza il campionamento per importanza per gestire le funzioni di partizione intrattabili e il Maximum Entropy RL per garantire la stabilità dell'addestramento.
4. Validazione Empirica: Valutazione completa su benchmark matematici e di coding che dimostra prestazioni superiori rispetto ai baseline.

Risultati Sperimentali

Gli autori hanno valutato rePIRL utilizzando Qwen2.5-3B-Instruct e Qwen3-4B-Base su sette benchmark standard (AIME-2024, AMC, Math-500, MinervaMath, OlympiadBench, Leetcode, LiveCodeBench).

• Prestazioni: rePIRL ha raggiunto risultati state-of-the-art tra i modelli di scala 3B e 4B. Ad esempio, su Qwen3-4B-Base, ha raggiunto una precisione media del 41,6% sui benchmark matematici e del 39,8% sui benchmark di coding, superando solidi baseline come PRIME, MCTS e il vanilla RLOO.
• Confronto: Ha mostrato miglioramenti assoluti del 3,2% (math) e del 4,3% (coding) rispetto al vanilla RLOO con Qwen3-4B-Base.
• Applicazioni:
  • Test-time Training: Quando addestrato esclusivamente su segnali PRM senza ricompense di esito, rePIRL ha eguagliato le prestazioni dei modelli addestrati con ricompense di esito verificabili.
  • Test-time Scaling: L'uso del PRM come verificatore per selezionare i migliori rollout ha prodotto miglioramenti monotoni delle prestazioni, superando il majority voting.
  • Problemi Difficili: Il PRM ha fornito segnali precoci informativi per problemi difficili dove le ricompense di esito erano scarse, portando a una convergenza più rapida rispetto all'addestramento basato solo sull'esito.
• Efficienza: rePIRL ha introdotto solo un overhead di tempo di addestramento di circa 1,4x rispetto all'RL basato sull'esito (ad es., RLOO), evitando al contempo la necessità di reti di valore separate.

Significato e Rivendicazioni

Il paper sostiene di essere il primo lavoro ad generalizzare l'Inverse Reinforcement Learning per il ragionamento degli LLM e il primo ad unificare i metodi esistenti di apprendimento dei PRM in un singolo framework.

• Assunzioni Minime: Gli autori sottolineano che rePIRL opera sotto le assunzioni minime necessarie, rendendolo più generalizzabile rispetto ai metodi che richiedono annotazioni a livello di token o l'accesso specifico alla policy esperta.
• Stabilità: Il framework affronta il problema del collasso dell'entropia comune nei metodi di ricompensa intrinseca sfruttando il Maximum Entropy RL, garantendo una dinamica di addestramento stabile.
• Utilità: Il PRM addestrato è presentato come un modulo riutilizzabile e standalone che beneficia le applicazioni downstream come il test-time training e lo scaling, particolarmente in scenari in cui le ricompense di esito non sono disponibili o sono scarse.

Gli autori concludono che, sebbene rePIRL aggiunga una prospettiva teorica unificata e utilità pratica, il lavoro futuro potrebbe esplorare la combinazione con metodi di ricompensa intrinseca, confronti sistematici con euristiche di RL online e l'estensione a compiti agentici e modelli più grandi.