Sintesi Tecnica: Maschere di Divergenza Predittiva per l'RL di LLM

1. Definizione del Problema

L'Apprendimento per Rinforzo (RL) per i Large Language Models (LLM) è intrinsecamente off-policy a causa delle discrepanze tra i motori di inferenza utilizzati per i rollout e gli stack di addestramento, nonché della obsolescenza della policy. Per stabilizzare questi aggiornamenti off-policy, i metodi standard come il Proximal Policy Optimization (PPO) impiegano maschere di regione di fiducia (trust-region masks). Queste maschere si basano su due criteri per decidere se mascherare (azzerare) il gradiente di un token:

1. Criterio di Prossimità: Determina se la policy si è spostata troppo rispetto alla policy di comportamento.
2. Criterio di Direzione: Determina se l'aggiornamento spinge la policy ulteriormente lontano dalla policy di comportamento (direzione verso l'esterno) o se la riporta indietro (correttivo).

Il lavoro recente, nello specifico DPPO, ha migliorato il criterio di prossimità sostituendo il rapporto di importanza a campione singolo con una divergenza distributiva (ad esempio, la divergenza KL) stimata sul top-K del vocabolario. Tuttavia, DPPO mantiene il criterio di direzione originale di PPO, che si basa esclusivamente sul rapporto di importanza del token campionato ($r_t - 1$).

Gli autori identificano un'incoerenza fondamentale in questo approccio ibrido: mentre il criterio di prossimità è ora definito da una divergenza distributiva globale, il criterio di direzione rimane un proxy a campione singolo. Gli autori osservano che il segno del cambiamento nel rapporto campionato può discostarsi dal segno del cambiamento effettivo nella divergenza. Di conseguenza, un token potrebbe essere mantenuto (perché il rapporto suggerisce un aggiornamento correttivo) anche se l'aggiornamento aumenta effettivamente la divergenza globale, o viceversa.

2. Metodologia

Il documento propone la Maschera di Divergenza Predittiva (Predictive Divergence Mask), un metodo che allinea il criterio di direzione con la natura distributiva del criterio di prossimità. Invece di chiedere "Il rapporto campionato si allontanerà da uno?", il metodo chiede "Il prossimo passo del gradiente aumenterà la divergenza?".

Derivazione Teorica

Per le policy softmax utilizzate negli LLM autoregressivi, gli autori derivano un'espressione in forma chiusa per la derivata direzionale della divergenza di Kullback-Leibler (KL), indicata con $\dot{D}$, lungo un passo del gradiente della policy.
$$\dot{D} = \underbrace{(\pi_k - \mu_k)}_{\text{Termine Locale}} + \underbrace{\sum_{i} \pi_i (\mu_i - \pi_i)}_{\text{Termine Globale}}$$
Dove:

• $\pi_k$ e $\mu_k$ sono le probabilità del token campionato $k$ sotto le policy di addestramento e di comportamento, rispettivamente.
• Il Termine Locale corrisponde esattamente alla quantità utilizzata nel classico criterio di direzione basato sul rapporto ($r - 1$).
• Il Termine Globale cattura l'accoppiamento di normalizzazione attraverso l'intero vocabolario. Questo termine è invisibile a un singolo rapporto campionato, ma è cruciale per determinare il reale cambiamento nella divergenza.

La maschera predittiva mantiene un aggiornamento solo se il segno di $\hat{A}_t \cdot \dot{D}_t$ è non positivo (indicando un movimento correttivo o neutro), a condizione che il token sia già al di fuori della regione di fiducia ($D_t > \delta$).

Stima Pratica

I motori di rollout in produzione espongono tipicamente solo le probabilità dei token top-$K$, non la distribuzione completa del vocabolario. Per calcolare $\dot{D}$ sotto questa vista troncata, gli autori propongono due stimatori leggeri per la coda non vista della distribuzione:

1. Stimatore della Coda Aggregata (Aggregated-Tail Estimator): Collassa l'intera massa di probabilità residua in un singolo token aggregato.
2. Stimatore della Coda Uniforme (Uniform-Tail Estimator): Distribuisce la massa residua equamente tra i token del vocabolario non visti.

Entrambi gli stimatori richiedono solo semplici riduzioni sulle probabilità top-$K$ trattenute e sulle masse residue, aggiungendo un carico computazionale trascurabile.

3. Contributi Chiave

• Identificazione dell'Incoerenza: Il documento evidenzia come mescolare un criterio di prossimità distributivo (divergenza) con un criterio di direzione a campione singolo (rapporto) crei un disallineamento in cui la maschera di direzione può fallire nel tracciare la reale variazione della divergenza.
• Maschera di Divergenza Predittiva: Una nuova regola di mascheramento che utilizza la derivata del primo ordine della divergenza stessa per determinare la direzione dell'aggiornamento.
• Derivazione in Forma Chiusa: La derivazione della derivata direzionale per le policy softmax, separando esplicitamente i componenti locali (campionati) e globali (di normalizzazione).
• Stimatori del Vocabolario Troncato: Due stimatori pratici e a basso overhead (code aggregate e uniformi) che permettono al metodo di funzionare con i vincoli standard di produzione dei rollout (viste top-$K$).

4. Risultati Sperimentali

Gli autori hanno valutato il metodo sul dataset DAPO-Math-17k utilizzando modelli Qwen3 (4B, 8B e 30B) con varie impostazioni di precisione (FP8).

• Stabilità dell'Addestramento: In tutte le scale dei modelli e nelle impostazioni di precisione (inclusi i test di stress con FP8), la Maschera di Divergenza Predittiva ha migliorato la stabilità dell'addestramento rispetto al baseline DPPO-TopK-KL. Al contrario, lo standard GRPO con trucchi di "clip-higher" ha mostrato instabilità e collasso nelle configurazioni di modelli più grandi.
• Allineamento del Criterio di Direzione: L'analisi a livello di token ha rivelato che il criterio di direzione basato sulla divergenza è meglio allineato con il cambiamento reale della divergenza rispetto al criterio basato sul rapporto.
  • Tra i token in cui i due criteri erano in disaccordo, il criterio basato sulla divergenza ha avuto un tasso di "mantenimento non sicuro" (unsafe-keep) inferiore (34,2% vs 36,9%), ovvero era meno propenso a mantenere aggiornamenti che in realtà aumentavano la divergenza.
  • Viceversa, ha mantenuto una frazione maggiore di aggiornamenti che hanno contratto con successo la divergenza (65,1% vs 62,9%).
• Robustezza: Il metodo ha dimostrato una migliore robustezza rispetto all'iperparametro della soglia della regione di fiducia ($\delta$), performando meglio del baseline anche sotto vincoli più stretti ($\delta = 0,05$).
• Equivalenza degli Stimatori: Gli stimatori della coda aggregata e della coda uniforme hanno prodotto risultati quasi identici, suggerendo che il metodo non è sensibile all'assunzione specifica di modellazione della coda.

5. Significato e Rivendicazioni

Il documento sostiene che la Maschera di Divergenza Predittiva risolva un'incoerenza strutturale negli attuali metodi di RL basati sulla divergenza per gli LLM. Garantendo che sia il criterio di prossimità che quello di direzione siano derivati dalla stessa divergenza distributiva, il metodo fornisce un controllo della regione di fiducia più fedele.

Gli autori sottolineano che, sebbene il miglioramento nei tassi di "mantenimento non sicuro" sia modesto (a causa del piccolo set di token in cui i criteri sono in disaccordo), la coerenza è teoricamente solida ed empiricamente benefica attraverso le scale. Concludono che il criterio di direzione dovrebbe ricevere lo stesso trattamento distributivo del criterio di prossimità per sfruttare appieno i benefici delle regioni di fiducia basate sulla divergenza. Il metodo è presentato come un sostituto diretto (drop-in replacement) per le implementazioni esistenti di DPPO che non richiede nuovi iperparametri oltre alla soglia di divergenza esistente.