Discrete Bayesian Sample Inference for Graph Generation

Dit artikel introduceert GraphBSI, een nieuw een-staps graf-generatiemodel gebaseerd op Bayesiaanse Steekproefinference dat via een stochastische differentiaalvergelijking discrete grafstructuren effectief genereert en state-of-the-art prestaties behaalt op moleculaire en synthetische benchmarks.

De kern

Stel je voor dat je een meesterchef bent die nieuwe, unieke recepten moet bedenken. Maar in plaats van ingrediënten zoals bloem en suiker, moet je moleculen of netwerken creëren. Deze zijn heel lastig te maken omdat ze uit losse, discrete stukjes bestaan (zoals atomen die wel of niet aan elkaar zitten) en geen vaste volgorde hebben.

Deze paper introduceert een nieuwe, slimme manier om deze complexe structuren te "koken" met een model genaamd GraphBSI. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude probleem: Het "Gokken"

Traditionele modellen proberen een nieuw molecuul te maken door één voor één atomen te plaatsen, alsof je een puzzel probeert te leggen terwijl je blindelings gokt of de volgende stukje past. Als je een fout maakt, moet je vaak helemaal opnieuw beginnen. Dit is traag en leidt vaak tot onbruikbare resultaten.

2. De nieuwe aanpak: GraphBSI (De "Geloofsbord" methode)

In plaats van direct het eindresultaat (het molecuul) te tekenen, begint GraphBSI met een geloof of een verwachting.

• De Metafoor van het Wazige Foto: Stel je voor dat je probeert een foto te maken van een onbekend dier. Aan het begin is je foto volledig wazig en vol ruis; je ziet alleen een vage vlek.
• Het Geloof: GraphBSI houdt niet de foto zelf vast, maar een lijst met vermoedens over hoe die foto eruit zou kunnen zien. Dit is hun "geloof".
• Het Verbeteren: In plaats van de foto direct te scherpen, past het model zijn vermoedens aan. Het zegt: "Oké, op basis van de ruis die ik zie, denk ik dat er waarschijnlijk een poot is hier, en een staart daar."
• Iteratief Scherpstellen: Het model doet dit stap voor stap. Het neemt een kleine "ruis" (een foutje of een hint), kijkt of zijn vermoeden klopt, en past zijn geloof aan. Naarmate het proces vordert, wordt het wazige beeld steeds scherper, totdat het vermoeden zo sterk is dat het model zeker weet: "Ah, dit is een koe!" en het plaatst de atomen op de juiste plek.

3. De Wiskundige Magie: De "Rijbaan" en de "Stuifmeel"

De auteurs beschrijven dit proces met wiskunde die lijkt op een Stochastische Differentiaalvergelijking (SDE). Laten we dit simpel houden:

• De Rijbaan (Deterministisch): Je kunt je voorstellen dat je een auto rijdt op een perfect gladde weg naar je bestemming. Dit is snel en voorspelbaar, maar als je een verkeerde afslag neemt, kom je vast te zitten.
• De Stuifmeel (Stochastisch/Chaos): Nu voegen ze een beetje "wind" of "ruis" toe. Je auto wordt nu een beetje op en neer geschud. Klinkt gek, maar dit helpt! Als je een verkeerde afslag neemt, kan die wind je terugblazen naar de goede weg.
• De Gouden Middenweg: GraphBSI heeft een knop (de parameter $\gamma$) waarmee ze precies kunnen instellen hoeveel "wind" er moet waaien.
  • Te weinig wind? Je blijft vastzitten in een slecht ontwerp.
  • Te veel wind? Je auto waait weg en je maakt geen enkele vooruitgang.
  • De paper laat zien dat met de juiste hoeveelheid "wind" (ruis), het model veel betere en creatievere moleculen maakt dan modellen die alleen op de gladde weg rijden.

4. Waarom is dit belangrijk?

Dit model is een doorbraak voor twee redenen:

1. Snelheid en Kwaliteit: Het kan binnen heel weinig stappen (soms maar 50) een perfect, chemisch geldig molecuul maken. Andere modellen hebben vaak honderden of duizenden stappen nodig.
2. Flexibiliteit: Het werkt niet alleen voor moleculen (voor nieuwe medicijnen), maar ook voor andere netwerken, zoals sociale netwerken of transportroutes.

Samenvatting in één zin

GraphBSI is als een slimme chef die niet blindelings probeert een gerecht te koken, maar eerst een duidelijk plan (een geloof) opbouwt en dat plan stap voor stap verfijnt met de juiste hoeveelheid creativiteit (ruis) om perfect nieuwe moleculen te "bakken" zonder te verbranden.

Het resultaat? Een model dat de staat van de kunst (SOTA) heeft bereikt op de belangrijkste benchmarks voor het genereren van moleculen, wat een enorme stap voorwaarts is voor de ontdekking van nieuwe medicijnen en materialen.

Technische samenvatting

Probleemstelling

Het genereren van graf-structuren (zoals moleculen, kennisgrafieken en sociale netwerken) is cruciaal voor toepassingen zoals drugontwikkeling en materiaalwetenschap. Echter, grafen vormen een uitdaging voor traditionele generatieve modellen vanwege hun discrete, ongeordende en variabele grootte.

• Discreetheid: Grafen bestaan uit discrete entiteiten (knooppunten en randen), terwijl veel succesvolle generatieve modellen (zoals Diffusion Models) zijn ontworpen voor continue data (zoals afbeeldingen).
• Permutatie-invariantie: Grafen hebben geen natuurlijke volgorde van knooppunten; een model moet invariant zijn tegen permutaties.
• Bestaande oplossingen: Autoregressieve modellen zijn traag en schenden permutatie-invariantie door een specifieke volgorde te vereisen. One-shot modellen (zoals Diffusion en Flow Matching) zijn sneller, maar het hanteren van discrete ruimtes vereist vaak complexe discretisatiestappen of Markov-ketens.

Methodologie: GraphBSI

Het paper introduceert GraphBSI (Graph Bayesian Sample Inference), een nieuw one-shot generatief model dat gebaseerd is op Bayesian Sample Inference (BSI). In plaats van samples direct te evolueren, evolueert het model een belief (overtuiging) over de verdeling van de samples in een continue ruimte van distributieparameters.

1. Kernconcept: Categorical BSI

• Van Samples naar Parameters: Het model werkt niet op de discrete grafen zelf, maar op de parameters van een categorische verdeling (logits) over de mogelijke knooppunt- en randcategorieën.
• Bayese Updates: Het proces wordt gezien als een sequentie van Bayese updates. Het model begint met een brede prior over de logits ($z_0$) en verrijkt deze iteratief met "ruisige metingen" ($y_i$) om de overtuiging over de uiteindelijke graaf te verfijnen.
• Theorema 1: Voor categorische data wordt bewezen dat de posterior-belief parameters ($z_{post}$) eenvoudig worden bijgewerkt door de prior logits op te tellen met de gewogen meting: $z_{post} = z + \alpha y$.
• Training: Het model wordt getraind door een Evidence Lower Bound (ELBO) te maximaliseren, wat resulteert in een verliesfunctie die lijkt op het voorspellen van de oorspronkelijke sample ($x$) vanuit de huidige belief ($z_t$) en tijd ($t$).

2. Formulering als Stochastische Differentiaalvergelijking (SDE)

De auteurs leiden af dat het discrete update-proces convergeert naar een continue SDE wanneer het aantal stappen naar oneindig gaat:
$$dz_t = \beta'(t)f_\theta(z_t, t)dt + \sqrt{\beta'(t)}dW_t$$
Hierbij is $f_\theta$ een neurale netwerkschatter en $W_t$ een Wiener-proces.

3. Gecontroleerde Stochasticiteit (De $\gamma$-parameter)

Een cruciale innovatie is de afleiding van een gegeneraliseerde familie van SDE's met een parameter $\gamma$ die de stochasticiteit regelt, terwijl de marginale verdelingen behouden blijven:
$$dz_t = \beta'(t)f_\theta(z_t, t)dt + \gamma^{-1/2}\beta'(t)\nabla_{z_t} \log p_t(z_t)dt + \sqrt{\gamma\beta'(t)}dW_t$$

• $\gamma = 0$: Resulteert in een deterministische Probability Flow ODE (vergelijkbaar met Flow Matching).
• $\gamma = 1$: Herstelt de originele SDE.
• $\gamma > 1$: Introduceert meer ruis, waardoor het model fouten in eerdere stappen kan corrigeren.
• Discretisatie: De auteurs introduceren twee sampling-algoritmen:
  1. Euler-Maruyama (EM): Standaard numerieke integratie.
  2. Ornstein-Uhlenbeck (OU): Een analytische oplossing binnen elk tijdsinterval door de reconstructie als constant te beschouwen. Dit algoritme is stabieler bij hogere ruisniveaus en vereist minder stappen voor hoge kwaliteit.

4. Toepassing op Grafen

• Grafen worden gemodelleerd als paren $(X, A)$ waarbij $X$ de knooppunten en $A$ de randen zijn, beide gecodeerd als one-hot vectors.
• Het model is permutatie-invariant dankzij het gebruik van een Graph Transformer als reconstructie-netwerk ($f_\theta$).
• Variabele grootte wordt gehanteerd door eerst het aantal knooppunten $N$ te sample en vervolgens de features te genereren, waarbij inactieve knooppunten worden gemaskeerd.

Belangrijkste Bijdragen

1. Theoretische Generalisatie: Afleiding van BSI voor categorische data, wat een vereenvoudigde interpretatie biedt ten opzichte van Bayesian Flow Networks (BFN) en het vermijdt om limietbenaderingen toe te passen bij Bayese updates.
2. SDE-Formulering: Het formuleren van categorische BSI als een SDE en het afleiden van een familie van SDE's met een controleerbare ruisparameter ($\gamma$) die het mogelijk maakt te interpoleren tussen deterministische stromen en stochastische samplers.
3. Efficiëntie en Prestaties: Demonstratie dat GraphBSI state-of-the-art (SOTA) resultaten bereikt met zeer weinig functionele evaluaties (zowel 50 als 500 stappen), wat aanzienlijk efficiënter is dan veel bestaande methoden.

Resultaten

Het model is geëvalueerd op twee belangrijke benchmarks voor moleculaire generatie (GuacaMol en MOSES) en synthetische grafen (Planar, Tree, SBM).

• Moleculaire Generatie (MOSES & GuacaMol):
  • GraphBSI presteert SOTA op de meeste metrieken (Validiteit, Uniekheid, Novelty, FCD, SNN).
  • Met slechts 50 stappen (Ornstein-Uhlenbeck discretisatie) overtreft het bestaande modellen zoals DeFoG en GraphBFN.
  • Bij 500 stappen bereikt het bijna perfecte validiteit (>99%) en overtreft het alle concurrenten op de GuacaMol-benchmarks.
  • De FCD (Fréchet ChemNet Distance) op MOSES wordt drastisch verlaagd van 1.07 (bestaande SOTA) naar 0.72 met de EM-discretisatie.
• Synthetische Grafen:
  • Het model bereikt 100% validiteit op planaire en boom-structuren.
  • Het presteert competitief op de Stochastic Block Model (SBM) taken.
• Ablatie Studies:
  • Ruiscontrole ($\gamma$): Het optimaliseren van het ruisniveau is cruciaal. Een $\gamma = 0$ (deterministisch) presteert slecht, terwijl een gematigde tot hoge ruis (afhankelijk van het aantal stappen) de prestaties aanzienlijk verbetert door foutcorrectie.
  • Discretisatie: De Ornstein-Uhlenbeck (OU) methode is stabieler en presteert over het algemeen beter dan Euler-Maruyama, vooral bij hogere ruisniveaus waar EM instabiel wordt.
  • Precisie-schema: Een exponentieel schema werkt het beste, maar het eindprecisie-niveau is de belangrijkste hyperparameter; deze moet hoog genoeg zijn voor een perfecte reconstructie, maar niet zo hoog dat iteraties worden verspild.

Betekenis en Conclusie

GraphBSI biedt een krachtig alternatief voor bestaande graf-generatiemodellen door de discrete aard van grafen te omzeilen via een continue ruimte van distributieparameters. De belangrijkste doorbraak is de flexibiliteit in stochasticiteit via de $\gamma$-parameter, die het model in staat stelt om te schakelen tussen snelle, deterministische stromen en robuuste, stochastische samplers die fouten kunnen corrigeren.

De resultaten tonen aan dat dit framework niet alleen theoretisch elegant is (met connecties naar BFN en Diffusion), maar ook praktisch superieur is in termen van kwaliteit en efficiëntie voor complexe discrete generatietaken zoals het ontwerpen van nieuwe moleculen. Dit opent de deur voor snellere en betrouwbaardere generatieve AI in de chemie en netwerkwetenschappen.