Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Il paper propone RapTB, un nuovo obiettivo di training per GFlowNet che risolve il collasso dei moduli e il bias di lunghezza assegnando crediti densi ai prefissi, e SubM, una strategia di replay submodulare che garantisce diversità, migliorando così la generazione di molecole con LLM.

L'essenziale

Immagina di avere un chef robotico (chiamato GFlowNet) il cui compito è creare piatti deliziosi (come nuove molecole per farmaci o frasi intelligenti) partendo da un ingrediente base. Il suo obiettivo non è solo trovare un piatto perfetto, ma esplorare tutti i possibili piatti buoni, creando una varietà enorme di opzioni, ognuna con una probabilità proporzionale alla sua bontà.

Il problema è che questo chef robotico tende a diventare pigro e ripetitivo. Invece di esplorare la cucina, si blocca su una sola ricetta base (ad esempio, "pasta al pomodoro") e la ripete all'infinito, ignorando tutte le altre possibilità. Questo fenomeno si chiama "crollo delle modalità" (mode collapse).

La carta scientifica che hai condiviso propone due soluzioni intelligenti per risolvere questo problema: RapTB e SubM. Ecco come funzionano, spiegate con metafore semplici.

1. Il Problema: Lo Chef che si blocca

Il robot ha due difetti principali:

• Crollo del Prefisso (Prefix Collapse): Fa i primi passi della ricetta (es. "prendi la pasta") sempre uguali per tutti i piatti, anche se poi dovrebbero essere diversi. È come se tutti i piatti iniziassero con "Metti l'acqua nel pentolino...", ma poi il robot non sa più cosa fare e si blocca.
• Pregiudizio sulla Lunghezza: Tende a fare piatti troppo corti (solo acqua) o troppo lunghi (un'opera d'arte che non finisce mai), perdendo il senso della ricetta.

Perché succede?

1. Ricompense tardive: Il robot riceve un voto ("buono" o "cattivo") solo alla fine del piatto. Se il piatto è lungo, non sa quali passi intermedi siano stati importanti. È come guidare una macchina senza specchietti retrovisori, aspettando di arrivare a destinazione per capire se hai preso la strada giusta.
2. Il Diario degli Errori (Replay Bias): Il robot impara leggendo un diario delle sue ricette migliori. Ma se il diario contiene solo 5 ricette quasi identiche (perché le ha trovate per caso), il robot impara solo quelle e smette di cercare novità.

---

2. La Soluzione 1: RapTB (Il Tutor che ti guida passo dopo passo)

RapTB è come avere un tutor esperto che cammina accanto allo chef robotico mentre cucina.

• Come funziona: Invece di aspettare la fine del piatto per dare un voto, il tutor guarda ogni singolo passo (ogni "prefisso" della ricetta).
• L'idea geniale (Assorbimento): Se il robot sta preparando un piatto e arriva a metà strada, il tutor non dice solo "bravo". Il tutor guarda anche il futuro di quella ricetta. Se sa che la ricetta porterà a un piatto eccellente, dice al robot: "Ehi, quel passo che hai fatto 30 secondi fa era ottimo, perché ti ha portato qui!".
• Il vantaggio: Questo dà al robot un feedback costante e chiaro su ogni singola decisione, anche quelle fatte all'inizio. Non deve più indovinare. Inoltre, il tutor è molto attento a non confondere il robot: non gli dice di cambiare il modo in cui decide quando fermarsi (evitando che il robot faccia piatti infiniti o troppo corti), ma si concentra solo sul migliorare gli ingredienti scelti.

In sintesi: RapTB trasforma un voto finale vago in una serie di consigli precisi e immediati per ogni passo della ricetta.

---

3. La Soluzione 2: SubM (Il Curatore del Museo)

SubM è come un curatore di un museo che deve scegliere quali quadri esporre.

• Il problema: Se il museo (il "buffer" di memoria dove il robot impara) è pieno solo di quadri rossi perché sono stati i più venduti, il pubblico (il robot) penserà che l'arte sia solo rossa.
• La soluzione (Submodularità): Il curatore usa una regola matematica intelligente per scegliere i quadri. Non sceglie solo i "più belli" (alta ricompensa), ma cerca un equilibrio:
  1. Qualità: Il quadro deve essere bello.
  2. Diversità: Il quadro deve essere diverso dagli altri (niente copie identiche).
  3. Copertura: Il museo deve avere quadri di tutte le dimensioni (piccoli, medi, grandi).
• Come funziona: Ogni volta che il robot crea nuovi quadri, il curatore li guarda e sceglie il gruppo migliore da mettere in mostra, assicurandosi che il museo non diventi una galleria di cloni.

In sintesi: SubM impedisce al robot di imparare solo dalle copie identiche, costringendolo a esplorare un mondo più vasto e vario.

---

Il Risultato Finale

Quando unisci RapTB (il tutor attento) e SubM (il curatore intelligente):

1. Il robot smette di fare sempre la stessa ricetta base.
2. Impara a creare piatti lunghi e complessi senza impazzire.
3. Produce una varietà incredibile di molecole o frasi, tutte di alta qualità.

L'analogia finale:
Immagina di insegnare a un bambino a scrivere una storia.

• Senza queste tecniche: Gli dici "Scrivi una storia" e lo premi solo se finisce con un "Fine" perfetto. Lui scriverà sempre "C'era una volta. Fine." perché è la via più sicura.
• Con RapTB: Gli dici, mentre scrive: "Ottima idea iniziare con 'C'era una volta'! E se aggiungessi un drago qui? Bravissimo!".
• Con SubM: Gli mostri un libro di storie che contiene avventure, commedie, drammi e storie di animali, non solo la stessa storia di un gatto ripetuta 100 volte.

Il risultato è che il bambino (il modello AI) diventa un vero scrittore creativo, capace di inventare storie lunghe, varie e fantastiche.

Sommario tecnico

Titolo: Rooted Absorbed Prefix Trajectory Balance con Replay Submodulare per l'Addestramento di GFlowNet

1. Il Problema: Collasso dei Modi nei GFlowNet basati su LLM

I Generative Flow Networks (GFlowNets) sono progettati per apprendere politiche stocastiche che campionano oggetti completi (come sequenze di testo o molecole) con una probabilità proporzionale alla loro ricompensa. Quando applicati ai Large Language Models (LLM) in una formulazione ad albero dei prefissi terminabili, i GFlowNets soffrono di due gravi problemi di collasso dei modi (mode collapse):

1. Collasso del Prefisso (Prefix Collapse): L'entropia dei token iniziali crolla drasticamente. Diverse traiettorie terminali ad alta ricompensa condividono prefissi quasi identici, limitando la diversità esplorativa fin dalle prime fasi.
2. Bias di Lunghezza (Length Bias): Il modello tende a generare sequenze sistematicamente troppo corte o troppo lunghe, non allineandosi con la distribuzione reale delle soluzioni ottimali.

Cause identificate:

• Assegnazione del Credito Debole: L'uso di ricompense fornite solo alla fine della traiettoria (terminal-only) crea un segnale di apprendimento ad alta varianza per i passi intermedi, rendendo difficile attribuire il merito ai prefissi corretti.
• Replay Bias: I buffer di esperienza tradizionali, che privilegiano le traiettorie ad alta ricompensa, inducono uno spostamento della distribuzione di addestramento verso un sottoinsieme ristretto di modi, aggravando il collasso.

2. Metodologia Proposta

Gli autori propongono una soluzione composta da due meccanismi complementari: RapTB (per l'assegnazione del credito) e SubM (per la gestione del replay).

A. RapTB: Rooted Absorbed Prefix Trajectory Balance

RapTB modifica l'obiettivo di addestramento per fornire segnali di apprendimento densi a livello di prefisso senza introdurre instabilità.

• Prefissi Radicati (Rooted Prefixes): A differenza del Subtrajectory Balance (SubTB) che impone vincoli su finestre sovrapposte arbitrarie (causando condizioni al contorno conflittuali), RapTB vincola solo i residui che partono dalla radice dell'albero ($s_0$). Questo elimina le condizioni al contorno eterogenee che causano il drift della probabilità di terminazione.
• Ricompense Assorbite dai Suffix (Absorbed Suffix Rewards): Per ridurre la varianza, la ricompensa terminale viene "assorbita" e propagata all'indietro verso i prefissi intermedi. Viene calcolato un target di ricompensa per ogni prefisso $s_{0:k}$ combinando:
  • Un backup Max ($u^{max}_k$): la massima ricompensa osservata nel suffisso.
  • Un backup Soft ($u^{soft}_k$): una media logaritmica pesata delle ricompense del suffisso, con penalità per la distanza.
  • Il target finale è una combinazione di questi due.
• Obiettivo Finale: RapTB mantiene il Trajectory Balance (TB) globale come ancora principale (per garantire la consistenza con la distribuzione target) e aggiunge un termine ausiliario che guida l'apprendimento dei prefissi radicati utilizzando i target assorbite, riducendo la varianza senza spostare la distribuzione di terminazione.

B. SubM: Submodular Replay

Per contrastare lo spostamento della distribuzione causato dal replay tradizionale, viene introdotta una strategia di aggiornamento del buffer basata sulla massimizzazione submodulare.

• Obiettivo: Selezionare un sottoinsieme di traiettorie dal buffer corrente e dal nuovo batch che massimizza una funzione submodulare composta da:
  1. Qualità/Ricompensa: Mantenere le traiettorie ad alto punteggio.
  2. Diversità: Massimizzare la copertura dello spazio degli stati (usando la funzione facility-location con similarità Tanimoto per SMILES o Jaccard per testo).
  3. Copertura della Lunghezza: Garantire che il buffer contenga esempi di diverse lunghezze, prevenendo il bias verso sequenze corte o lunghe.
• Efficienza: L'algoritmo utilizza una strategia greedy con garanzia di near-ottimalità, con un costo computazionale trascurabile.

3. Contributi Chiave

1. Caratterizzazione del Collasso: Identificazione empirica del collasso dei modi negli LLM-GFlowNet come una combinazione riproducibile di collasso del prefisso e bias di lunghezza, guidata da un'assegnazione del credito ad alta varianza e bias di replay.
2. RapTB: Un nuovo obiettivo che combina vincoli di prefisso radicati con ricompense assorbite dai suffissi. Fornisce segnali di addestramento densi, riduce la varianza e previene il termination drift (sbilanciamento delle probabilità di fine sequenza) tipico dei metodi SubTB applicati agli LLM.
3. SubM: Una strategia di refresh del buffer di esperienza che bilancia ricompensa, diversità e copertura delle lunghezze in un unico obiettivo submodulare, migliorando la stabilità e la copertura dell'esplorazione.
4. Validazione Empirica: Dimostrazione che la combinazione RapTB + SubM supera gli stati dell'arte in termini di qualità, diversità e validità su compiti complessi.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre compiti principali: generazione di molecole (SMILES), generazione di espressioni aritmetiche (Expr24) e generazione di frasi (CommonGen).

• Generazione di Molecole (SMILES):

  • RapTB + SubM ottiene il miglior compromesso tra qualità (QED) e diversità, mantenendo un'alta validità chimica (>98%).
  • Il metodo TB standard soffre di collasso del prefisso e bias verso sequenze corte.
  • Il metodo SubTB mostra un drastico crollo della validità chimica a causa del drift di terminazione.
  • RapTB + SubM riesce a generare molecole valide anche su orizzonti temporali lunghi (fino a 15 token), dove TB fallisce completamente.

• Expr24 (Espressioni Matematiche):

  • Il compito è sparsamente ricompensato (solo se il risultato è 24).
  • RapTB + SubM raddoppia la copertura normalizzata delle soluzioni uniche rispetto ai baseline migliori, mantenendo un'accuratezza quasi perfetta (>99%).
  • L'analisi del log-probability di terminazione mostra che SubTB soffre di un drastico termination drift (valori estremamente negativi), mentre RapTB mantiene una calibrazione corretta.

• CommonGen (Generazione di Testi):

  • Dimostra che RapTB mantiene la calibrazione naturale della lunghezza e delle probabilità di terminazione, mentre SubTB devia catastroficamente verso lunghezze massime artificiali.

5. Significato e Impatto

Questo lavoro risolve una limitazione fondamentale nell'applicazione dei GFlowNet agli LLM: la tensione tra la necessità di segnali di apprendimento densi (per evitare l'alta varianza) e la stabilità della struttura dell'albero dei prefissi.

• Robustezza: RapTB fornisce un approccio stabile che non richiede di sacrificare la consistenza globale (TB) per ottenere segnali locali.
• Esplorazione Efficiente: SubM dimostra che la diversità e la copertura dello spazio delle soluzioni possono essere gestite attivamente durante l'addestramento, prevenendo il fenomeno "rich-get-richer" (i ricchi diventano più ricchi) tipico dei replay buffer basati solo sulla ricompensa.
• Applicabilità: Le tecniche proposte sono direttamente applicabili alla scoperta di farmaci (generazione di molecole), alla sintesi di codice e ad altri compiti di generazione sequenziale dove la diversità e la validità sono critiche.

In sintesi, il paper propone un framework unificato che combina un'architettura di obiettivo di addestramento migliorata (RapTB) con una strategia di gestione dei dati (SubM) per abilitare l'addestramento robusto ed efficiente di GFlowNet su modelli linguistici su larga scala.