Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

Deze studie presenteert een exascale-werkstroom voor materiaalsoverkapping die, door middel van een op HydraGNN gebaseerd multi-task grafisch fundamenteel model getraind op meer dan 544 miljoen atoomstructuren, in staat is om 1,1 miljard structuren in 50 seconden te screenen en zo jaren aan eerste-principe-berekeningen te vervangen.

De kern

Stel je voor dat je een enorme bibliotheek hebt met recepten voor het bouwen van nieuwe materialen, zoals batterijen, medicijnen of supersterk staal. In het verleden was het vinden van het perfecte recept als het zoeken naar een naald in een hooiberg, waarbij je elk stukje hooi één voor één moest controleren met een vergrootglas. Dit proces, genaamd "eerste-principes-berekening", was zo traag dat het jaren zou duren om zelfs maar een klein deel van de bibliotheek te doorzoeken.

Deze paper beschrijft een revolutionaire nieuwe manier om die bibliotheek te doorzoeken: een supersnelle, intelligente robot die in 50 seconden 1,1 miljard recepten kan lezen en begrijpen.

Hier is hoe ze dat hebben gedaan, vertaald naar alledaagse taal:

1. De "Super-Leraar" (Het Model)

Stel je voor dat je een student wilt opleiden om chemie te begrijpen. Normaal gesproken geef je die student één boek (één dataset) en laat je hem dat uit het hoofd leren. Maar in de echte wereld zijn er duizenden boeken met verschillende soorten informatie: sommige zijn heel nauwkeurig maar klein, andere zijn grof maar enorm groot.

De onderzoekers hebben een super-leraar (een AI-model genaamd HydraGNN) getraind op 16 verschillende boeken tegelijk. Dit zijn in totaal meer dan 544 miljoen moleculaire structuren.

• Het probleem: Als je een student alles door elkaar laat lezen, raakt hij in de war. De grote boeken domineren en de kleine, belangrijke details verdwijnen.
• De oplossing: Ze hebben een slimme methode gebruikt (Multi-Task Learning). Het is alsof de student één grote, gedeelde kennisbank heeft voor de basisregels van de natuurkunde, maar voor elk specifiek boek een eigen, gespecialiseerde assistent heeft. Zo kan hij leren van de grote boeken zonder de kleine details van de andere te vergeten.

2. De "Supersnelle Bibliotheek" (De Hardware)

Om dit te doen, hadden ze de krachtigste computers ter wereld nodig: de Frontier-supercomputer (een van de snelste ter wereld).

• Ze gebruikten 16.384 grafische kaarten (GPUs) tegelijk. Dat is alsof je 16.000 mensen hebt die elk een pagina van een boek lezen, allemaal perfect op elkaar afgestemd.
• Ze bouwden een slim systeem om de data (de boeken) razendsnel naar die mensen te sturen, zodat niemand hoeft te wachten. Dit is de "ADIOS2/DDStore" pipeline: een soort super-snel transportband die voorkomt dat de leraar vastloopt in de file.

3. De "Proeflezing" (Hyperparameter Optimalisatie)

Voordat ze de definitieve robot bouwden, deden ze een enorme proef. Ze lieten de computer 6 grote zoektochten uitvoeren om te ontdekken welke bouwplaat (architectuur) het beste werkte.

• Het was alsof ze 6 verschillende soorten auto's bouwden en ze allemaal tegelijk een race lieten rijden om te zien welke het snelste en zuinigste was.
• Uiteindelijk won een ontwerp gebaseerd op PaiNN. Dit was de "winnaar": snel, nauwkeurig en niet te groot.

4. Het Grote Resultaat: Van Jaren naar Seconden

Dit is het meest indrukwekkende deel.

• Vroeger: Om 1,1 miljard moleculen te testen, zou een wetenschapper met traditionele methoden 6,7 jaar lang non-stop moeten rekenen op een hele supercomputer.
• Nu: Met hun nieuwe model duurt het 50 seconden.
• De analogie: Het is alsof je vroeger één voor één alle woorden in een bibliotheek moest lezen om een zin te vinden, en nu heb je een magische bril die de hele bibliotheek in één oogopslag scant en je precies vertelt waar de zin staat.

5. Waarom is dit belangrijk?

Dit model is niet alleen snel, het is ook slim en flexibel:

• Finetuning: Stel je hebt een heel specifiek probleem (bijvoorbeeld: "Hoe werkt dit medicijn in de menselijke maag?") en je hebt maar weinig data. Normaal zou een AI hier niet goed in zijn. Maar omdat onze "Super-Leraar" al alles over de basis van de chemie weet, moet hij alleen nog maar een klein beetje bijleren. Dit werkt als een universitair student die al alles over biologie weet en nu snel een specialisatie in geneeskunde kan doen.
• Betrouwbaarheid: Ze hebben getest of het model nog steeds werkt als je het "sneller" maakt door minder nauwkeurige berekeningen te gebruiken (zoals het verschil tussen een schets en een fotorealistische tekening). Ze hebben gevonden dat je voor de allerbelangrijkste taken de "fotorealistische" (hoge precisie) modus moet gebruiken, maar voor het snelle doorzoeken van miljoenen opties werkt de "schets" modus prima.

Samenvatting

Deze paper beschrijft hoe onderzoekers een AI-superheld hebben gecreëerd die in staat is om de chemische wereld te verkennen met een snelheid die voorheen onmogelijk was. Ze hebben een manier gevonden om enorme hoeveelheden verschillende data samen te voegen zonder dat de AI in de war raakt, en ze hebben bewezen dat je in 50 seconden kunt doen wat vroeger jaren kostte.

Dit opent de deur voor het snel vinden van nieuwe materialen voor schone energie, betere medicijnen en duurzame technologie, zonder dat we decennia hoeven te wachten op de resultaten. Het is de overgang van "langzaam en voorzichtig zoeken" naar "razendsnel en slim ontdekken".

Technische samenvatting

1. Het Probleem

De ontdekking van nieuwe materialen wordt traditioneel vertraagd door de hoge rekentijd van eerste-principes methoden, zoals Dichtheidsfunctionaaltheorie (DFT). Deze methoden zijn te duur om op industriële schaal uitgebreide chemische ruimtes te screenen. Hoewel Machine Learning Interatomische Potentiaal (MLIP) modellen en atomaire Graph Foundation Models (GFMs) een oplossing bieden, blijven er aanzienlijke uitdagingen bestaan:

• Data-ongelijkheid en multi-fidelity: Geen enkel dataset of elektronische structuur-instelling is uniform betrouwbaar voor organische, anorganische en hybride systemen. Het trainen op een mengsel van datasets met verschillende nauwkeurigheidsniveaus (fidelity) en schalen leidt tot onbalans in de optimalisatie en label-ruis.
• Schaalbaarheid: Het trainen van modellen op honderden miljoenen structuren vereist exascale rekenkracht, maar bestaande workflows kampen vaak met schaalbaarheidsproblemen bij het verwerken van heterogene data.
• Aanpassing: Veel wetenschappelijke taken hebben te weinig gelabelde data voor training van modellen van nul af aan, wat de noodzaak van effectieve fine-tuning van voorgeprogrammeerde modellen versterkt.

2. Methodologie

Het paper presenteert een geïntegreerde exascale-workflow gebaseerd op HydraGNN, die is ontworpen om deze uitdagingen aan te pakken via de volgende pijlers:

• Multi-Task Learning (MTL) Architectuur:

  • In plaats van één enkel model te trainen op een gemengde dataset, gebruikt de workflow een gedeelde "message-passing" backbone die fundamentele atomaire interacties leert.
  • Er zijn specifieke "output heads" per dataset die de label-kenmerken en fideliteitsverschillen van die specifieke dataset absorberen. Dit vermindert gradient-interferentie tussen datasets en stabiliseert de training op onbalans data.
  • Een Composition-Conditioned Branch Weighting mechanisme (een kleine MLP) leert dynamisch welke dataset-specifieke experts het meest relevant zijn voor een nieuwe chemische samenstelling, waardoor het model kan interpoleren tussen trainingsdomeinen.

• Data-Infrastructuur en Pipeline:

  • Gebruik van ADIOS2 en DDStore (Distributed Data Store) voor efficiënte data-inname en caching van trainingschards dicht bij de compute-nodes. Dit minimaliseert I/O-wachttijden en voorkomt overbelasting van het gedeelde bestandssysteem.
  • De workflow is ontworpen om te werken op 16.384 GPU's (Frontier) en ondersteunt sharded distributed training om het geheugendruk per apparaat te verlagen.

• Hyperparameter-Optimalisatie (HPO) en Modelselectie:

  • Er werden zes grote HPO-campagnes uitgevoerd op de Frontier-supercomputer (1.024 nodes per campagne) om zes verschillende MPNN-architecturen te vergelijken: EGNN, SchNet, DimeNet, MACE, PaiNN en PNAEq.
  • De selectie van het beste model was niet alleen gebaseerd op validatiefout, maar ook op "time-to-epoch" en rekenefficiëntie onder productiebeperkingen.

• Precisie-analyse:

  • Systematische karakterisering van de impact van numerieke precisie (BF16, FP32, FP64) op zowel training als inferentie. Dit is cruciaal voor wetenschappelijke nauwkeurigheid, vooral bij het berekenen van krachten (derivaten van energie).

3. Belangrijkste Bijdragen

• Eerste Exascale Multi-Task Workflow: De eerste demonstratie van het gezamenlijk trainen van een atomaire GFM op 16 open-source datasets, goed voor meer dan 544 miljoen structuren en 85+ elementen.
• Schalbaarheid en Portabiliteit: De workflow toont naadloze sterke en zwakke schaling aan op drie verschillende leidersklasse-supercomputers: Frontier (AMD MI250X), Aurora (Intel XPU) en Perlmutter (NVIDIA A100).
• Innovatieve Inferentie-Optimalisatie: Een reeks optimalisaties voor de inferentie-pipeline, waaronder:
  • Encoder reuse: De gedeelde encoder wordt slechts één keer uitgevoerd in plaats van 16 keer.
  • Branch skipping: Het overslaan van onbelangrijke output-branches.
  • Fused gradient: Het combineren van alle 16 energie-predicties in één gewogen som en het berekenen van krachten in één autograd-call.
  • Torch.compile: Kernel-fusie voor extra snelheid.
• Data-efficiëntie: Bewijs dat fine-tuning van een exascale-voorgeladen model aanzienlijk beter presteert dan training van nul af aan, zelfs in data-schaarse regimes.

4. Resultaten

• Schaalbaarheid:
  • Bereikte bijna lineaire sterke schaling tot 2.048 GPU's op Perlmutter, 6.144 op Aurora en 1.024 op Frontier.
  • De workflow bleef stabiel en efficiënt over verschillende hardware-architecturen.
• Modelprestaties:
  • De PaiNN-architectuur werd geselecteerd als leidende model vanwege de beste balans tussen nauwkeurigheid en rekentijd.
  • Bij fine-tuning op 12 downstream-taken (o.a. QM9, MD17, OQMD) leverde het "unfrozen" fine-tuning van het voorgeprogrammeerde model tot een orde van grootte lagere validatiefouten op vergeleken met training van nul af aan of het bevriezen van de backbone.
• Inferentie-snelheid (Screening):
  • Op de Frontier-supercomputer (9.300 compute nodes) werd het mogelijk om 1,1 miljard atomaire structuren te screenen in slechts 50 seconden.
  • Dit komt neer op een doorvoer van ongeveer 21,8 miljoen structuren per seconde (of 293 structuren/s/GPU).
  • Voor vergelijking: dezelfde taak zou met traditionele DFT-methoden ongeveer 6,7 jaar continue rekentijd vereisen.
• Precisie vs. Snelheid:
  • FP64 biedt de hoogste nauwkeurigheid en is essentieel voor wetenschappelijke toepassing.
  • FP32 introduceert een constante energie-afwijking van ~0,021 eV, terwijl BF16 ~0,080 eV afwijking toont, maar met aanzienlijk hogere doorvoer.

5. Betekenis en Impact

Dit werk markeert een doorbraak in computationele materialenwetenschap door de kloof tussen hoge nauwkeurigheid en hoge doorvoer te overbruggen.

• Versnelling van Ontdekking: Het maakt het mogelijk om chemische ontwerpruimtes te verkennen die voorheen ontoegankelijk waren vanwege de combinatorische complexiteit en de kosten van eerste-principes berekeningen.
• Praktische Toepasbaarheid: Het resulterende model is compact (12,1 miljoen parameters, ~92 MB) en kan direct worden ingezet in productiesimulaties zonder compressie.
• Wetenschappelijke Robuustheid: Door het trainen op FP64 en het analyseren van precisie-effecten, garandeert de workflow dat de resultaten wetenschappelijk betrouwbaar zijn, zelfs bij gebruik van lagere precisie voor inferentie.
• Toekomstperspectief: Het bewijst dat exascale-trained foundation modellen niet alleen "grote modellen" zijn, maar praktische, schaalbare tools die de basis vormen voor toekomstige high-throughput wetenschappelijke workflows.

Samenvattend biedt deze studie een bewezen, reproduceerbare en hoogpresterende workflow die de materialenontdekking transformeert van een beperkte, handmatige zoektocht naar een snelle, datagedreven screening op exascale-niveau.