Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

Riassunto Tecnico: SHEAR (Span-level Hidden state Enabled Advantage Reweighting)

1. Il Problema: Il limite dell'assegnazione del credito (Credit Assignment)

Nel campo del Reinforcement Learning con ricompense verificabili (RLVR), algoritmi come GRPO (Group Relative Policy Optimization) vengono utilizzati per migliorare il ragionamento dei Large Language Models (LLM). Tuttavia, questi metodi soffrono di un problema fondamentale di assegnazione del credito a grana grossa: quando un modello genera una sequenza di ragionamento lunga, riceve un unico feedback binario (corretto/errato) alla fine. GRPO assegna lo stesso vantaggio (advantage) a tutti i token della sequenza, indipendentemente dal fatto che siano stati fondamentali per il successo o che siano stati l'origine di un errore.

Le alternative attuali, come i Process Reward Models (PRM), tentano di assegnare ricompense a ogni singolo passaggio, ma richiedono annotazioni umane costose o l'addestramento di modelli ausiliari, che possono essere rumorosi o non allineati con la politica in evoluzione.

2. Metodologia: L'intuizione delle rappresentazioni latenti

Gli autori propongono di estrarre un segnale di "qualità del ragionamento" direttamente dagli stati nascosti (hidden states) del modello stesso, senza bisogno di modelli esterni o annotazioni passo-passo.

L'osservazione empirica:
Il paper dimostra che le distribuzioni degli stati nascosti contengono un segnale intrinseco: nei punti in cui il ragionamento diverge tra traiettorie corrette e errate, la distanza tra le loro distribuzioni latenti aumenta drasticamente.

Il metodo SHEAR:
L'algoritmo SHEAR modifica il processo di aggiornamento di GRPO attraverso quattro fasi principali:

1. Decomposizione in Span: La sequenza di token viene suddivisa in segmenti sovrapposti (spans) utilizzando una finestra scorrevole.
2. Distanza di Wasserstein a livello di Span: Per ogni span, viene calcolata la Distanza di Wasserstein ($W_1$) rispetto agli span del gruppo opposto (es. lo span di una traiettoria errata rispetto agli span delle traiettorie corrette). La distanza di Wasserstein è scelta perché cattura differenze nella forma e nella struttura della distribuzione, non solo uno spostamento della media.
3. Aggregazione tramite Max-Pooling: Il segnale di divergenza viene riportato dal livello di "span" al livello di "token" utilizzando il max-pooling, garantendo che ogni token riceva il peso della divergenza più forte tra gli span che lo includono.
4. Riassegnazione del Vantaggio (Reweighting): Il vantaggio di GRPO viene riscalato per ogni token:
   $$\tilde{A}_t^{(i)} = A^{(i)} \cdot \omega_t^{(i)}$$
   Dove $\omega_t^{(i)}$ è il peso derivato dalla distanza di Wasserstein. Questo amplifica gli aggiornamenti sui token che si trovano in regioni di forte divergenza (dove l'errore è avvenuto o dove il successo è stato consolidato).

3. Contributi Chiave

• Scoperta Empirica: Dimostrazione che la divergenza della distanza di Wasserstein tra stati nascosti è un proxy affidabile per la qualità locale del ragionamento.
• Fondazione Teorica: Viene presentato un Teorema di Separazione che prova matematicamente che, sotto ipotesi strutturali minime, gli span successivi a un errore mostrano distanze di Wasserstein significativamente maggiori rispetto agli span precedenti, permettendo la distinzione tra segnale e rumore di campionamento.
• Algoritmo SHEAR: Un metodo di self-supervision che non richiede modelli di ricompensa aggiuntivi, minimizzando l'overhead computazionale (inferiore al 16%).

4. Risultati Sperimentali

Il metodo è stato testato su benchmark di ragionamento matematico (MATH500, AIME, ecc.) e generazione di codice (HumanEval, MBPP, ecc.) utilizzando diversi modelli (Qwen2.5, Llama 3.1).

• Ragionamento Matematico: SHEAR supera costantemente GRPO e, sorprendentemente, supera i modelli basati su PRM (Process Reward Models) supervisionati. Questo accade perché il segnale di SHEAR è intrinsecamente allineato con la politica in addestramento, evitando il "mismatch" tipico dei PRM esterni.
• Generazione di Codice: Si osservano miglioramenti consistenti, specialmente su task complessi e lunghi (come LiveCodeBench), dove l'assegnazione del credito è più critica.
• Efficienza: L'incremento del tempo di addestramento è contenuto e decresce all'aumentare della dimensione del modello.

5. Significato e Implicazioni

Il lavoro di Chen et al. sposta il paradigma dell'assegnazione del credito nel Reinforcement Learning per LLM: invece di cercare di "insegnare" al modello cosa sia giusto tramite etichette esterne, si impara a sfruttare la consapevolezza interna del modello (le sue rappresentazioni latenti) per identificare i momenti critici del ragionamento. Questo approccio apre la strada a un RL più efficiente, scalabile e capace di auto-correggersi senza la necessità di costosi dataset annotati passo-passo.