OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

Dit artikel introduceert OmniMol, een state-of-the-art machine-learned interatomair potentiaal voor kleine moleculen die een Point-Edge Transformer-architectuur en kennisoverdracht vanuit de hoge-energiefysica benut om uitstekende prestaties te bereiken met minimale fine-tuning en uniek snelle inferentie.

De kern

Stel je twee heel verschillende werelden voor: de ene is de chaotische, supersnelle wereld van deeltjesfysica (waar wetenschappers atomen tegen elkaar aan slaan om te zien wat eruit vliegt), en de andere is de ingewikkelde, kleverige wereld van moleculaire chemie (waar atomen aan elkaar plakken om medicijnen, materialen en leven te vormen).

Lange tijd gebruikten wetenschappers in deze twee gebieden volledig verschillende hulpmiddelen om hun werelden te begrijpen. Maar in dit artikel introduceren de auteurs OmniMol, een nieuw hulpmiddel dat probeert de experts in de deeltjesfysica chemie te leren begrijpen door gebruik te maken van een "fundamenteel model" dat ze al hadden gebouwd.

Hier is de eenvoudige uitleg van hoe ze dit deden en wat ze ontdekten:

1. De "Meesterkok"-analogie

Stel je het oorspronkelijke model, genaamd Omnilearned, voor als een meesterkok die jarenlang heeft gekookt met deeltjesjets.

• De ingrediënten: In de deeltjesfysica is een "jet" een sproei van subatomaire deeltjes (zoals protonen en neutronen) die uit een botsing vliegen.
• De vaardigheid: Deze kok leerde patronen in deze sproei's herkennen. Ze weten hoe deeltjes met elkaar interageren, hoe ze clusteren en hoe ze voorspellen wat er als nächst gebeurt. Ze werden getraind op één miljard verschillende deeltjesproei's.

Nu vroegen de auteurs zich af: Kan dezezelfde kok een moleculair maaltje bereiden?

• De nieuwe ingrediënten: In plaats van subatomaire deeltjes zijn de "ingrediënten" atomen (zoals Koolstof, Zuurstof, Waterstof) in een molecuul.
• De uitdaging: Atomen gedragen zich anders dan subatomaire deeltjes, maar ze delen een vergelijkbare structuur: het zijn gewoon punten in de ruimte met specifieke types.

2. De "Universele Vertaler" (De architectuur)

Om dit werk te laten doen, bouwden ze geen nieuwe kok van scratch. Ze namen de bestaande "Meesterkok" (Omnilearned) en gaven hen een nieuwe set gereedschappen:

• De Point-Edge Transformer (PET): Stel je de kok voor die naar een bord eten kijkt. In plaats van maar één ingrediënt tegelijk te bekijken, laat deze tool hen elk ingrediënt tegelijk bekijken en zien hoe elk enkelvoudig ingrediënt met elk ander ingrediënt samenhangt.
• De "Fysica-bias": Dit is de geheime saus. Het model heeft een ingebouwd "regelsboek" dat het vertelt: "Hé, deze twee deeltjes/atomen zitten dicht bij elkaar, dus ze moeten meer aandacht aan elkaar besteden." Dit helpt het model zich te focussen op de belangrijkste relaties zonder verward te raken door ruis.

3. Het experiment: Fine-tuning

De auteurs namen dit voor deeltjes getrainde model en gaven het een "crashcursus" in chemie met behulp van een dataset genaamd oMol (een verzameling van miljoenen moleculen).

• Het doel: Ze wilden dat het model fungeerde als een Machine-Learned Interatomic Potential (MLIP). In gewone taal betekent dit dat het model twee dingen moet voorspellen voor elke groep atomen:
  1. Energie: Hoeveel "lijm" houdt ze bij elkaar?
  2. Kracht: Als je één atoom duwt, hoe hard zal het terugduwen?

4. De resultaten: Snel en verrassend goed

Het artikel vond enkele spannende dingen:

• De "Few-Shot"-superkracht: Normaal gesproken vereist het leren van chemie aan een computer enorme hoeveelheden data. Maar omdat OmniMol begon met de "kennis" van de deeltjesfysica, leerde het chemie zeer snel. Zelfs met een relatief kleine hoeveelheid nieuwe data (zoals 100.000 moleculen), presteerde het bijna even goed als modellen die op miljoenen waren getraind. Het is alsof een meesterkok een nieuwe keuken kan leren met slechts een paar recepten, omdat hij de basis van smaak en hitte al begrijpt.
• Snelheid: OmniMol is ongelooflijk snel. Terwijl andere modellen misschien lang nodig hebben om te berekenen hoe een molecuul beweegt, doet OmniMol dit in een flits. De auteurs merken op dat OmniMol voor elke uur rekentijd drie keer meer moleculen kan simuleren dan sommige van zijn concurrenten.
• De afweging: Toen ze enorme hoeveelheden data hadden (miljoenen moleculen), vervaagde het voordeel van het beginnen met kennis uit de deeltjesfysica enigszins. Dit suggereert dat de "kennis uit de deeltjesfysica" werkt als een sterke voorsprong, maar als je genoeg tijd en data hebt om een model van scratch te trainen, maakt die voorsprong minder uit.

5. Het grote plaatje

Het artikel concludeert dat OmniMol de eerste keer is dat een "fundamenteel model" dat is gebouwd voor één wetenschappelijke discipline (de deeltjesfysica) succesvol is overgedragen naar een volledig andere discipline (chemie).

Ze bewezen dat als je een slim model hebt dat begrijpt hoe punten in de ruimte met elkaar interageren in het ene veld, het kan worden aangepast om te begrijpen hoe punten in de ruimte met elkaar interageren in een ander veld, waardoor tijd en rekenkracht worden bespaard.

Samenvattend: De auteurs namen een super slimme AI die was getraind op botsingen van deeltjes met hoge energie, pasten zijn brein aan om atomen in plaats van deeltjes te begrijpen, en ontdekten dat het een bliksemsnel, uiterst nauwkeurig hulpmiddel werd voor het voorspellen van het gedrag van moleculen, vooral wanneer data schaars is.

Technische samenvatting

Technische Samenvatting: OmniMol

Probleemstelling

Machine learning (ML) heeft de representatie en simulatie van complexe fysische systemen getransformeerd, met name in de deeltjesfysica en moleculaire chemie. Hoewel deze domeinen sterk verschillen in energieschalen, delen ze een fundamentele datastructuur: variabel grote verzamelingen deeltjes (of atomen) in de faseruimte, die effectief gestructureerde puntwolken vormen.

De belangrijkste uitdaging die wordt aangepakt, is de ontwikkeling van efficiënte machine-learned interatomaire potentialen (MLIP's). Traditionele methoden zoals Dichtheidsfunctionaaltheorie (DFT) zijn computatievriendelijk duur, wat grootschalige en langdurige moleculaire dynamica (MD) simulaties beperkt. MLIP's streven ernaar potentie-energieoppervlakken en krachten te benaderen voor een fractie van deze kosten. Het trainen van robuuste MLIP's vereist echter doorgaans enorme datasets en aanzienlijke rekenkracht. Het artikel hypotheert dat een foundation model dat is vooringeschat op puntwolken in de deeltjesfysica (specifiek deeltjesjets) kan worden overgebracht naar moleculaire dynamica, wat de optimalisatie kan versnellen en de nauwkeurigheid in regimes met weinig data kan verbeteren.

Methodologie

Architectuur: Point-Edge Transformer (PET)

OmniMol is gebouwd door Omnilearned aan te passen, een foundation model dat oorspronkelijk was ontworpen voor het classificeren en genereren van deeltjesjets in de hoge-energiefysica (HEP). De kernarchitectuur is een Point-Edge Transformer (PET), die lokale aandacht over $k$-naaste buren koppelt aan globale all-to-all transformerblokken.

Belangrijke architectonische componenten zijn:

1. Input Embeddings: Atomen worden ingebed in een tokenruimte die positionele informatie ($\vec{r}$), discrete atoomnummers ($Z$) en extra kenmerken (lading, spin) combineert.
2. Lokale Aandachtsblok: Voor elk atoom wordt een lokaal buurschap geconstrueerd met behulp van $K$-naaste buren ($K=15$ voor moleculen, vergeleken met $K=10$ voor jets). Paarsgewijze fysische kenmerken worden berekend, inclusief afstandstermen, inverse machten van afstand en geleerde functies van atomaire embeddings. Deze worden verwerkt door een kleine lokale transformer om een lokaal embeddingvector te creëren.
3. Globale Aandacht met Interactieve Bias: Het globale zelf-aandachtsmechanisme integreert een expliciete bias afgeleid van paarsgewijze fysische kenmerken. De attentie-logits worden gemodificeerd als $A^*_{ij} = A_{ij} + B_{ij}$, waarbij $B_{ij}$ een door een MLP ingebedde bias-term is. Deze "interactiematrix-aandachtsbias" injecteert paarsgewijze fysicapriors direct in de transformer, waardoor het netwerk wordt gestuurd naar fysisch betekenisvolle buurschappen zonder expressiviteit te offeren.
4. Output Heads: De generatieve head van Omnilearned wordt hergebruikt voor twee taken:
   • Krachtpredictie: Een permutatie-equivariante head die krachten per atoom voorspelt.
   • Energiepredictie: Een head die energiecorrecties per atoom voorspelt, die worden opgeteld om de totale moleculaire energie te verkrijgen, waarbij extensive priors worden behouden.

Invariantie en Behoudsbeperkingen

Om aan fysische beperkingen te voldoen, adresseren de auteurs twee vereisten:

• Energiebehoud: Krachten worden niet direct voorspeld, maar berekend via backpropagatie van de energie-uitvoer ($\vec{F}_i = \nabla_{\vec{r}_i} E$). Dit zorgt voor exact energiebehoud, maar verhoogt de rekenkosten tijdens training (vereist dubbele backpropagatie). Derhalve wordt deze beperking alleen toegepast op de "kleine" modelvariant.
• Rotatie-Equivariantie: De standaardarchitectuur is niet inherent equivariant omdat ruwe coördinaatverschillen worden ingevoerd in MLP's. Om dit op te lossen, introduceren de auteurs een "equivariante en conservatieve" variant. Deze versie verwijdert directe coördinaatverschiltermen uit de paarsgewijze kenmerken en integreert in plaats daarvan hoekinformatie (cosinussen van hoeken gevormd door vectoren tussen naburige atomen) in het lokale blok. Deze modificatie behoudt equivariantie terwijl het prestatieverlies dat gepaard gaat met het verwijderen van coördinatentermen aanzienlijk wordt teruggewonnen.

Trainings- en Fine-Tuning Strategieën

Het model wordt fine-getuned op de oMol-dataset (specifiek de oMol-25, oMol-4M, oMol-100M en oMol-140M subsets). Twee strategieën worden onderzocht:

1. LoRA (Low-Rank Adaptation): De vooringeschatte PET-ruggegraatgewichten worden bevroren. Low-rank adapters worden alleen geïntroduceerd voor de transformer-lichaamsmatrices ($W_Q, W_K, W_V, W_O, W_{MLP}$), naast het trainen van de moleculaire input-encoders, de bias-MLP en de task-heads. Er wordt ook een "embedding adapting"-laag toegevoegd om geleerde embeddings te modificeren.
2. Volledige Fine-Tuning: Alle gewichten in het lichaam en de input-encoders worden ontdooid en getraind, terwijl de task-heads vanaf nul worden getraind.

Het trainingsdoel minimaliseert de som van de Mean Absolute Errors (MAE) voor energieën en krachten, waarbij krachten zwaarder wegen ($\lambda_F = 10$).

Belangrijkste Resultaten

Prestaties op oMol

• Volledige Training: Wanneer getraind op grote datasets (oMol-4M en oMol-100M/140M), bereikt OmniMol concurrerende prestaties met state-of-the-art MLIP's. Bijvoorbeeld, op oMol-140M bereikt het OmniMol-large model een energie-MAE van 1,04 meV/atoom en een kracht-MAE van 13,59 meV/Å.
• Regime met weinig data: De grootste winsten worden waargenomen wanneer de trainingsdata beperkt is. Wanneer fine-getuned op slechts 100k moleculen of met zeer weinig epochs (2 passes) over oMol-4M, presteren de vooringeschatte OmniMol-varianten aanzienlijk beter dan modellen die vanaf nul worden getraind.
  • Op een subset van 100k verbeterde pre-training de energie-MAE met maximaal 29,4% en de kracht-MAE met 26,9% voor het medium model.
  • Met slechts twee epochs op oMol-4M toonde het medium model een verbetering van 54,6% in energie-MAE en 56,9% in kracht-MAE vergeleken met zijn niet-vooringeschatte tegenhanger.
• Equivariante/Conservatieve Variant: De equivariante en conservatieve modelvariant toont aanzienlijk verbeterde prestaties (vooral voor krachten) in regimes met weinig data, hoewel dit voordeel afneemt naarmate de datasetgrootte toeneemt.

Schaling en Inferentiesnelheid

• Schaling: OmniMol volgt een schone power-law-schaling met modelgrootte, zonder tekenen van verzadiging tot 1 miljard parameters, consistent met recente bevindingen over transformer-gebaseerde MLIP's.
• Inferentiesnelheid: Ondanks grote parameteraantallen demonstreert OmniMol uniek snelle inferentiesnelheden dankzij hardware-optimalisaties voor transformers. Op een A100 GPU voor systemen van ~100 atomen is OmniMol-medium ongeveer 3x sneller dan vergelijkbare Graph Neural Network (GNN) baselines (eSEN-md-d en AllScAIP-md) terwijl het concurrerende nauwkeurigheid behoudt (slechts ~0,7 meV/atoom hogere energiefout dan AllScAIP-md).

Betekenis en Claims

Het artikel claimt de eerste demonstratie van cross-disciplinaire overdracht voor wetenschappelijke puntwolk foundation-modellen te presenteren. Door een model aan te passen dat is vooringeschat op deeltjesjets uit de hoge-energiefysica naar moleculaire dynamica, demonstreren de auteurs dat:

1. Cross-Domein Overdracht Haalbaar is: Een foundation model gebouwd voor deeltjesfysica kan effectief worden overgebracht naar moleculaire chemie, wat suggereert dat de onderliggende puntwolkstructuren leerbare kenmerken delen over sterk verschillende fysische schalen.
2. Inductieve Bias Leerproces Versnelt: De vooringeschatting fungeert als een sterke inductieve bias. Net zoals equivariantie helpt wanneer data schaars is, maakt de "bittere les" van vooringeschatting snelle optimalisatie en verbeterde nauwkeurigheid mogelijk wanneer trainingsdata beperkt is.
3. Efficiëntie: De architectonische overdracht maakt uniek snelle inferentiesnelheden mogelijk, wat cruciaal is voor toepassingen die snelle verkenning van ontwerpruimtes vereisen, zoals de ontdekking van kleine moleculaire geneesmiddelen.

De auteurs concluderen dat hoewel de studie zich richt op MLIP's, de lessen met betrekking tot puntwolk foundation-modellen wijdverspreide bruikbaarheid kunnen hebben in wetenschappelijke domeinen waar systemen worden beschreven als ongeordende verzamelingen van interagerende lichamen. Ze claimen geen universele superioriteit boven alle bestaande methoden in alle regimes, maar benadrukken de specifieke voordelen in scenario's met weinig data en inferentiesnelheid.