Evaluating GFlowNet from partial episodes for stable and flexible policy-based training

Questo lavoro colma il divario tra gli approcci basati su valori e su politiche nei GFlowNet introducendo un obiettivo di bilanciamento della valutazione su episodi parziali che migliora l'affidabilità e la flessibilità dell'addestramento, permettendo l'uso di politiche di ritorno parametriche e di dati offline.

L'essenziale

Immagina di dover esplorare un labirinto gigantesco e buio per trovare le stanze più preziose (quelle con il "premio" più alto). Il problema è che il labirinto è così grande che non puoi vederlo tutto, e le stanze preziose sono sparse in modo casuale, nascoste dietro angoli bui.

Questo è il problema che risolvono le GFlowNets (Reti di Flusso Generativo). Sono come un esploratore intelligente che impara a camminare nel labirinto per trovare queste stanze preziose, non una alla volta, ma imparando a mappare l'intero percorso in modo da visitarle tutte con la frequenza giusta.

Ecco di cosa parla questo nuovo lavoro, spiegato in modo semplice:

1. Il Problema: L'Esploratore Confuso

Per far funzionare questo esploratore (chiamato "politica" o policy), gli scienziati usano due metodi principali:

• Metodo "Valore" (Value-based): Come un cartografo che disegna una mappa delle "correnti" del labirinto. Se il flusso d'acqua (i dati) è bilanciato, la mappa è corretta. È robusto, ma a volte lento e rigido.
• Metodo "Politica" (Policy-based): Come un allenatore che dice all'esploratore: "Sbagli, vai meglio qui". Per farlo, l'allenatore ha bisogno di un giudice (chiamato Evaluator o Critic) che gli dica quanto è bravo l'esploratore in ogni punto del labirinto.

Il problema: In passato, questo "giudice" era molto difficile da addestrare. Era come avere un arbitro che sbaglia spesso i fischietti: l'allenatore (l'esploratore) si confonde, impara male e il sistema diventa instabile.

2. La Soluzione: Il "Bilanciatore di Episodi Parziali" (Sub-EB)

Gli autori di questo paper (Niu, Wu e Qian) hanno avuto un'idea brillante. Hanno detto: "E se usassimo la logica della mappa delle correnti (il metodo Valore) per addestrare il nostro giudice (il metodo Politica)?"

Hanno creato una nuova regola chiamata Sub-EB (Subtrajectory Evaluation Balance).

L'analogia della Bilancia:
Immagina di dover pesare un oggetto, ma non hai una bilancia perfetta.

• Il vecchio metodo cercava di pesare l'oggetto intero alla fine del viaggio. Se il viaggio era lungo, il peso era difficile da calcolare e pieno di errori.
• Il nuovo metodo Sub-EB pesa l'oggetto a pezzi (pezzi di viaggio, o "episodi parziali").
• Dice: "Se il peso che ho calcolato per il primo pezzo del viaggio più il peso del secondo pezzo è uguale al peso totale atteso, allora il mio giudice sta funzionando bene!"

In pratica, invece di chiedere al giudice di indovinare il punteggio finale subito, lo costringono a essere coerente passo dopo passo. Questo rende il giudice molto più affidabile e stabile.

3. I Vantaggi Magici

Grazie a questo nuovo metodo, succedono due cose fantastiche:

1. Maggiore Flessibilità (Il Giudice che Impara):
   Prima, il "giudice" doveva essere molto semplice e rigido. Ora, grazie a Sub-EB, possiamo usare giudici molto più sofisticati e complessi (chiamati backward policies parametrizzate). È come passare da un arbitro che guarda solo il pallone, a un arbitro che usa la telecronaca, i replay e l'AI per prendere decisioni perfette. Questo permette al sistema di adattarsi meglio a compiti difficili.

2. Uso di Vecchi Dati (Allenamento Offline):
   Prima, l'esploratore doveva imparare solo camminando nel labirinto in tempo reale (online). Se sbagliava, perdeva tempo.
   Con Sub-EB, l'esploratore può anche studiare le mappe di altri esploratori che hanno già camminato nel labirinto (dati offline). È come se potesse guardare le registrazioni di altri viaggiatori per imparare dai loro errori senza doverli commettere di nuovo. Questo rende l'addestramento molto più veloce ed efficiente.

4. I Risultati: Ha Funzionato?

Gli autori hanno testato questa idea su diversi "labirinti":

• Griglie virtuali: Come puzzle matematici complessi.
• Design di molecole: Come cercare di inventare nuovi farmaci o materiali chimici (dove ogni "stanza" è una molecola diversa).
• Reti Bayesiane: Come organizzare informazioni complesse (usato in intelligenza artificiale medica o finanziaria).

In tutti questi casi, il nuovo metodo Sub-EB ha dimostrato di essere:

• Più stabile: Non va in crash o si blocca durante l'addestramento.
• Più veloce: Impara a trovare le soluzioni migliori in meno tempo.
• Più creativo: Trova soluzioni diverse e innovative, non si blocca sulle prime risposte che trova.

In Sintesi

Questo paper è come se avessimo dato all'esploratore del labirinto un nuovo allenatore che usa le regole della fisica (il flusso) per insegnargli a camminare meglio. Invece di dire "hai sbagliato alla fine", l'allenatore controlla ogni singolo passo, assicurandosi che tutto sia in equilibrio. Il risultato è un'intelligenza artificiale che impara più velocemente, è più sicura e riesce a risolvere problemi complessi (come creare nuovi farmaci) in modo molto più efficace.

Sommario tecnico

1. Il Problema

Le Generative Flow Networks (GFlowNets) sono modelli generativi progettati per campionare candidati combinatori (come grafi, sequenze o strutture) con una probabilità proporzionale a una funzione di ricompensa $R(x)$. Il processo di generazione è modellato come traiettorie su un Grafo Aciclico Diretto (DAG).

Esistono due principali paradigmi di addestramento:

• Metodi basati sul valore (Value-based): Si basano sul bilanciamento dei flussi (flow balance) lungo le traiettorie o sottotrajettorie. Sono stabili ma spesso richiedono politiche di raccolta dati specifiche (off-policy) e non ottimizzano direttamente la divergenza della politica.
• Metodi basati sulla politica (Policy-based): Utilizzano un approccio Actor-Critic per minimizzare la divergenza di Kullback-Leibler (KL) tra la distribuzione delle traiettorie forward e quella backward. Tuttavia, questi metodi affrontano una sfida fondamentale: la stima affidabile della funzione di valutazione $V(s)$, che misura la divergenza tra le sottotrajettorie forward e backward a partire da uno stato $s$. Le stime attuali (es. basate su $\lambda$-TD) sono spesso instabili o richiedono politiche backward fisse, limitando la flessibilità e l'efficienza.

Il problema centrale è come apprendere in modo stabile e affidabile la funzione di valutazione $V(s)$ per guidare l'aggiornamento della politica forward, specialmente in spazi combinatori vasti, senza sacrificare la flessibilità (es. l'uso di politiche backward parametriche o dati offline).

2. Metodologia: Sub-EB (Subtrajectory Evaluation Balance)

Gli autori propongono un nuovo approccio che colma il divario tra i metodi basati sul valore e quelli basati sulla politica, introducendo il concetto di Subtrajectory Evaluation Balance (Sub-EB).

Concetti Chiave

1. Connessione Teorica: Il paper dimostra teoricamente che, per una politica forward $\pi_F$ fissata, la soluzione dell'equazione di bilanciamento del flusso logaritmico coincide esattamente con la funzione di valutazione $V(s)$ che misura la divergenza KL tra le distribuzioni forward e backward.

2. Condizione Sub-EB: Viene definita una nuova condizione di equilibrio per la funzione di valutazione $V$. Per ogni sottotrajettoria parziale (da uno stato $s_i$ a uno stato $s_j$), la condizione richiede che:
   $$\mathbb{E}_{P_F} \left[ \log \left( P_F(\tau_{i:j}|s_i) \exp(V(s_i)) \right) \right] = \mathbb{E}_{P_F} \left[ \log \left( P_B(\tau_{i:j}|s_j) \exp(V(s_j)) \right) \right]$$
   Questa condizione garantisce che la differenza tra i valori di divergenza appresi agli stati $s_i$ e $s_j$ corrisponda alla vera divergenza sulle sottotrajettorie tra questi stati.

3. Ottimizzazione (Sub-EB Objective): Viene introdotto un nuovo obiettivo di perdita ($L_V$) per apprendere $V(\cdot; \phi)$ minimizzando l'errore quadratico medio della condizione Sub-EB su tutte le sottotrajettorie di un episodio.

   • A differenza degli obiettivi $\lambda$-TD tradizionali che considerano solo mismatch a livello di singolo passo (edge-wise) o eventi che iniziano in un punto specifico, Sub-EB utilizza mismatch su intere sottotrajettorie (subtrajectory-wise).
   • Questo permette un apprendimento più bilanciato e stabile di $V(s)$, incorporando informazioni sia dagli stati precedenti che successivi.

4. Flessibilità e Addestramento Offline:

   • Politiche Backward Parametriche: A differenza dei metodi precedenti che richiedevano una politica backward $\pi_B$ fissa, Sub-EB permette di aggiornare congiuntamente $\pi_B$ e $V$ durante l'ottimizzazione, rendendo il processo più efficiente e adattivo.
   • Addestramento Offline: Viene proposto un flusso di lavoro offline (Algoritmo 2) che utilizza una politica di raccolta dati $\pi_D$ diversa da $\pi_F$. Viene definita una funzione di valutazione inversa $W$ e un obiettivo di bilanciamento backward, permettendo l'integrazione di tecniche di raccolta dati offline (come la ricerca locale) per esplorare stati ad alta ricompensa senza compromettere la stabilità dell'addestramento della politica.

3. Contributi Chiave

• Teoria: Stabilisce un legame formale e rigoroso tra la funzione di flusso $F(s)$ (usata nei metodi basati sul valore) e la funzione di valutazione $V(s)$ (usata nei metodi basati sulla politica), dimostrando che la condizione di bilanciamento del flusso è una condizione sufficiente per la corretta valutazione della divergenza.
• Nuovo Obiettivo (Sub-EB): Propone l'obiettivo Sub-EB per l'apprendimento affidabile di $V$, che supera le limitazioni degli obiettivi $\lambda$-TD tradizionali in termini di stabilità e varianza.
• Flessibilità Operativa: Dimostra che Sub-EB supporta nativamente politiche backward parametriche e permette l'addestramento offline, risolvendo problemi di esplorazione e sfruttamento (exploration-exploitation) tipici dei metodi GFlowNet.
• Validazione Sperimentale: Fornisce evidenze empiriche su dataset sintetici (Hypergrids) e reali (design di sequenze biologiche/molecolari, apprendimento di strutture di Bayesian Network).

4. Risultati Sperimentali

Gli esperimenti confrontano Sub-EB con metodi basati su politica (RL, CV) e basati su valore (Sub-TB, Q-Much).

• Hypergrids (Ambienti Sintetici):

  • Su griglie 2D e 3D di grandi dimensioni, Sub-EB mostra una convergenza più rapida e una stabilità superiore rispetto ai metodi RL tradizionali (basati su $\lambda$-TD).
  • Sub-EB supera i metodi basati su valore (Sub-TB) e le varianti empiriche (CV) in termini di accuratezza della distribuzione finale (misurata con DTV e DJSD).
  • L'uso di politiche backward parametriche con Sub-EB (Sub-EB-P) porta alle prestazioni migliori, confermando la capacità del metodo di adattarsi dinamicamente.

• Design di Sequenze (Biologiche e Molecolari):

  • Su dataset come SIX6, PHO4, QM9 e sEH, Sub-EB dimostra un'efficacia superiore nel modellare la distribuzione e nel scoprire modalità (mode discovery) ad alta ricompensa.
  • Le varianti offline (Sub-EB-B) che integrano tecniche di ricerca locale riescono a trovare più stati terminali ad alta ricompensa, sebbene con un leggero compromesso sulla diversità della distribuzione, dimostrando la flessibilità del framework.

• Apprendimento di Strutture Bayesiane (BN):

  • Su spazi di ricerca enormi (fino a $10^{35}$ strutture), Sub-EB raggiunge i punteggi di ricompensa più alti e converge più velocemente di Sub-TB e Q-Much.
  • Mantiene un'alta diversità nelle soluzioni generate, superando i metodi basati su valore che tendono a convergere prematuramente su strutture subottimali.

• Design di Grafi Molecolari:

  • Su task complessi come LogP e JNK3, Sub-EB ottiene il miglior compromesso tra ricompensa media e diversità, superando i metodi concorrenti in termini di velocità di convergenza.

5. Significato e Impatto

Questo lavoro rappresenta un passo significativo per l'adozione delle GFlowNets in scenari reali complessi:

1. Stabilità: Risolve il problema dell'instabilità nell'addestramento basato su politica, rendendo i metodi Actor-Critic per GFlowNets competitivi e spesso superiori rispetto ai metodi basati sul valore.
2. Generalizzazione: La capacità di gestire politiche backward parametriche e dati offline apre la strada all'uso di GFlowNets in contesti dove la raccolta dati è costosa o dove è necessaria un'esplorazione profonda guidata da euristiche esterne.
3. Unificazione Teorica: Fornisce una base teorica solida che unifica le prospettive di "flusso" e "divergenza", offrendo nuovi strumenti per lo sviluppo di algoritmi di generazione combinatoria più robusti.

In sintesi, il paper introduce Sub-EB come un nuovo standard per l'addestramento delle GFlowNets, offrendo un equilibrio ottimale tra stabilità, flessibilità e prestazioni su spazi combinatori su larga scala.