Bayesian Flow Is All You Need to Sample Out-of-Distribution Chemical Spaces

Dit artikel toont aan dat het ChemBFN-model, een Bayesiaanse flow-netwerk met een semi-autoregressieve strategie en versterkingslering, superieur is aan bestaande methoden voor het genereren van hoogwaardige moleculen buiten de trainingsverdeling voor *de novo* drugdesign.

De kern

Bayesian Flow: De Magische Deeltjesversneller voor Nieuwe Medicijnen

Stel je voor dat je een gigantische bibliotheek hebt, vol met recepten voor medicijnen. De meeste recepten in deze bibliotheek zijn al bekend en getest door wetenschappers. Maar wat als je een hele nieuwe medicijnrecept wilt bedenken dat beter werkt dan alles wat we nu hebben? Dat is het probleem waar deze wetenschappers mee worstelen: hoe vind je iets nieuws in een ruimte die zo groot is dat je er nooit doorheen kunt lopen?

Deze paper introduceert een nieuwe manier om die nieuwe recepten te vinden, genaamd Bayesian Flow Networks (BFN). Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Kopieer-App"

De meeste huidige AI-modellen voor het ontwerpen van moleculen werken als een perfecte fotokopieerapp. Ze kijken naar de bestaande recepten (de trainingsdata) en proberen zo dicht mogelijk bij die originele kopieën te blijven.

• Het nadeel: Als je een fotokopieerapp vraagt om iets nieuws te maken, zal hij altijd iets produceren dat op het origineel lijkt. Hij durft niet ver genoeg af te wijken om echt innovatieve, betere medicijnen te vinden. Dit noemen ze "in-distribution" generatie (binnen de bekende grenzen).

2. De Oplossing: De "Magische Klei"

De auteurs tonen aan dat hun nieuwe model, ChemBFN, werkt als magische klei in plaats van een fotokopieerapp.

• In plaats van alleen te kijken naar wat er al is, leert dit model de regels van de klei zelf. Het begrijpt hoe atomen zich kunnen verbinden, maar het is niet bang om de klei in vormen te duwen die nog nooit eerder zijn gemaakt.
• Dit maakt het model van nature goed in het vinden van gebieden buiten de bekende bibliotheek ("out-of-distribution"). Het durft te experimenteren.

3. De Versnellers: Hoe maak je het sneller en slimmer?

Het model was al goed, maar het was nog wat traag en soms produceerde het "rommel" (ongeldige moleculen). De auteurs hebben twee slimme trucjes toegevoegd:

• De "Feedback-Coach" (Reinforcement Learning):
  Stel je voor dat je een kind leert tekenen. Als het kind een lelijke klad maakt, zeg je: "Nee, dat is geen hond." In dit model krijgt de AI een kleine "schop onder de kont" (een beloning) als ze een geldig, goed molecuul maakt. Hierdoor leert het model sneller om alleen de goede vormen te maken en minder tijd te verspillen aan onzin.

• De "Express-trein" (ODE-solver):
  Normaal gesproken moet je een molecuul stap voor stap "ontmaskeren", alsof je een beeld langzaam uit de mist haalt. Dit duurt lang (duizenden stappen). De auteurs hebben een nieuwe methode bedacht die werkt als een express-trein. In plaats van elke kleine stap te nemen, rijdt de trein in grote sprongen naar het einddoel. Hierdoor kunnen ze in 10 tot 100 stappen doen wat anders 1000 stappen kostte. Dit maakt het mogelijk om zelfs op een gewone laptop te werken, zonder dure supercomputers.

4. De "Leesrichting" (Semi-autoregressief)

Normaal leest een AI een zin (of molecuul) als een boek: van links naar rechts, maar hij kijkt ook terug naar wat er al geschreven is én vooruit naar wat er nog komt.
De auteurs hebben de AI echter een causale bril opgezet. De AI leest nu alleen naar voren, alsof hij een verhaal schrijft waarbij hij niet mag kijken naar de toekomst.

• Waarom is dit slim? Dit dwingt de AI om zich te concentreren op de lokale structuur (hoe atomen direct naast elkaar zitten) in plaats van te proberen het hele molecuul in één keer te raden. Het resultaat? De AI wordt creatiever en maakt betere, nieuwere structuren die verder afwijken van de bekende voorbeelden.

5. De Resultaten: Nieuwe Werelden

Toen ze dit model testten, gebeurde er iets wonderlijks:

• Het vond moleculen die beter werkten (beter tegen ziektes) dan alles wat in de trainingsdata zat.
• Het vond moleculen die makkelijker te maken waren in het lab.
• Het deed dit ook voor eiwitten (grote moleculen die het lichaam gebruikt), wat een enorme stap is voor het ontwerpen van nieuwe biologische medicijnen.

Conclusie

Kortom: Deze paper zegt dat we niet hoeven te wachten tot AI perfect is in het kopiëren van het oude. Met Bayesian Flow hebben we een tool die van nature creatief is, snel leert, en ons helpt om de grenzen van de chemie te verleggen. Het is alsof we van een fotokopieerapp zijn gegaan naar een magische 3D-printer die nieuwe, betere medicijnen uit het niets kan toveren.

Technische samenvatting

Probleemstelling

Het genereren van nieuwe moleculen met eigenschappen die beter zijn dan die van de trainingsdata (zogenaamde Out-of-Distribution of OOD-generatie) is cruciaal voor de novo drugdesign. Bestaande methoden, zoals diffusiemodellen (Diffusion Models), zijn echter vaak beperkt in deze taak. Deze modellen zijn ontworpen om de verdeling van de trainingsdata zo nauwkeurig mogelijk na te bootsen, wat leidt tot:

1. Moeite met het genereren van hoogst nieuwe monsters met de gewenste eigenschappen.
2. Moeilijkheden bij multi-objectieve optimalisatie.
3. Een neiging om naar valse positieve gebieden te springen bij overconfidente geleiding.

Het doel van dit onderzoek is een model te ontwikkelen dat intrinsiek in staat is om hoge-kwaliteit OOD-monsters te genereren die buiten de trainingsruimte vallen, maar wel chemisch betekenisvol blijven.

Methodologie

De auteurs gebruiken ChemBFN (Chemical Bayesian Flow Network), een model dat gebaseerd is op Bayesian Flow Networks (BFN). In tegenstelling tot diffusiemodellen, die een omgekeerde stochastische differentiaalvergelijking (SDE) simuleren, optimaliseert BFN direct de parameters van een verdeling naar een informatieve richting. Dit maakt het toepasbaar op zowel continue als discrete data.

Om de prestaties van ChemBFN voor OOD-generatie te verbeteren en de sampling te versnellen, introduceert het artikel drie kernmethodes:

1. Versterkingsleer (RL) Term:
   Een extra term wordt toegevoegd aan de trainingsverliesfunctie, geïnspireerd door het REINFORCE-algoritme. Deze term straalt de kans dat de output op elk tijdstip $t$ een geldig molecuul is, om zo de validiteit van de gegenereerde SMILES-strings te verhogen.

2. ODE-achtig Sampling Proces:
   Om het aantal benodigde sampling-steps drastisch te verminderen (van 1000 naar bijvoorbeeld 10 of 100), wordt een Ordinary Differential Equation (ODE) solver-achtige aanpak gebruikt in de latente ruimte. Door de randomiteit te schalen met een temperatuurcoëfficiënt ($\tau$), wordt de validiteit van de gegenereerde objecten verbeterd zonder de diversiteit volledig te verliezen.

3. Semi-Autoregressieve (SAR) Strategie:
   Hoewel BFN-modellen oorspronkelijk bidirectioneel werken (gebruikmakend van zowel linker- als rechtercontext), introduceert het artikel een semi-autoregressieve aanpak. Hierbij worden causale masks toegepast in de attention-mechanismen, zodat een token alleen wordt bijgewerkt op basis van voorgaande tokens (niet toekomstige).

   • Er worden vier strategieën getest die variëren in het gebruik van SAR tijdens training en sampling (Normal/Normal, Normal/SAR, SAR/Normal, SAR/SAR).
   • De theorie stelt dat SAR de "locality" van het model versterkt, wat essentieel is voor het leren van substructuren en het combineren ervan tot nieuwe, unieke moleculen.

Belangrijkste Bijdragen

• Intrinsieke OOD-capaciteit: Het bewijs dat BFN, en specifiek ChemBFN, van nature een geschikte sampler is voor OOD-generatie, zonder complexe modificaties.
• SAR-Integratie: De introductie van een semi-autoregressieve strategie die de prestaties verbetert en de modelcapaciteit om buiten de trainingsverdeling te genereren, versterkt.
• Efficiëntie: Een combinatie van RL en ODE-like sampling die het aantal stappen reduceert tot ~10-100, waardoor generatie mogelijk wordt op laptops zonder GPU's.
• Theoretische Analyse: Een wiskundige onderbouwing waarom BFN met een specifieke nauwkeurigheidsplanning ($\beta(t)$) en SAR leidt tot betere locality en creatieve generatie vergeleken met Masked Diffusion Language Models (MDLM).

Resultaten

Het model is getest op zowel kleine moleculen (MOSES, GuacaMol, ZINC250k datasets) als eiwitsequenties.

• Kleine Moleculen (Unconditional):

  • De combinatie van RL en ODE-sampling behoudt hoge validiteit en diversiteit, zelfs bij zeer weinig stappen (10 stappen).
  • SAR-strategieën leiden tot een grotere afstand tot de trainingsdata (hoge FCD-waarden), wat aangeeft dat het model succesvol nieuwe ruimtes verkent.

• Kleine Moleculen (Conditional / Multi-objective):

  • Bij geleide generatie (gericht op hoge QED, lage synthetische toegankelijkheid SA, en goede docking scores DS voor 5 verschillende proteïnen) presteert ChemBFN met SAR-strategieën superieur aan State-of-the-Art (SOTA) modellen zoals REINVENT, MORLD, en MOOD.
  • Novel Hit Ratio: ChemBFN (met SAR) behaalde de beste resultaten op 4 van de 5 doelen.
  • Docking Scores: Alle ChemBFN-modellen overtroffen SOTA-methoden in het genereren van moleculen met de beste 5% docking scores.
  • Het gebruik van SELFIES (in plaats van SMILES) in combinatie met RL en ODE verhoogde de "novel hit ratio" aanzienlijk (van <6% naar >25%).

• Eiwitgeneratie:

  • Het model slaagde erin om eiwitsequenties te genereren met een hoger percentage $\beta$-sheets en een groter oppervlak (SASA) dan de trainingsdata, terwijl de "naturalness" (log-likelihood) vergelijkbaar bleef met natuurlijke eiwitten.

Significantie

Dit onderzoek toont aan dat Bayesian Flow Networks een krachtig alternatief zijn voor diffusiemodellen in het domein van de chemische generatie. De belangrijkste implicaties zijn:

1. Versnelling van Drugdesign: De mogelijkheid om snel (in enkele stappen) en effectief nieuwe, hoogwaardige kandidaat-moleculen te genereren die verder liggen dan de huidige trainingsdata, versnelt het de novo ontwerpproces aanzienlijk.
2. Kostenefficiëntie: Door de sampling-steps drastisch te verlagen, wordt de berekeningskosten verlaagd, waardoor geavanceerde generatieve modellen toegankelijker worden.
3. Universele Toepasbaarheid: De methode werkt zowel voor kleine organische moleculen als voor complexe biologische systemen zoals eiwitten, wat het een veelzijdig hulpmiddel maakt voor materiaalwetenschap en biologie.

Samenvattend bewijst dit werk dat "Bayesian Flow" voldoende is om effectief te samplen in ongeziene chemische ruimtes, mits gecombineerd met de juiste strategieën (SAR, RL, en ODE-sampling).