Zatom-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials

Zatom-1 is het eerste open-source foundationmodel dat generatieve en voorspellende taken voor zowel 3D-moleculen als materialen verenigt via een multimodale flow-matching-architectuur, wat leidt tot verbeterde prestaties en snellere inferentie dan gespecialiseerde modellen.

De kern

Stel je voor dat chemie een gigantische bibliotheek is. In deze bibliotheek staan twee soorten boeken: de ene soort beschrijft kleine, losse moleculen (zoals medicijnen), en de andere soort beschrijft grote, kristalachtige materialen (zoals batterijen of zonnepanelen).

Tot nu toe hadden wetenschappers twee verschillende bibliothecarissen:

1. De ene kon alleen nieuwe medicijnen uitvinden (genereren), maar kon niet goed voorspellen hoe ze werkten.
2. De andere kon alleen de eigenschappen van bestaande materialen voorspellen, maar kon geen nieuwe ontwerpen maken.

Ze spraken niet met elkaar, en ze gebruikten verschillende talen. Dat maakte het vinden van nieuwe uitvindingen traag en lastig.

ZATOM-1 is de nieuwe, super-bibliothecaris.

Hier is hoe dit werkt, vertaald naar alledaagse taal:

1. De "Alles-in-Één" Super-App

ZATOM-1 is de eerste computerprogramma dat zowel nieuwe moleculen kan ontwerpen als hun eigenschappen kan voorspellen, en dat voor zowel kleine medicijnen als grote materialen. Het is alsof je één app hebt die zowel een architect is (die nieuwe huizen tekent) als een ingenieur (die controleert of het dak niet instort).

2. Hoe leert het? (De "Kleurenplaat" Methode)

Stel je voor dat je een kind een tekening van een huis geeft, maar de muren zijn grijs en de ramen zijn leeg. Het kind moet de kleuren invullen.

• De oude manier: Wetenschappers gaven het kind eerst een boek met alleen de muren, dan een boek met alleen de ramen, en daarna een boek met de dakpannen. Het kind moest alles apart leren.
• De ZATOM-1 manier: Dit model leert door te "ontmaskeren". Het krijgt een tekening waar de atomen (de bouwstenen) en hun posities (waar ze staan) een beetje wazig zijn. De taak van het model is om te raden: "Wat zou hier perfect passen?"

Door dit miljoenen keren te doen met zowel medicijnen als materialen, leert het model de "regels van de natuur" zonder dat iemand het hardop uitlegde. Het leert dat een waterstofatoom hier niet past, maar een koolstofatoom daar wel.

3. De "Tijdmachine" voor Ontwerpen

Een van de coolste dingen aan ZATOM-1 is dat het extreem snel is.

• Vroeger: Het ontwerpen van een nieuw materiaal duurde alsof je een auto probeerde te bouwen door elke bout één voor één met de hand vast te draaien. Het kon dagen duren om één ontwerp te testen.
• Nu: ZATOM-1 is als een 3D-printer die in seconden duizenden ontwerpen kan "printen" en direct kan zeggen: "Deze 10 zijn goed, die 990 zijn rot." Het is 10 tot 100 keer sneller dan de beste oude methoden.

4. De "Magische Transfer"

Dit is het meest fascinerende stukje. Omdat ZATOM-1 heeft geleerd hoe kristallen (materialen) werken, wordt het plotseling ook beter in het begrijpen van medicijnen (moleculen).

• Analogie: Stel je voor dat je eerst hebt geleerd hoe je een complex horloge repareert (materialen). Als je daarna een fiets moet repareren (moleculen), begrijp je plotseling veel beter hoe tandwielen en veren werken, zelfs als je nog nooit een fiets hebt gezien.
• In de paper zien ze dat het model, door te oefenen met materialen, medicijnen beter kan voorspellen dan modellen die alleen op medicijnen zijn getraind.

5. Waarom is dit belangrijk?

Vandaag de dag duurt het jaren om een nieuw medicijn of een nieuwe batterij te vinden. Wetenschappers moeten duizenden pogingen doen, veel ervan mislukken.
Met ZATOM-1 kunnen we:

• Medicijnen sneller vinden: Misschien een nieuwe behandeling voor kanker of een virus.
• Betere materialen maken: Batterijen die sneller laden, of materialen die zonnepanelen veel efficiënter maken.
• Kosten besparen: Minder tijd in het lab betekent minder geld uitgeven aan mislukte experimenten.

Samenvattend

ZATOM-1 is als een universele scheikundige-assistent die nooit slaapt. Hij heeft de "taal" van de atoomwereld geleerd door te spelen met duizenden voorbeelden. Hij kan nu niet alleen nieuwe dingen bedenken, maar ook vertellen of die dingen werken, en hij doet dit zo snel dat hij de hele bibliotheek van de chemie in een dag kan doorzoeken.

Het is een grote stap vooruit in hoe we de wereld van morgen ontwerpen, van betere medicijnen tot schoner energie.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "ZATOM-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials", vertaald en samengevat in het Nederlands.

Probleemstelling

Huidige AI-aanpakpen voor chemische modellering zijn vaak gespecialiseerd in één domein (ofwel moleculen, ofwel materialen) en één taak (ofwel generatie, ofwel voorspelling). Deze fragmentatie beperkt het delen van representaties en transfer learning tussen deze domeinen. Bestaande methoden voor 3D-generatie maken vaak gebruik van complexe twee-staps processen (zoals latente diffusie met auto-encoders) of handgemaakte priors, wat leidt tot trage training, trage inferentie en beperkte schaalbaarheid. Er is een gebrek aan een universeel, open-source foundation model dat zowel generatief als voorspellend leren kan voor 3D-moleculen en kristallijne materialen in één architectuur.

Methodologie: ZATOM-1

ZATOM-1 is het eerste end-to-end foundation model dat generatief en voorspellend leren unificert voor 3D-moleculen en materialen. De kern van de methode bestaat uit de volgende componenten:

1. Representatie en Tokenisatie
Het model werkt op het fundamentele niveau van atomaire granulariteit. Het representeert systemen via vijf gedefinieerde modaliteiten:

• Discrete modale: Atoomtypes ($A$).
• Continue modaliteiten: 3D-coördinaten ($X$), fractionele coördinaten ($F$), roosterverlengingen ($L_{len}$) en roosterhoeken ($L_{ang}$).
• Voor niet-periodieke moleculen worden de rooster- en fractionele coördinaten gemaskeerd; voor periodieke materialen worden de 3D-coördinaten gemaskeerd om rotatie-invariantie te garanderen.

2. Architectuur: Trunk-based Flow Transformer (TFT)

• Het model gebruikt een gestandaardiseerde Transformer-architectuur (gebaseerd op Query-Key Normalization, Flash Attention en SwiGLU feed-forward netwerken).
• Het is een multimodale flow matching model. In plaats van een latente ruimte te gebruiken (zoals bij VAEs of Latent Diffusion), voert het model generatieve flow matching direct uit in de Euclidische ruimte ($\mathbb{R}^3$).
• Dit elimineert de noodzaak voor vooraf getrainde auto-encoders, wat aanzienlijke snelheidswinst oplevert.

3. Trainingstraject
Het model wordt in twee fasen getraind:

• Fase 1: Generatief Pre-training (Multimodal Flow Matching):
  • Het model leert een gezamenlijke verdeling van atoomtypes en geometrie door een conditionele stroom te leren die samples van een ruisverdeling naar de data-verdeling transporteert.
  • Het gebruikt Euclidean Flow Matching voor continue variabelen en Discrete Flow Matching voor atoomtypes.
  • De loss-functie combineert MSE voor continue modaliteiten en cross-entropy voor discrete atoomtypes.
  • Dit dient als een universele, zelftoezichtende (label-vrije) initialisatie.
• Fase 2: Voorspellend Finetuning:
  • De gewichten van de "Trunk" (de Transformer-encoder) worden bevroren.
  • Extra "heads" en auxiliary lagen worden toegevoegd om eigenschappen, energieën en krachten te voorspellen.
  • Dit stelt het model in staat om multi-task learning uit te voeren.

4. Sampling

• Sampling gebeurt via iteratieve Euler-stappen (of andere ODE-oplossers) om ruis te verwijderen.
• Het model kan zowel moleculen als kristallijne materialen genereren zonder dat er aparte modellen nodig zijn.

Belangrijkste Bijdragen

1. Eerste Unificatie: ZATOM-1 is het eerste volledig open-source foundation model dat zowel generatie als voorspelling voor 3D-moleculen en materialen in één architectuur combineert.
2. Schaalbaarheid en Snelheid: Door het vermijden van latente ruimtes en het gebruik van een gestandaardiseerde Transformer, bereikt het model een 10-voudige versnelling in inferentie-tijd ten opzichte van state-of-the-art baselines (zoals ADiT) en een 3-voudige reductie in GPU-trainingstijd.
3. Positieve Transfer Learning: Het paper demonstreert voor het eerst dat generatief pre-training op zowel moleculen als materialen leidt tot verbeterde voorspellende prestaties in het andere domein (bijv. het modelleren van materialen tijdens pre-training verbetert de eigendomsvoorspelling voor moleculen).
4. State-of-the-Art Resultaten: Het model bereikt nieuwe staat-van-de-kunst resultaten op generatieve benchmarks (QM9, GEOM-Drugs, LeMat-GenBench) en multi-task eigendomsvoorspelling.

Resultaten

• Generatie van Materialen (LeMat-GenBench): ZATOM-1 genereert stabiele, unieke en nieuwe kristalstructuren. Het presteert beter dan eerdere flow-matching methoden (zoals ADiT) en benadert de prestaties van modellen die complexe relaxatie-stappen gebruiken, maar dan zonder die extra kosten.
• Generatie van Moleculen (QM9 & GEOM-Drugs):
  • Op QM9 bereikt ZATOM-1 een validiteit van ~95% en slaagt ~99% van de gegenereerde moleculen alle PoseBusters sanity checks (fysieke realisme checks).
  • Op GEOM-Drugs (grote drug-achtige moleculen) bereikt het een validiteit van 94%, wat een nieuwe benchmark is voor generatieve modellen op deze schaal.
• Voorspellende Prestaties:
  • ZATOM-1 bereikt state-of-the-art resultaten voor multi-task eigendomsvoorspelling op QM9.
  • Het toont aan dat gezamenlijke pre-training leidt tot betere generalisatie dan domein-specifieke training.
• Efficiëntie:
  • Generatie van 10.000 samples duurt minder dan 4 minuten op één NVIDIA A100 GPU (voor het 300M parameter model), vergeleken met veel langzamere baselines.
  • Het model schaalt voorspelbaar: grotere modelgroottes leiden tot betere prestaties (validiteit en verlies).

Betekenis en Impact

ZATOM-1 markeert een belangrijke mijlpaal in Scientific Machine Learning (SciML). Het bewijst dat generatieve modellering niet alleen nuttig is voor het creëren van nieuwe structuren, maar ook als een krachtige, zelftoezichtende pre-training strategie kan dienen voor voorspellende taken.

• Versnelling van Ontdekking: Door de enorme snelheidswinst en de capaciteit om zowel moleculen als materialen te modelleren, kan dit model de cyclus van chemisch ontwerp en ontdekking aanzienlijk versnellen.
• Universele Representaties: Het succes van cross-domein transfer learning suggereert dat fundamentele chemische kennis gedeeld kan worden tussen moleculen en materialen, wat leidt tot robuustere en generaliseerbaardere AI-modellen.
• Open Science: De beschikbaarheid van de code en gewichten als open-source project democratiseert de toegang tot geavanceerde foundation modellen voor chemici en materialwetenschappers.

Kortom, ZATOM-1 biedt een schaalbare, snelle en krachtige basis voor de toekomst van computationele chemie en materiaalkunde, waarbij generatie en voorspelling niet langer als gescheiden taken worden behandeld, maar als complementaire aspecten van één fundamenteel model.