AceFF: A State-of-the-Art Machine Learning Potential for Small Molecules

AceFF is een geavanceerd, vooraf getraind machine learning potentiaalmodel dat, gebaseerd op een geoptimaliseerde TensorNet2-architectuur en getraind op een uitgebreide dataset van geneesmiddelachtige verbindingen, DFT-nauwkeurigheid combineert met hoge doorvoersnelheid voor de ontdekking van kleine moleculen.

De kern

AceFF: De "Super-kracht" voor het Ontwerpen van Nieuwe Medicijnen

Stel je voor dat chemici en medicijnontwikkelaars een enorme bouwset hebben om nieuwe medicijnen te maken. Om te weten of een nieuw medicijn werkt, moeten ze precies begrijpen hoe atomen (de bouwstenen) met elkaar bewegen en interageren.

Vroeger hadden ze twee opties, maar beide hadden grote nadelen:

1. De "Snelle maar onnauwkeurige" methode: Dit is als een schets op een servet. Het gaat heel snel, maar de details kloppen niet altijd. In de echte wereld noemen we dit klassieke krachtvelden. Ze zijn snel, maar kunnen soms de complexe chemie van nieuwe medicijnen verkeerd voorspellen.
2. De "Uiterst nauwkeurige maar trage" methode: Dit is als het bouwen van een schaalmodel van elke steen in een kathedraal, tot op de atomaire detail. Het is perfect, maar het duurt eeuwen om één steen te berekenen. Dit is de kwantummechanica (DFT). Voor een heel medicijn is dit te langzaam om te gebruiken.

AceFF is de nieuwe "Gouden Tussenweg". Het is een kunstmatige intelligentie (AI) die de snelheid van de schets combineert met de precisie van het schaalmodel.

---

Wat is AceFF precies?

AceFF is een Machine Learning Potentiaal (MLIP). Je kunt je dit voorstellen als een super-intelligente voorspeller die is getraind op miljoenen voorbeelden van hoe atomen zich gedragen.

• De "Leerling" die alles snapt: De AI is getraind op een enorme dataset van "drug-achtige" moleculen. Het kent de belangrijkste bouwstenen (waterstof, koolstof, stikstof, zuurstof, etc.) en kan zelfs omgaan met geladen moleculen (zoals een batterij die positief of negatief is).
• De "Nieuwe Motor" (TensorNet2): De onderliggende technologie van AceFF is verbeterd. De auteurs noemen dit TensorNet2.
  • Analogie: Stel je voor dat de oude AI (AceFF 1.0) een auto was die goed reed op een rechte weg, maar moeite had met bochten of regen. De nieuwe AceFF 2.0 is diezelfde auto, maar nu met een nieuwe motor en een 4x4-systeem. Het kan nu beter omgaan met complexe situaties, zoals geladen moleculen of vreemde vormen, zonder vast te lopen.

Waarom is dit een doorbraak?

De auteurs hebben AceFF getest tegen andere bekende methoden en het scoort fantastisch op twee belangrijke punten:

1. Snelheid: Het is duizenden keren sneller dan de traditionele, nauwkeurige methoden.
   • Vergelijking: Als de oude nauwkeurige methode een uur nodig heeft om één seconde beweging van een molecuul te simuleren, doet AceFF dat in een fractie van een seconde. Hierdoor kunnen wetenschappers nu simulaties draaien die ze eerder nooit durfden te proberen.
2. Nauwkeurigheid: Het is bijna net zo nauwkeurig als de zware kwantumberekeningen.
   • Vergelijking: Het is alsof je een GPS hebt die niet alleen de snelste route aangeeft, maar ook precies weet waar de glibberige plekken op de weg zijn, zonder dat je urenlang de verkeerskaarten moet bestuderen.

De "Testritten" (Benchmarks)

De auteurs hebben AceFF op de proef gesteld in verschillende scenario's, net als een auto die getest wordt op een circuit:

• De "Twist-test" (Torsion Scans): Ze lieten moleculen draaien en knijpen. AceFF voorspelde precies hoe veel energie daarvoor nodig was, zelfs bij geladen moleculen waar andere AI-modellen (zoals ANI-2x) faalden.
• De "Strakke Knoop" (Strained Conformers): Ze keken naar moleculen die in een onnatuurlijke, strakke vorm gedwongen werden. AceFF hield het hoofd koel en gaf de juiste energie, terwijl andere methoden de verkeerde conclusies trokken.
• De "Grote Stad" (Proteïne-complexen): Ze testten hoe AceFF werkt in een groot systeem: een medicijnmolecuul dat vastzit aan een eiwit in het lichaam. AceFF bleek stabiel en snel genoeg om dit te simuleren, zelfs in combinatie met traditionele methoden voor de rest van het systeem.

Wat betekent dit voor de toekomst?

Met AceFF kunnen medicijnontwikkelaars:

• Sneller nieuwe medicijnen vinden: Ze kunnen duizenden mogelijke moleculen in een dag testen in plaats van in een maand.
• Betere medicijnen maken: Omdat de simulaties nauwkeuriger zijn, is de kans groter dat het medicijn in de praktijk ook echt werkt.
• Kosten besparen: Minder tijd in het lab betekent minder geld uitgeven aan experimenten die mislukken.

Kort samengevat

AceFF is als het nieuwe, slimme navigatiesysteem voor de wereld van medicijnontwikkeling. Het combineert de snelheid van een racewagen met de precisie van een chirurg. Het maakt het mogelijk om complexe moleculaire dansjes te simuleren die voorheen te traag of te moeilijk waren om te berekenen, en opent zo de deur voor snellere en betere medicijnen.

De code en het model zijn nu gratis beschikbaar voor iedereen die het wil gebruiken, zodat de hele wetenschappelijke wereld kan meedansen op dit nieuwe ritme.

Technische samenvatting

Probleemstelling

Atomaire simulaties zijn essentieel voor het begrijpen van moleculaire systemen, maar de keuze van het krachtveld (force field) is kritiek voor de nauwkeurigheid.

• Klassieke moleculaire mechanica (MM): Methoden zoals GAFF en AMBER zijn computerefficiënt, maar missen vaak de voorspellende nauwkeurigheid voor diverse, drug-achtige moleculen, vooral bij zeldzame functionele groepen of wanneer kwantumeffecten en polarisatie een rol spelen.
• Kwantummechanica (QM): Methoden zoals DFT (Density Functional Theory) bieden hoge nauwkeurigheid, maar zijn te duur voor routine simulaties van biomoleculaire systemen.
• Machine Learning Interatomic Potentials (MLIPs): Hoewel MLIPs een veelbelovend alternatief zijn dat QM-nauwkeurigheid benadert tegen een lagere kost, blijven generalisatieproblemen bestaan. Bestaande modellen hebben vaak beperkingen:
  • ANI-2x: Snel, maar beperkt tot 8 elementen en alleen neutrale moleculen.
  • MACE-OFF23 / OrbMol: Zeer nauwkeurig en breed toepasbaar, maar aanzienlijk trager.
  • AIMNet2: Ondersteunt geladen moleculen, maar is langzamer dan de snelste opties.

Er is een behoefte aan een model dat een optimale balans biedt tussen snelheid (voor high-throughput drug discovery) en nauwkeurigheid, met volledige ondersteuning voor essentiële geneesmiddelchemie-elementen en geladen toestanden.

Methodologie

Het paper introduceert AceFF-2, een vooraf getraind MLIP dat is geoptimaliseerd voor kleine moleculen in de geneesmiddelenontwikkeling.

1. Architectuur: TensorNet2
AceFF-2 bouwt voort op de TensorNet-architectuur met een nieuwe versie genaamd TensorNet2. De belangrijkste innovaties zijn:

• Ladingsequilibratie (NQE): In tegenstelling tot eerdere versies die ladingen simpelweg toevoegden, implementeert TensorNet2 een methode geïnspireerd op AIMNet2. Het model voorspelt per atoom partiële ladingen ($\vec{q}_i$) en gewichten ($\vec{w}_i$) via een neuraal netwerk. Vervolgens wordt een Neutral Charge Equilibration (NQE) stap uitgevoerd om te garanderen dat de som van de partiële ladingen gelijk is aan de totale moleculaire lading ($Q$).
• Coulomb-interactie: De voorspelde ladingen worden gebruikt om een langeafstands-Coulomb-energieterm ($E_{Coulomb}$) te berekenen, naast de standaard korteafstands-energie. Dit stelt het model in staat om geladen moleculen en ionische interacties correct te modelleren.
• Equivariantie: De architectuur behoudt O(3)-equivariantie (rotatie en translatie) en gaat over naar SO(3)-equivariantie in interactielagen. Dit is cruciaal voor het onderscheiden van chirale moleculen (enantiomeren), wat essentieel is in de farmacie.

2. Software-optimalisaties
Om de inferentiesnelheid te maximaliseren, zijn er aanzienlijke softwareverbeteringen doorgevoerd in de TorchMD-Net-implementatie:

• NVIDIA Warp Kernels: Aangepaste kernels zijn geïmplementeerd voor de decompositie en compositie van tensorcomponenten (scalair, vector, symmetrisch spoorloos) en de message-passing-routines.
• Efficiëntie: Dit resulteert in een 3x snelheidswinst en een 3x reductie in GPU-geheugengebruik door het vermijden van onnodige uitbreidingen naar $3 \times 3$ matrices en het elimineren van grote tussentijdse tensors voor Autograd.
• CUDA Graphs: Het model is geoptimaliseerd voor CUDA Graphs, wat essentieel is voor lage latentie bij kleine systemen (cruciaal voor MLIP/MM hybride simulaties).

3. Dataset
Het model is getraind op een intern QM-dataset van ongeveer 2 miljoen unieke drug-achtige moleculen (tot 30 atomen) met 12 miljoen conformaties.

• Elementen: H, B, C, N, O, F, Si, P, S, Cl, Br, I.
• Ladingen: [-2, -1, 0, +1, +2].
• Theorie: DFT-niveau $\omega$B97M-V/def2-TZVPPD.

Belangrijkste Bijdragen

1. AceFF-2 Model: Een state-of-the-art MLIP dat specifiek is ontworpen voor drug discovery, met ondersteuning voor alle relevante elementen en geladen toestanden.
2. TensorNet2 Architectuur: Een verbeterde versie van TensorNet met geavanceerde ladingsequilibratie voor betere modellering van elektrostatische interacties.
3. Hoge Snelheid: Door geoptimaliseerde kernels en CUDA Graphs ondersteuning, is AceFF-2 een van de snelste nauwkeurige MLIPs, geschikt voor high-throughput toepassingen.
4. Uitgebreide Validatie: Een robuust benchmarking-protocol dat AceFF-2 vergelijkt met toonaangevende modellen (ANI-2x, AIMNet2, OrbMol, MACE, UMA) en traditionele methoden (GAFF, GFN2-XTB).
5. Open Source Beschikbaarheid: De modelgewichten, inferentiecode en tutorials zijn beschikbaar via HuggingFace en GitHub, inclusief integratie met ASE en OpenMM-ML.

Resultaten

De prestaties van AceFF-2 zijn getest op diverse benchmarks:

• Torsie-Scans (Sellers & Behara):

  • AceFF-2 behaalt een zeer hoge nauwkeurigheid, direct na OrbMol (de meest nauwkeurige, maar langzamere, concurrent).
  • Het presteert aanzienlijk beter dan ANI-2x en semi-empirische methoden (GFN2-XTB).
  • In tegenstelling tot AceFF-1.0 en MACE-OFF23 (die voor neutrale moleculen zijn ontworpen), toont AceFF-2 een gebalanceerde prestatie voor zowel neutrale als geladen moleculen.

• Gestrande Conformaties (Wiggle150):

  • AceFF-2 verbetert significant ten opzichte van AceFF-1.0 (MAE daalt van 2.73 naar 1.76 kcal/mol).
  • Het presteert beter dan alle geteste MLIPs (behalve OrbMol en UMA) en semi-empirische methoden, ondanks dat het is getraind op een kleiner dataset dan OMol25.

• Drug-achtige Moleculen (Schrödinger Ligands):

  • Test op grote, onbekende moleculen (>30 atomen). AceFF-2 toont een sterke generalisatievermogen.
  • De krachtfouten (force errors) liggen voor de meeste moleculen onder de 0.05 eV/Å, wat een verbetering is ten opzichte van AceFF-1.0, vooral voor geladen species.

• Stabiliteit en Snelheid in MD:

  • Potentiële Energie Oppervlak: AceFF-2 volgt de DFT-curve nauwkeurig, zelfs bij extreme uitrekkingen en compressies, en toont geen onfysische stappen (in tegenstelling tot OrbMol).
  • Hybride MLIP/MM Simulaties: In simulaties van een proteïne-ligand complex (Tyk2) is AceFF-2 stabiel met tijdstappen van 1fs en 2fs.
  • Snelheid: AceFF-2 is aanzienlijk sneller dan MACE en AIMNet2. Hoewel het iets trager is dan AceFF-1.0 (door de extra ladingstermen), is het nog steeds zeer snel (36.7 ns/day op een RTX4090 voor een MLIP/MM systeem).
  • Batching: Door het gebruik van torch.compile en batching kunnen meerdere conformaties gelijktijdig worden geminimaliseerd, wat de doorvoer met bijna een factor 10 verhoogt voor batches van 32 conformaties.

• Moleculaire Minimalisatie (Platinum Diverse):

  • AceFF-2 slaagt erin om 99.93% van de conformaties binnen 10 Å te minimaliseren.
  • Het faalt alleen bij extreem geladen moleculen (lading > -2), wat buiten het trainingsdomein valt. OrbMol slaagt hierin vanwege zijn grotere dataset, maar AceFF-2 is sneller en toont vergelijkbare nauwkeurigheid voor de meeste gevallen.

Betekenis en Conclusie

AceFF-2 vult een cruciale kloof in het landschap van MLIPs. Het combineert de snelheid die nodig is voor high-throughput screening en routine moleculaire dynamica met de nauwkeurigheid van DFT-niveau, specifiek voor de complexe en geladen moleculen die voorkomen in de geneesmiddelenontwikkeling.

• Strategische Positie: Het model bevindt zich op het Pareto-front tussen de extreem snelle maar beperkte ANI-2x en de uiterst nauwkeurige maar trage OrbMol/MACE-modellen.
• Toepasbaarheid: Het maakt nauwkeurige simulaties van kleine drug-achtige moleculen en hybride MLIP/MM systemen (waarbij alleen het ligand met MLIP wordt gemodelleerd) haalbaar.
• Impact: Door de beschikbaarheid van een robuust, snel en nauwkeurig krachtveld, kan AceFF-2 de efficiëntie en voorspellende kracht van computergestuurde drug discovery aanzienlijk verbeteren, terwijl het de afhankelijkheid van kostbare QM-berekeningen vermindert.