Riepilogo Tecnico: f-Trajectory Balance

Enunciato del Problema

L'allineamento e il fine-tuning dei modelli generativi, inclusi i Large Language Models (LLM) e le Generative Flow Networks (GFlowNets), si basano tipicamente sulla minimizzazione della discrepanza tra la distribuzione di un modello e una distribuzione target. Nel fine-tuning tramite Reinforcement Learning (RL), questo è spesso formulato come massimizzazione del reward atteso soggetto a una penalità di divergenza di Kullback-Leibler (KL).

Una sfida significativa in questo dominio è il compromesso tra comportamenti mode-seeking (ricerca di modi) e mode-covering (copertura dei modi). Gli obiettivi standard, come la divergenza KL inversa (spesso implementata tramite Trajectory Balance o perdita KL al quadrato), tendono a essere mode-seeking, causando il collasso dei modelli su pochi modi ad alta probabilità. Questo è subottimale per compiti che richiedono esplorazione, come la scoperta di farmaci o la generazione di candidati diversificati, dove la "copertura dei modi" (catturare l'intera diversità della distribuzione target) è preferita.

Inoltre, i metodi esistenti spesso faticano con la validità off-policy. Mentre stimatori come REINFORCE o PPO sono comuni, soffrono di alta varianza. Gli approcci recenti che utilizzano la perdita KL al quadrato ($KL_{sq}$) offrono gradienti a bassa varianza on-policy e validità off-policy, ma sono intrinsecamente legati alle proprietà della divergenza KL inversa, limitando la capacità di controllare il compromesso esplorazione-sfruttamento.

Metodologia

Gli autori propongono f-Trajectory Balance, una famiglia di funzioni di perdita surrogate che generalizza la perdita KL al quadrato all'intera famiglia delle f-divergenze.

Fondamento Teorico

Il contributo teorico centrale è l'instaurazione di una corrispondenza uno-a-uno tra funzioni di perdita invarianti per traslazione sulle log-probabilità e f-divergenze.

1. Da Divergenza a Perdita: Per qualsiasi funzione convessa $f$ che definisce una f-divergenza $D_f(p_\theta \| p^*)$, gli autori derivano una funzione di perdita $L_f$ definita sulla differenza delle log-probabilità $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$:
   $$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$$
   Questa perdita possiede due proprietà critiche:

   • Equivalenza del Gradiente On-Policy: Quando i campioni sono estratti dalla politica del modello ($\mu = p_\theta$), i gradienti attesi auto-differenziati di $L_f$ corrispondono ai gradienti della corrispondente f-divergenza $D_f$.
   • Validità Off-Policy: Quando valutata su dati off-policy (campionati da una distribuzione di comportamento $\mu$), la perdita rimane un obiettivo valido con lo stesso minimizzatore globale ($p_\theta = p^*$), a condizione che il supporto del target sia contenuto nel supporto della politica di comportamento.

2. Da Perdita a Divergenza: Viceversa, qualsiasi perdita $\ell$ strettamente convessa e invariante per traslazione sulle log-probabilità corrisponde alla minimizzazione di una specifica f-divergenza. Il documento dimostra che queste mappature sono inverse l'una dell'altra.

Implementazioni Pratiche

Per affrontare le sfide pratiche nell'addestramento, gli autori introducono diverse estensioni:

• Perdita DevGrad (Normalizzazione per Batch): Per gestire le funzioni di partizione intrattabili (comuni nelle GFlowNets e nei target non normalizzati), gli autori generalizzano lo stimatore VarGrad. Invece di utilizzare la varianza del batch, utilizzano la deviazione generalizzata per batch delle differenze di log-probabilità. Questo comporta la stima della costante di normalizzazione $\log Z$ minimizzando la perdita sul batch, centrando efficacemente i coefficienti della funzione di punteggio e riducendo la varianza.
• Perdita Temprata: Per prevenire l'overflow numerico quando si trattano parametri di temperatura $\beta$ piccoli (dove la distribuzione target è $p^* \propto \pi_{ref} \exp(r/\beta)$), gli autori propongono una perdita temprata. Questa scala la perdita per $1/\beta$ e la applica alla distribuzione temprata, garantendo che i gradienti rimangano limitati e stabili.
• f-Trajectory Balance per GFlowNets: Il framework è applicato alle GFlowNets imponendo vincoli di equilibrio delle traiettorie utilizzando la perdita $L_f$. Questo generalizza il Trajectory Balance standard (che corrisponde alla KL inversa) a f-divergenze arbitrarie.

Contributi Chiave

1. Famiglia di Perdite Generalizzata: Derivazione di una famiglia generale di perdite surrogate invarianti per traslazione ($L_f$) i cui gradienti on-policy corrispondono a quelli della corrispondente f-divergenza, mantenendo al contempo la validità off-policy.
2. Equivalenza Teorica: Dimostrazione della relazione inversa tra perdite di log-probabilità invarianti per traslazione e f-divergenze, permettendo la progettazione sistematica di nuove funzioni di perdita con proprietà desiderate (ad esempio, copertura dei modi).
3. Estensione a Target Non Normalizzati: Generalizzazione della tecnica VarGrad a f-divergenze arbitrarie tramite la perdita DevGrad, affrontando il problema della funzione di partizione intrattabile nel fine-tuning RL.
4. Controllo dell'Esplorazione: Dimostrazione che la scelta della f-divergenza (parametrizzata da $\alpha$ nelle $\alpha$-divergenze) controlla direttamente il comportamento mode-seeking vs mode-covering. Valori di $\alpha$ più bassi (ad esempio, Forward KL, Hellinger) promuovono la copertura dei modi, mentre valori più alti (ad esempio, Pearson $\chi^2$) promuovono la ricerca dei modi.

Risultati Sperimentali

Gli autori valutano la famiglia di perdite su quattro compiti distinti:

1. Griglia Sintetica (GFlowNets): Su un compito su griglia 2D progettato per testare la copertura dei modi, le perdite con $\alpha$ più basso (Forward KL, Hellinger) hanno scoperto con successo tutti e quattro i modi della distribuzione target significativamente più velocemente e in modo più affidabile rispetto al Trajectory Balance standard (KL inversa). Al contrario, il Pearson $\chi^2$ ($\alpha=2$) ha mostrato un forte comportamento mode-seeking, bloccandosi nel primo modo trovato.
2. Scoperta di Molecole (SynFlowNet): Applicato al framework SynFlowNet per la generazione di molecole sinteticamente accessibili. Gli autori hanno rilevato che la regolazione di $\alpha$ permetteva di controllare la diversità delle molecole generate. Un programma di annealing (iniziando con $\alpha$ più basso per l'esplorazione e aumentando fino a $\alpha=1$ per lo sfruttamento) ha prodotto il miglior compromesso, generando molecole con reward più alto e maggiore diversità rispetto al Trajectory Balance standard.
3. Modelli Diffusivi: Nel campionamento condizionato di cifre MNIST (dispari vs pari), il Trajectory Balance standard ha fallito nel campionare uniformemente attraverso i modi target (sovracampionando cifre specifiche come 0 e 6). Le perdite f-Trajectory Balance, in particolare quelle con $\alpha$ più basso, hanno raggiunto un campionamento più uniforme attraverso i modi target.
4. Fine-tuning Asincrono di LLM: Le perdite sono state applicate al fine-tuning di LLM (famiglie Qwen e OLMo) su problemi di matematica (GSM8k, MATH) con un ritardo asincrono di 50 passi. Le perdite basate su f-divergenze hanno dimostrato un addestramento stabile e compromessi costanti tra entropia e reward, mentre le implementazioni PPO standard hanno mostrato instabilità nelle stesse condizioni asincrone.

Significato e Affermazioni

Il documento afferma che l'efficacia della perdita KL al quadrato non è un fenomeno isolato, ma un caso specifico di una famiglia più ampia di perdite. Estendendo la proprietà di "corrispondenza del gradiente on-policy" e "validità off-policy" all'intera famiglia delle f-divergenze, gli autori forniscono un framework unificato per il tuning dei modelli generativi.

Il significato principale risiede nella capacità di controllare esplicitamente il compromesso esplorazione-sfruttamento semplicemente selezionando una diversa funzione di perdita surrogate (o regolando un parametro $\alpha$), senza modificare l'algoritmo di addestramento sottostante. Questo permette ai praticanti di scegliere tra comportamenti mode-seeking (per la massimizzazione del reward) e mode-covering (per diversità ed esplorazione), mantenendo la stabilità e la compatibilità off-policy richieste per l'addestramento moderno, su larga scala e asincrono, di LLM e GFlowNets. Il lavoro suggerisce che questo framework è particolarmente prezioso per applicazioni come la scoperta di farmaci e agenti esplorativi, dove coprire lo spazio dei candidati ad alto reward è critico.