OXtal: An All-Atom Diffusion Model for Organic Crystal Structure Prediction

In dit paper wordt OXtal geïntroduceerd, een groot all-atoom diffusiemodel dat experimenteel realiseerbare 3D-molecuulkristalstructuren voorspelt uit 2D-chemische grafieken en hiermee aanzienlijk beter presteert dan eerdere methoden voor kristalstructuurvoorspelling.

De kern

OXTAL: De Digitale Kristalbouwer

Stel je voor dat je een 2D-tekening van een legpuzzel hebt. Je ziet de vorm van de stukjes en hoe ze aan elkaar verbonden zijn, maar je weet niet hoe ze in het echt in elkaar passen om een compleet, driedimensionaal plaatje te vormen. In de chemie is dit precies het probleem: wetenschappers hebben de platte chemische formule van een molecuul (zoals een medicijn of een materiaal), maar ze weten niet hoe die moleculen zich in de echte wereld stapelen tot een kristal.

Dit is belangrijk, want de manier waarop moleculen zich stapelen, bepaalt of een medicijn werkt, of een batterij lang meegaat, of een scherm helder is. Maar het vinden van deze stapeling is als het zoeken naar een naald in een hooiberg, waarbij de hooiberg zelf ook nog eens voortdurend verandert.

Het probleem: Een gigantische zoektocht
Vroeger probeerden computers dit op te lossen door miljoenen mogelijke stapelingen te "gokken" en ze vervolgens met zware, dure wiskunde (zoals kwantummechanica) te testen. Dit is als proberen een huis te bouwen door elke mogelijke combinatie van bakstenen uit te proberen. Het kost enorm veel tijd, energie en geld. Vaak vinden ze wel een stapeling die werkt, maar niet de beste stapeling die in de natuur voorkomt.

De oplossing: OXTAL
In dit paper introduceren de auteurs OXTAL. Je kunt OXTAL zien als een super-slimme, digitale architect die niet meer hoeft te gokken, maar die gewoon weet hoe het moet.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Kopieer-Plak" Architect (Diffusiemodel)

Stel je voor dat je een foto van een kristal hebt, maar je gooit er zand overheen totdat je alleen nog maar ruis ziet. OXTAL is getraind om dit proces omgekeerd te doen. Het begint met een wazig, willekeurig beeld van atomen en "ontraadt" het langzaam, stap voor stap, totdat er een perfect kristal uit komt.

In plaats van te rekenen aan elke mogelijke stapeling, heeft OXTAL 600.000 echte kristallen bestudeerd. Het heeft geleerd: "Ah, als deze moleculen hier staan, dan staan die andere daar." Het leert de patronen van de natuur, net zoals een kind leert hoe blokken in elkaar passen door te spelen, in plaats van door de zwaartekrachtsformules uit te rekenen.

2. De "Schaal" Methode (S4)

Een groot probleem bij het modelleren van kristallen is dat ze oneindig lijken te herhalen. OXTAL gebruikt een slimme truc genaamd S4 (Stoichiometric Stochastic Shell Sampling).

Stel je voor dat je een grote stad wilt begrijpen. Je hoeft niet de hele stad tegelijk te bekijken. Je kijkt naar één huisje en kijkt dan naar de buren, dan naar de buren van de buren, en zo verder. OXTAL doet dit met moleculen. Het pakt één molecuul en kijkt naar de "schillen" (bollen) van moleculen eromheen. Door te leren hoe deze lokale buren zich gedragen, kan het model de hele stad (het kristal) voorspellen zonder dat het de hele stad in één keer hoeft te "rekenen". Dit maakt het veel sneller en schaalbaarder.

3. Geen "Vaste Regels", maar "Intuïtie"

Veel oude methoden probeerden strikte wiskundige regels (symmetrie) op te leggen aan de computer. OXTAL doet dit niet. In plaats daarvan geeft het de computer een enorme hoeveelheid data en laat het de computer de regels zelf ontdekken. Het is alsof je een kind leert rijden door het in een auto te zetten en te laten oefenen, in plaats van het eerst een heel boek over de wetten van de beweging te laten lezen.

Waarom is dit een doorbraak?

• Snelheid en Kosten: Waar traditionele methoden duizenden euro's aan computerkracht kosten om één kristal te voorspellen, kost OXTAL een fractie daarvan. Het is duizenden keren goedkoper.
• Nauwkeurigheid: OXTAL kan de juiste stapeling vinden in slechts 30 pogingen. Traditionele methoden moeten vaak duizenden pogingen doen om iets vergelijkbaars te bereiken.
• Flexibiliteit: Het werkt niet alleen voor stijve moleculen, maar ook voor "slappe" moleculen die kunnen bewegen, wat veel moeilijker te voorspellen is.

De conclusie
OXTAL is als het geven van een superkracht aan chemici. In plaats van urenlang te zoeken naar de juiste manier om moleculen te stapelen, kan de computer dit in een flits doen. Dit opent deuren voor het sneller ontwikkelen van nieuwe medicijnen, betere batterijen en geavanceerde materialen. Het is een stap van "gokken en rekenen" naar "leren en voorspellen".

Technische samenvatting

1. Het Probleem: Kristalstructuurvoorspelling (CSP)

Het voorspellen van de experimenteel realiseerbare 3D-kristalstructuur van een organisch molecuul, gegeven alleen zijn 2D-chemische graaf, staat bekend als Crystal Structure Prediction (CSP). Dit is een langdurige uitdaging in de computationele chemie.

• Complexiteit: Kristalvorming wordt bepaald door de competitie tussen intramoleculaire interacties (die de conformatie van het molecuul bepalen) en zwakke, langetermijn intermoleculaire krachten (die bepalen hoe moleculen zich in een periodiek rooster packen).
• Energie-landschap: Het energie-landschap is extreem ruw met vele lokale minima. Experimentele structuren vertegenwoordigen vaak lokale minima die kinetisch toegankelijk zijn, niet noodzakelijk het globale thermodynamische minimum.
• Huidige beperkingen:
  • Traditionele methoden (DFT): Gebaseerd op kwantumchemie (zoals DFT) zijn extreem rekenintensief. Ze vereisen het genereren en optimaliseren van duizenden tot honderdduizenden structuren per molecuul, wat ze onpraktisch maakt voor high-throughput screening.
  • Bestaande ML-methoden: Veel machine learning-modellen zijn beperkt tot anorganische kristallen (kleine eenheidscellen, sterke bindingen) of gebruiken vereenvoudigde representaties die de complexe conformationele flexibiliteit en stoichiometrie van organische moleculen niet goed kunnen vangen.

2. Methodologie: OXTAL

OXTAL is een groot, all-atom diffusiemodel (100M parameters) dat de gezamenlijke verdeling van intramoleculaire conformaties en periodieke packing leert, uitsluitend geconditioneerd op de 2D-moleculaire graaf.

A. Architectuur en Representatie

In tegenstelling tot eerdere modellen die expliciete equivariante architecturen en handmatige parameterisatie van roostervectoren gebruikten, kiest OXTAL voor een andere aanpak:

• Cartesiaanse Coördinaten: Het model traint direct op 3D-coördinaten in plaats van roostervectoren.
• Data Augmentation: In plaats van inductieve bias via symmetrie-eigenschappen (equivariantie) in de architectuur te bouwen, worden symmetrieën (rotatie, translatie, permutatie) aangeleerd via data-augmentatie (SE(3) augmentatie).
• Transformer-architectuur: Het model is gebaseerd op de AlphaFold3-architectuur, maar aangepast voor kristallen. Het bestaat uit:
  1. Een Atom Encoder die fysische en structurele informatie (atoomnummer, lading, bindingsinformatie) inbedt.
  2. Een Pairformer Trunk (geïnspireerd door AlphaFold3) die informatie tussen atomen verspreidt via driehoekige self-attention.
  3. Een Diffusion Module (een grote Transformer) die het denoising-proces uitvoert om de uiteindelijke atoomposities te voorspellen.

B. STOICHIOMETRIC STOCHASTIC SHELL SAMPLING (S4)

Een van de kerninnovaties is het S4-trainingschema, ontworpen om langeafstandsinteracties in kristallen efficiënt te modelleren zonder expliciete roosterparameterisatie.

• Concept: Kristallisatie is een lokaal-naar-globaal proces. S4 simuleert dit door "schillen" van moleculen rondom een centraal molecuul te selecteren.
• Implementatie:
  • Er wordt een centraal molecuul $m_c$ gekozen.
  • Moleculen worden gegroepeerd in concentrische schillen $S_k$ op basis van de minimale intermoleculaire afstand ($d_{min}$).
  • Het model traint op een willekeurige subset van deze schillen (een "crop"), waarbij de stoichiometrie (de verhouding van verschillende molecuulsoorten) strikt wordt behouden.
  • Dit verkleint de token-grootte (wat schaalbaarheid bevordert) en fungeert als natuurlijke data-augmentatie, terwijl het toch de langeafstandsperiodiciteit van het kristal leert.
• Voordeel: Het vermijdt de bias van traditionele kNN- of centroid-cropping-methoden en behandelt anisotrope packing-motieven beter.

3. Belangrijkste Bijdragen

1. Eerste schaalbare all-atom diffusiemodel voor organische CSP: OXTAL is het eerste model dat direct kristalpacking genereert vanuit 2D-graaf voor zowel rigide als flexibele moleculen, inclusief co-kristallen en solvaten.
2. S4 Training Scheme: Een nieuwe, rooster-vrije trainingsmethode die langeafstandsinteracties efficiënt captureert zonder de complexiteit van expliciete roosterparameterisatie, wat schaalbaarheid mogelijk maakt.
3. Uitgebreide Dataset: Training op een dataset van 600.000 experimenteel gevalideerde kristalstructuren uit de Cambridge Structural Database (CSD), wat een ongeëvenaarde diversiteit biedt.
4. Prestaties: Het model levert prestaties die meerdere ordes van grootte beter zijn dan bestaande ML-methoden en concurrerend zijn met traditionele DFT-methoden, maar tegen een fractie van de kosten.

4. Resultaten

OXTAL is geëvalueerd op rigide en flexibele moleculen, evenals op de CCDC CSP Blind Tests (de 5e, 6e en 7e editie).

• Nauwkeurigheid:
  • Conformer Recovery: OXTAL herstelt experimentele solid-state conformaties met een RMSD < 0.5 Å voor de meeste rigide en flexibele moleculen.
  • Packing Similarity: Het bereikt een packing-similariteit rate van >80% (PacC) binnen slechts 30 samples.
  • Blind Tests: OXTAL presteert consistent beter dan andere ML-baselines en bereikt vergelijkbare "approximate solve" rates als DFT-methoden, maar gebruikt daarvoor slechts 30 samples in plaats van duizenden.
• Efficiëntie en Kosten:
  • Schaalbaarheid: OXTAL is ordes van grootte goedkoper dan DFT-methoden. Terwijl DFT-baselines in de 7e blind test naar schatting 46 miljoen CPU-core-uren gebruikten, kost OXTAL's inferentie slechts een fractie daarvan (gebaseerd op cloud-prijzen).
  • Sample Efficiency: Het model vindt vaak de juiste packing-motief binnen de eerste 10-30 samples, terwijl DFT-methoden vaak duizenden structuren moeten genereren om één succesvol resultaat te vinden.
• Chemische Interpretatie:
  • Het model kan complexe interacties correct voorspellen, waaronder waterstofbruggen, $\pi$-$\pi$-stacking, en halogeenbindingen.
  • Het is in staat om verschillende polymorfen (verschillende kristalvormen van hetzelfde molecuul) en co-kristallen te genereren, wat aantoont dat het zowel thermodynamische als kinetische landschappen leert.

5. Betekenis en Impact

OXTAL markeert een paradigmaverschuiving in de kristalstructuurvoorspelling:

• Van Search naar Generatie: Het verlegt de focus van het genereren en ranken van enorme aantallen kandidaten (zoals bij DFT) naar het direct genereren van hoogkwalitatieve, experimenteel realistische structuren.
• Toepassingsgebied: De technologie heeft directe implicaties voor de farmaceutische industrie (oplossen van oplosbaarheid, biologische beschikbaarheid en stabiliteit van medicijnen) en de ontwikkeling van organische halfgeleiders en materialen.
• Kostenverlaging: Door de inferentiekosten drastisch te verlagen, maakt OXTAL high-throughput screening van nieuwe materialen en medicijnen haalbaar, wat eerder onmogelijk was door de rekenkosten van kwantumchemische methoden.

Kortom, OXTAL bewijst dat grote, all-atom diffusiemodellen, getraind op enorme datasets met innovatieve sampling-technieken, de complexiteit van organische kristallisatie kunnen modelleren met een nauwkeurigheid en efficiëntie die de huidige state-of-the-art methoden overtreffen.