Riepilogo Tecnico: Ripensare la Regolarizzazione della Divergenza nel RL per LLM

1. Definizione del Problema

L'apprendimento per rinforzo (RL) è una pietra angolare del post-training dei grandi modelli linguistici (LLM) per allinearli alle preferenze umane e potenziare le capacità di ragionamento. Tuttavia, l'addestramento RL moderno per gli LLM è intrinsecamente off-policy a causa delle discrepanze tra i motori di training e di inferenza (discrepanze numeriche) e della pratica di suddividere i rollout in più mini-batch. Ciò rende necessario l'uso di Trust Region Policy Optimization (TRPO) per garantire aggiornamenti stabili, vincolando la divergenza tra la politica attuale ($\pi$) e la politica di comportamento ($\mu$).

I metodi mainstream attuali affrontano due limitazioni primarie in questo contesto:

1. Vincoli Basati sul Rapporto (PPO, GRPO, SPO): Questi metodi utilizzano il rapporto di importanza $r_t = \pi(y_t|s_t)/\mu(y_t|s_t)$ per definire la regione di fiducia. Negli LLM con vocabolari ampi e a coda lunga, il rapporto è un cattivo indicatore dello spostamento distributivo. Un piccolo aumento di probabilità su un token a bassa probabilità può generare un rapporto enorme con uno spostamento di massa trascurabile, mentre un cambiamento moderato del rapporto su un token ad alta probabilità può indurre cambiamenti significativi nella politica. Ciò porta a sovra-vincolare i token rari e a sotto-vincolare quelli comuni.
2. Mascheramento Rigido (DPPO): Lavori recenti come il Divergence Proximal Policy Optimization (DPPO) affrontano il problema del rapporto utilizzando una maschera basata sulla divergenza (specificamente la Binary-TV, che misura lo spostamento assoluto della probabilità $|\pi - \mu|$). Tuttavia, DPPO impone la regione di fiducia tramite una maschera binaria rigida. Una volta che un token supera il confine in una direzione dannosa, il suo gradiente viene scartato (impostato a zero). Questo impedisce l'ulteriore divergenza ma non fornisce un segnale correttivo per riportare la politica verso la regione di fiducia, portando a un'ottimizzazione fragile vicino alla soglia.

2. Metodologia: Divergence Regularized Policy Optimization (DRPO)

Gli autori propongono DRPO, un metodo che sostituisce la maschera rigida di DPPO con un regolarizzatore quadratico liscio e pesato dal vantaggio sullo spostamento della politica, preservando al contempo la stessa geometria della regione di fiducia Binary-TV.

Formulazione Core

DRPO deriva il proprio obiettivo riscrivendo il vincolo Binary-TV $|\pi(y_t|s_t) - \mu(y_t|s_t)| \leq \delta$ come un limite di rapporto adattivo al token. Invece di un clipping, applica un regolarizzatore liscio all'obiettivo surrogato:

$$\mathcal{L}_{DRPO}(x, \pi) = \mathbb{E}_{y \sim \mu} \left[ \sum_{t=1}^{|y|} \left( r_t \hat{A}_t - \frac{|\hat{A}_t|}{2\delta} \mu(y_t|s_t) (r_t - 1)^2 \right) \right]$$

Dove:

• $r_t$ è il rapporto di importanza.
• $\hat{A}_t$ è la stima del vantaggio.
• $\delta$ è la soglia della regione di fiducia.
• Il termine $\mu(y_t|s_t)$ scala il regolarizzatore, convertendo la penalità da una forma $\chi^2$ basata sul rapporto (come in SPO) a una forma $\ell_2^2$ basata sullo spostamento della probabilità assoluta.

Comportamento del Gradiente

Il gradiente dell'obiettivo DRPO introduce un peso continuo $w_t$ per il gradiente della politica di ogni token:
$$w_t = 1 - \text{sign}(\hat{A}_t(r_t - 1)) \frac{|\pi(y_t|s_t) - \mu(y_t|s_t)|}{\delta}$$

Questo peso esibisce tre comportamenti chiave:

1. All'interno della Regione di Fiducia: Il peso è positivo ma viene attenuato gradualmente man mano che lo spostamento si avvicina a $\delta$, permettendo aggiornamenti ma con magnitudo ridotta.
2. Fuori dalla Regione di Fiducia (Divergenza): Quando l'aggiornamento si allontana ulteriormente dalla politica di comportamento ($|\pi - \mu| > \delta$), il peso diventa negativo. Questo inverte il gradiente, fornendo un segnale correttivo che riporta la politica verso il confine, diversamente dallo zero-gradiente delle maschere rigide.
3. Aggiornamenti di Convergenza: Quando l'aggiornamento torna verso la politica di comportamento, il peso viene amplificato ($>1$), incoraggiando un ritorno fluido alla regione di fiducia.

Distinzioni Teoriche

• vs. SPO: SPO utilizza un regolarizzatore di tipo $\chi^2$ (dividendo per $\mu$), il che causa una crescita illimitata dei pesi del gradiente nella coda lunga dei token a bassa probabilità. DRPO utilizza un regolarizzatore di tipo $\ell_2^2$ (moltiplicando per $\mu$), garantendo che i pesi del gradiente rimangano limitati indipendentemente dalla probabilità del token.
• vs. DPPO: DRPO preserva la geometria Binary-TV di DPPO ma sostituisce la maschera discontinua con un regolarizzatore liscio e differenziabile.

3. Contributi Chiave

1. Algoritmo DRPO: Un nuovo metodo di ottimizzazione della politica che sostituisce le maschere di divergenza rigide con un regolarizzatore quadratico liscio basato sullo spostamento della probabilità assoluta (Binary-TV).
2. Pesi del Gradiente Limitati: Dimostrazione teorica ed empirica che i pesi del gradiente di DRPO rimangono limitati nei vocabolari a coda lunga, evitando l'instabilità causata dai metodi basati sul rapporto (PPO, SPO) dove i pesi possono esplodere per i token a bassa probabilità.
3. Meccanismo Correttivo: Introduzione di un segnale correttivo continuo per gli aggiornamenti che eccedono la regione di fiducia, prevenendo la "morte del gradiente" osservata negli approcci a maschera rigida come DPPO.
4. Principi di Design: Identificazione di tre criteri pratici per regolarizzatori efficaci nel RL per LLM:
   • Indurre un confine stabile allineato con lo spostamento distributivo (non solo con i rapporti).
   • Mantenere i pesi del gradiente per token limitati nei vocabolari a coda lunga.
   • Fornire segnali correttivi fluidi quando la politica diverge.

4. Risultati Sperimentali

Gli autori hanno valutato DRPO in diversi setting, tra cui:

• Modelli: Qwen3-4B, Qwen3-30B-A3B, Qwen3.5-35B-A3B, e DeepSeek-R1-Distill-Qwen-1.5B.
• Architetture: Dense e Mixture-of-Experts (MoE).
• Precisione: BF16, FP8 (solo rollout), e FP8 (end-to-end).
• Benchmark: AIME 2024 e AIME 2025 (compiti di ragionamento matematico).

Risultati Principali:

• Stabilità: DRPO consente costantemente un addestramento stabile in tutti i setting. Al contrario, i metodi basati sul rapporto (GRPO, SPO) hanno spesso sofferto di instabilità o collasso, particolarmente nei setting a bassa precisione (FP8).
• Performance: DRPO ha eguagliato o superato la migliore accuratezza di valutazione di tutti i baseline. Ha superato i metodi a maschera rigida (DPPO) in velocità di convergenza e accuratezza finale.
• Ablazioni:
  • Rimuovendo il peso del vantaggio assoluto $|\hat{A}_t|$ dal regolarizzatore, si sono verificati cali di performance e instabilità, confermando che il confine della regione di fiducia deve essere indipendente dalla scala del reward.
  • Misure di divergenza alternative (penalità KL, TV) hanno ottenuto prestazioni inferiori rispetto a DRPO, suggerendo che la specifica geometria del regolarizzatore Binary-TV è cruciale.
  • Applicare il regolarizzatore solo al di fuori della regione di fiducia (Mask-DRPO) ha prodotto performance comparabili a DRPO completo, indicando che il guadagno principale deriva dalla regolarizzazione correttiva al di fuori del confine.

5. Significato e Rivendicazioni

Il paper sostiene che la geometria del gradiente indotta da un regolarizzatore è più critica della misura nominale di divergenza nella funzione obiettivo. Mentre i lavori precedenti si concentrano spesso su divergenze standard (KL, JS), questo lavoro dimostra che:

• I vincoli basati sul rapporto sono fondamentalmente difettosi per gli LLM a causa dei vocabolari a coda lunga.
• Le maschere rigide, pur essendo migliori dei rapporti, introducono fragilità scartando i segnali correttivi.
• DRPO offre una soluzione unificata che combina la correttezza geometrica dei vincoli basati sulla divergenza (Binary-TV) con la stabilità di ottimizzazione dei regolarizzatori lisci.

Gli autori posizionano DRPO non solo come un algoritmo specifico, ma come una lente empirica per progettare obiettivi di ottimizzazione della politica stabili, enfatizzando che i regolarizzatori efficaci devono fornire segnali di gradiente limitati, continui e correttivi per gestire le sfide uniche del RL per LLM.