GFlowState: Visualizing the Training of Generative Flow Networks Beyond the Reward

Probleemstelling

Generative Flow Networks (GFlowNets) zijn een krachtig probabilistisch raamwerk dat objecten genereert evenredig aan een beloningsfunctie (reward function). Ze worden succesvol toegepast in domeinen zoals de ontdekking van nieuwe moleculen en materialen. Echter, het trainen van GFlowNets blijft moeilijk te interpreteren en te diagnosticeren.

Bestaande tools voor machine learning (zoals TensorBoard of Weights & Biases) zijn beperkt tot het volgen van geaggregeerde metrieken (zoals verlies of totale beloning) en bieden geen inzicht in:

1. Hoe het model de steekproefruimte (sample space) verkent.
2. Hoe steekproeftrajecten worden geconstrueerd.
3. Hoe de kansen op het selecteren van bepaalde states veranderen tijdens het trainen.
4. Het detecteren van specifieke trainingsproblemen zoals "mode collapse" (waar het model vastloopt in één hoge-belooningsregio en andere verwaarloost).

De onderliggende structuren van GFlowNets groeien combinatorisch snel, wat handmatige inspectie en conventionele monitoring onmogelijk maakt. Er is een kritieke behoefte aan visuele analysetools die specifiek zijn ontworpen voor de unieke dynamiek van GFlowNets.

Methodologie: GFlowState

De auteurs introduceren GFlowState, een interactief visueel analytics-systeem dat is ontworpen om het trainingsproces van GFlowNets te doorgronden. Het systeem is ontwikkeld in nauwe samenwerking met onderzoekers van het Mila – Quebec AI Institute en bestaat uit vier hoofdweergaven die zijn gekoppeld aan specifieke analyse-taken:

1. Sample Ranking (Stalen Ranking):

   • Een "bump chart" die de best en slechtst gerangschikte gegenereerde objecten toont over de iteraties.
   • Helpt bij het identificeren van momenten waarop het model nieuwe, hoogwaardige regio's ontdekt (dynamiek van exploratie) of wanneer het vastloopt in mode collapse (veel vlakke lijnen).

2. State Projection (Staatprojectie):

   • Visualiseert de verkenning van de steekproefruimte door dimensiereductie (bijv. UMAP) toe te passen op gegenereerde samples.
   • Data wordt geaggregeerd in hexagonale bins.
   • Kleurcodering kan gebaseerd zijn op beloning, verlies, of de correlatie tussen steekproefkans en beloning.
   • Vergelijkt de gegenereerde verdeling met een referentiedataset (validatieset) om te zien of het model de volledige ruimte dekt of bepaalde regio's ondervertegenwoordigt.

3. DAG View (Gerichte Acyclische Graaf):

   • Visualiseert de trajecten als een DAG, waarbij states als knopen en acties als randen worden weergegeven.
   • Gebruikt interactieve expansie (vanaf de wortel) en aggregatie (het samenvoegen van lineaire ketens) om de enorme graaf leesbaar te houden.
   • Hiermee kunnen ontwikkelaars specifieke trajecten inspecteren, fouten in overgangen identificeren en begrijpen hoe samples worden opgebouwd.

4. Transition Heatmap (Overgangswarmtekaart):

   • Toont overgangen (edges) gerangschikt op een specifieke metriek (bijv. overgangskans, variantie of frequentie) over de tijd (iteraties).
   • Helpt bij het detecteren van verschuivingen in het beleid (policy) en het identificeren van "mode collapse" door te zien welke overgangen stoppen of dominant worden.

Technische Implementatie:

• Gebouwd in Python met Plotly Dash.
• Gegevens worden opgeslagen in een SQL-database (knooppunten en randen).
• Het systeem is modulair ontworpen zodat het kan worden geïntegreerd in aangepaste omgevingen door vier specifieke functies te implementeren (database-formaat, feature-representatie, single-state visualisatie, multi-state visualisatie).

Belangrijkste Bijdragen

1. Identificatie van Analyse-taken: Een gestructureerd overzicht van de specifieke behoeften van ontwikkelaars, waaronder het analyseren van trajecten, het verkennen van de gegenereerde steekproefruimte en het monitoren van trainingsdynamiek.
2. GFlowState Prototype: Een interactief systeem dat deze taken adresseert via de vier bovenstaande weergaven, met een sterke focus op de interactie tussen trajecten, states en beloningen.
3. Validatie via Case Studies: Een uitgebreide evaluatie in twee domeinen:
   • Grid Environment: Een synthetisch 2D-rooster om de fundamentele dynamiek te testen.
   • Crystals Environment: Een complex, toepassinggericht scenario voor het genereren van kristalstructuren (samenstelling, ruimtegroep, roosterparameters) gericht op katalysatorontwerp.

Resultaten

De case studies, uitgevoerd met experts van Mila, tonen aan dat GFlowState aanzienlijke verbeteringen biedt ten opzichte van bestaande methoden:

• Snellere Debugging: Experts konden snel iteraties identificeren waarin nieuwe hoog-belooningsmodi werden ontdekt (in de Grid-omgeving) of waar het model vastliep in één specifieke structuur (in de Crystal-omgeving).
• Inzicht in Mode Collapse: Het systeem maakte het mogelijk om te zien hoe het model oorspronkelijk vastliep in de regio dichtstbij de start, en later leerde om diverse regio's te verkennen. De DAG View en Transition Heatmap maakten dit visueel duidelijk.
• Validatie van Verdeling: Met behulp van de State Projection en een validatieset konden experts bevestigen of het model de volledige steekproefruimte dekte of dat er "gaten" waren (ondervertegenwoordigde regio's).
• Workflow-versnelling: De interactieve verkenning van trajecten (DAG View) werd door experts beschreven als intuïtief en een grote versnelling ten opzichte van het handmatig laden van checkpoints en het inspecteren van ruwe data.

Betekenis en Toekomstperspectief

GFlowState vult een cruciale lacune in het ecosysteem van generatieve AI. Het maakt de "black box" van GFlowNet-training transparant door de structurele dynamiek zichtbaar te maken. Dit is essentieel voor de toepassing van GFlowNets in wetenschappelijke ontdekkingen, waar het vinden van diverse, hoogwaardige oplossingen (in plaats van slechts één optimale) kritiek is.

Toekomstige verbeteringen die door de auteurs worden voorgesteld, omvatten:

• Betere visualisatie van overgangen in complexe, heterogene ruimtes (zoals bij kristallen).
• Geavanceerde clustering-mogelijkheden voor states.
• Ondersteuning voor continue omgevingen zonder voorafgaande discretisatie (door dynamische precisie-aanpassing).
• Streaming van de DAG tijdens het trainen (momenteel computatie-intensief).
• Betere schatting van correlatiemetrieken via validatiedata in plaats van alleen trainingsdata.

Samenvattend versterkt GFlowState de interpretatie en diagnose van GFlowNets, wat leidt tot robuustere modellen en snellere doorbraken in domeinen zoals materiaalwetenschap en geneesmiddelenontwikkeling.