All-in-one foundational models learning across quantum chemical levels

Deze paper introduceert het all-in-one ANI-model, een multimodale architectuur die meerdere niveaus van kwantumchemie kan leren en zo een schaalbare, universele en nauwkeurige oplossing biedt voor het voorspellen van moleculaire eigenschappen.

De kern

Stel je voor dat je een superintelligente kok wilt hebben die voor elke maaltijd de perfecte smaak kan voorspellen. In de wereld van de chemie is die "smaak" de energie van een molecuul, en de "recepten" zijn de wiskundige formules (de kwantumchemische niveaus) die we gebruiken om die energie te berekenen.

Tot nu toe hadden we een groot probleem: elke kok was gespecialiseerd in één recept.

• Wil je een snelle, simpele maaltijd? Dan gebruik je een snelle, maar minder nauwkeurige kok (zoals semi-empirische methoden).
• Wil je een Michelin-ster maaltijd? Dan moet je een dure, trage kok inhuren die urenlang rekent (zoals DFT of Coupled Cluster).

Je kon niet gewoon één kok hebben die beide kon doen. Als je een snelle kok wilde, moest je die apart trainen. Als je een perfecte kok wilde, weer een andere. Dat kostte tijd, geld en veel computerkracht.

De doorbraak: De "Alles-in-één" (AIO) Chef

In dit nieuwe onderzoek hebben Yuxinxin Chen en Pavlo Dral een revolutionaire nieuwe chef-kok bedacht: de AIO-ANI.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De Super-Kok die alles kan

Stel je een neural netwerk (een soort computerhersenen) voor als een kok. Normaal gesproken leer je die kok één ding: "Hoe smaakt dit als ik het kook volgens recept A?"

Deze nieuwe AIO-kok is anders. Je geeft hem niet alleen de ingrediënten (de vorm van het molecuul), maar je zegt hem ook: "Kook dit volgens recept A, of recept B, of recept C."

• Recept A is een snelle, simpele methode.
• Recept B is een geavanceerde, dure methode.
• Recept C is de ultieme, perfecte methode.

De kok leert in één keer alle recepten tegelijk. Hij leert de patronen van de snelle methoden én de fijne kneepjes van de dure methoden. Het resultaat? Je hebt nu één enkel model dat in een fractie van een seconde kan voorspellen hoe een molecuul zich gedraagt, of je nu een snelle schatting wilt of een ultra-precieze berekening.

2. Waarom is dit beter dan de oude manier? (De "Transfer Learning" vergelijking)

Vroeger probeerden wetenschappers dit te doen met een methode die ze "Transfer Learning" noemden. Dat was alsof je eerst een kok traint om pizza's te maken (de snelle methode), en hem daarna opnieuw moet trainen om sushi te maken (de dure methode).

• Je hebt dan twee aparte koks nodig.
• Het duurt lang om ze allebei te trainen.
• Het is lastig om te bepalen hoe je de training moet aanpassen.

De AIO-methode is als het trainen van één "Master Chef" die in één les alle gerechten van de hele wereld leert. Het is sneller, simpeler en de chef wordt zelfs iets beter in zijn werk dan de oude methode.

3. De "Delta-Learning" Truc: De Perfecte Combinatie

De auteurs gaan nog een stapje verder. Ze zeggen: "Waarom kiezen we? Laten we het beste van beide werelden combineren."

Ze gebruiken de AIO-kok om een correctie te maken.

• Stel, je gebruikt de snelle, simpele methode (recept A) omdat die snel is.
• Maar je weet dat die niet 100% perfect is.
• De AIO-kok kijkt naar het verschil tussen de snelle methode en de perfecte methode en zegt: "Ah, de snelle methode mist een beetje zout. Voeg dit toe."

Dit noemen ze Δ-learning (Delta-learning). Het is alsof je een goedkoop, snel gerecht neemt en er met een magische lepel een paar druppels van de "sterrenchef-saus" aan toevoegt. Het resultaat is een maaltijd die net zo snel is als het goedkope gerecht, maar net zo lekker als het duurste gerecht.

Wat betekent dit voor de wereld?

• Snelheid: Berekeningen die normaal dagen duren, gaan nu in seconden.
• Toegang: Wetenschappers hoeven geen dure supercomputers meer te huren om nauwkeurige resultaten te krijgen.
• Flexibiliteit: Je kunt nu in één keer schakelen tussen "snelle schatting" en "ultra-precies" zonder nieuwe software te hoeven installeren.

Kortom:
De auteurs hebben een "universele vertaler" gebouwd voor de chemie. In plaats van een woordenboek te hebben voor elke taal apart, hebben ze één boek gemaakt dat alle talen tegelijk vertaalt, en dat bovendien nog sneller en accurater is dan de oude methoden. Dit opent de deur voor snellere medicijnontwikkeling, nieuw materiaalontwerp en een beter begrip van de chemische wereld om ons heen.

De code en de modellen zijn nu beschikbaar voor iedereen om te gebruiken, zodat deze "Master Chef" overal ter wereld aan het werk kan gaan.

Technische samenvatting

Probleemstelling

Machine Learning Interatomische Potentiaal (MLIP) modellen richten zich doorgaans op één enkel niveau van kwantumchemie (QC). Hoewel er methoden zijn ontwikkeld voor multi-fidelity learning (leren op verschillende nauwkeurigheidsniveaus), bieden deze geen schaalbare oplossingen voor fundamentele modellen die op grote datasets kunnen worden getraind. Bestaande benaderingen zoals Transfer Learning (TL) vereisen vaak aparte modellen voor elk niveau of een tweestapsproces (pre-training en fine-tuning), wat leidt tot inefficiëntie. Andere methoden, zoals $\Delta$-learning, zijn beperkt tot specifieke combinaties van een basis- en doelniveau en vereisen evaluatie op het QC-niveau tijdens de inferentie. Er ontbreekt een schaalbare, uitbreidbare modelarchitectuur die het mogelijk maakt om op grote schaal te trainen met een willekeurig aantal QC-methoden (variërend van semi-empirisch tot coupled cluster) in één enkel model.

Methodologie: De AIO-ANI Architectuur

De auteurs introduceren een All-in-One (AIO) modelarchitectuur gebaseerd op multimodaal leren. De kern van deze methode is het behandelen van het referentieniveau van de theorie (level of theory) als een invoerfeature, naast de geometrische informatie.

• Architectuur: Het model is gebaseerd op de ANI-type (Atomistic Neural Network) modificatie van Behler-Parrinello neurale netwerken.
• Features:
  • Geometrie: Gecodeerd via ANI-type atomaire omgevingvectoren (AEV) die gebaseerd zijn op het elementtype.
  • QC-niveau: De theorie (bijv. DFT, CCSD(T)) wordt omgezet in een one-hot encoding en aan de AEV-features toegevoegd.
• Training: Het model wordt getraind op een dataset die ongeveer 4,5 miljoen energieën en krachten bevat op het DFT-niveau ($\omega$B97X/def2-TZVPP) en ongeveer 0,5 miljoen energieën op het high-level CCSD(T)/CBS-niveau. Ook semi-empirische niveaus (GFN2-xTB en ODM2) zijn opgenomen in latere experimenten.
• Inferentie: Eenmaal getraind, kan het model voorspellingen doen voor elk getraind niveau door simpelweg de gewenste theorie als string in te voeren. De totale energie wordt berekend als de som van de atomaire energieën plus zelf-atomaire energieën ($E_{SAE}$) en expliciete dispersiecorrecties ($E_{D4}$).
• $\Delta$-Learning Integratie: Het AIO-model kan ook worden gebruikt om $\Delta$-learning correcties te genereren. Door het verschil te nemen tussen de voorspellingen van het AIO-model op twee niveaus (bijv. CC en DFT), kan dit worden toegevoegd aan een echte DFT-basisberekening, wat resulteert in een robuustere en nauwkeurigere methode.

Belangrijkste Bijdragen

1. Universeel Fundamenteel Model: Ontwikkeling van AIO-ANI-UIP, een enkel fundamenteel model dat voorspellingen kan doen voor verschillende QC-niveaus (van semi-empirisch tot coupled cluster) zonder dat er aparte modellen nodig zijn voor elk niveau.
2. Superioriteit ten opzichte van Transfer Learning: De auteurs tonen aan dat AIO-leren sneller convergeert (1000 epochen vs. 3750 epochen voor TL) en nauwkeuriger is dan traditionele transfer learning, terwijl het tegelijkertijd flexibeler is voor het toevoegen van nieuwe niveaus.
3. Ontwikkeling van $\Delta$-AIO-ANI: Een verbeterd model dat het AIO-gebaseerde correctieconcept combineert met een DFT-basis. Dit resulteert in een methode met aanzienlijk hogere nauwkeurigheid en robuustheid dan zowel het pure ML-model als standaard DFT-methoden.
4. Open Beschikbaarheid: De code en de fundamentele modellen zijn beschikbaar gesteld via GitHub en zullen worden geïntegreerd in de MLatom-pakketten en de UAIQM-bibliotheek (Universal and Updatable AI-enhanced QM), toegankelijk via het XACS-cloudcomputingplatform.

Resultaten

• Nauwkeurigheid: Het AIO-ANI-UIP model bereikt een nauwkeurigheid die vergelijkbaar is met semi-empirische methoden (zoals GFN2-xTB) en DFT (B3LYP-D4/6-31G*), maar met de snelheid van een ML-potentieel. De gewogen gemiddelde absolute afwijking (WTMAD-2) op de GMTKN55-benchmarkset is ongeveer 1,2 kcal/mol voor zowel DFT- als CC-niveaus binnen de trainingsdistributie.
• Prestatie van $\Delta$-AIO-ANI: Het $\Delta$-AIO-ANI model (DFT-basis + AIO-correctie) presteert aanzienlijk beter dan de pure ML-alternatieven. De WTMAD-2 daalt naar 4,69 kcal/mol, wat de helft is van de fout van het standaard AIO-ANI-UIP en significant lager is dan GFN2-xTB of B3LYP/6-31G*.
• Generalisatie: Het model toont een sterke generalisatiecapaciteit over de chemische ruimte van organische moleculen.
• Validatie en Overfitting: De studie identificeerde dat generalisatiefouten onstabiel kunnen zijn en afhankelijk zijn van het aantal trainingsepochs. Door gebruik te maken van een externe validatieset (S30L) voor niet-covalente interacties, werd de training gestabiliseerd en overfitting voorkomen.

Betekenis en Impact

Dit werk markeert een doorbraak in de ontwikkeling van MLIP's door de barrière tussen verschillende kwantumchemische niveaus te doorbreken met één enkel, schaalbaar model.

• Efficiëntie: Het elimineert de noodzaak om meerdere modellen te trainen voor verschillende nauwkeurigheidsniveaus, wat tijd en rekenkracht bespaart.
• Flexibiliteit: Het biedt een eenvoudige manier om $\Delta$-learning correcties te creëren voor willekeurige combinaties van basis- en doelniveaus, wat leidt tot "AI-versterkte" kwantummechanische methoden.
• Toekomstperspectief: De AIO-architectuur legt de basis voor fundamentele modellen die kunnen leren van heterogene data en zich kunnen aanpassen aan nieuwe QC-methoden naarmate deze beschikbaar komen, wat essentieel is voor de volgende generatie computationele chemie.

De auteurs benadrukken dat het gebruik van eenvoudigere architecturen met meer data (zoals hier gedaan) praktischer kan zijn dan het jagen op extreem data-efficiënte, complexe netwerken. De beschikbaarheid van deze modellen via open-source platforms democratiseert de toegang tot high-level kwantumchemische berekeningen.