Metatensor and metatomic: foundational libraries for interoperable atomistic machine learning

Dit artikel introduceert metatensor en metatomic, twee fundamentele softwarebibliotheken die de interoperabiliteit tussen traditionele atomaire simulatiehulpmiddelen en moderne machine learning-frameworks mogelijk maken door een gemeenschappelijke opslag- en uitwisselingsstructuur voor data en modellen te bieden.

De kern

Stel je voor dat de wereld van het simuleren van atomen (de kleinste bouwstenen van onze wereld) bestaat uit twee grote, gescheiden steden die al eeuwenlang niet met elkaar praten.

Stad 1: De Oude Wereld (Traditionele Simulaties)
Dit is een stad gebouwd van stevige, betonnen blokken. De huizen zijn gebouwd in talen zoals Fortran, C en C++. Het zijn oude, zeer snelle en betrouwbare machines die al decennia lang werken. Ze zijn perfect voor het simuleren van hoe atomen bewegen, maar ze zijn wat stijf en moeilijk te verbouwen.

Stad 2: De Nieuwe Wereld (Machine Learning)
Dit is een stad van glazen torens en neonlichten, gebouwd in Python. Hier werken slimme algoritmen die leren van data. Ze zijn ongelooflijk flexibel en kunnen patronen zien die de oude machines niet eens dromen. Maar ze spreken een andere taal en hebben een heel andere manier van werken.

Het Probleem:
Tot nu toe konden deze twee steden niet goed samenwerken. Als je een slimme nieuwe machine (Stad 2) wilde gebruiken in een oude simulator (Stad 1), moest je elke keer een nieuwe, ingewikkelde brug bouwen. Dat kostte tijd, geld en veel geduld. Vaak lukte het gewoon niet, waardoor wetenschappers vastzaten in hun eigen stad.

De Oplossing: Metatensor en Metatomic
De auteurs van dit paper hebben twee nieuwe "universale vertalers" en "standaardverpakkingen" bedacht om deze twee steden te verbinden.

1. Metatensor: De Slimme Verpakking

Stel je voor dat je een pakketje moet versturen.

• Vroeger: Je gooide de inhoud (de data) in een doos en hoopte dat de ontvanger wist wat erin zat. Als je ook nog de instructies nodig had (bijvoorbeeld: "dit is de kracht op atoom X"), moest je die op een los vel papier schrijven. Als de ontvanger die losse pagina verloor, was de data waardeloos.
• Met Metatensor: Dit is een slimme, zelfverklarende verpakking.
  • Het is een doos die niet alleen de data bevat, maar ook een uitgebreid etiket. Op dat etiket staat precies: "Dit is de energie van atoom X", "Dit is de kracht", en zelfs "Dit is hoe de kracht verandert als je het atoom een beetje verschuift".
  • Het is zo slim dat het alleen de belangrijke dingen in de doos stopt en de lege ruimte (die bij atoom-data vaak heel groot is) weglaat. Dit maakt het super snel en licht.
  • Het resultaat: Of je nu Python, C of Fortran gebruikt, iedereen kan deze doos openen en precies begrijpen wat erin zit. Het is alsof iedereen plotseling dezelfde taal spreekt over data.

2. Metatomic: De Universele Stekker

Stel je voor dat je een nieuw apparaat (een AI-model) hebt dat je wilt aansluiten op een stopcontact (een simulatieprogramma).

• Vroeger: Elk apparaat had een eigen stekker. Je moest voor elk nieuw model een nieuw stopcontact bouwen. Dat was een nachtmerrie.
• Met Metatomic: Dit is een universele stekker.
  • Het zorgt ervoor dat elk AI-model, ongeacht hoe complex of nieuw het is, precies dezelfde vorm heeft aan de buitenkant.
  • Het zegt tegen de simulator: "Ik heb atoomposities nodig, en ik geef je energie en krachten terug."
  • Het resultaat: Je kunt nu elk nieuw AI-model (of het nu gemaakt is door een onderzoeker in Zwitserland of in de VS) direct in elk bestaand simulatieprogramma steken. Geen gedoe meer met speciale kabels.

Waarom is dit geweldig? (De Analoge Wereld)

Dankzij deze twee uitvindingen gebeurt er nu magie:

• Snelheid en Flexibiliteit: Wetenschappers kunnen nu de snelheid van de oude, betonnen machines combineren met de slimme voorspellingen van de nieuwe AI.
• Voorbeelden uit het paper:
  • PET-MAD: Een "universele atoom-voorspeller" die voor bijna elk materiaal werkt. Dankzij deze nieuwe tools kan deze nu in elke simulator worden gebruikt, van simpele tests tot super-complexe experimenten.
  • FlashMD: Een methode die de toekomst van atomen direct voorspelt in plaats van stap voor stap te rekenen. Dit versnelt simulaties met wel 100 keer!
  • NMR-spectroscopie: Chemici kunnen nu heel snel voorspellen hoe moleculen eruitzien in een laboratorium, wat normaal dagen duurt, maar nu in seconden gaat.

Conclusie

Kortom: Metatensor zorgt ervoor dat de data (de informatie) altijd duidelijk en compleet is, en Metatomic zorgt ervoor dat de modellen (de slimme algoritmen) overal kunnen worden aangesloten.

Het is alsof ze een nieuwe internationale taal hebben bedacht voor atomen. Hierdoor kunnen onderzoekers zich niet meer bezighouden met het bouwen van bruggen en kabels, maar zich volledig focussen op het ontdekken van nieuwe materialen, medicijnen en energieoplossingen. De barrières zijn weg, en de samenwerking tussen de oude en nieuwe wereld van de wetenschap is nu een feit.

Technische samenvatting

Probleemstelling

De integratie van machine learning (ML) in atomaire simulaties heeft de nauwkeurigheid en efficiëntie van materialenwetenschap en chemie aanzienlijk verbeterd. Echter, de huidige software-ecosystemen zijn sterk gefragmenteerd:

• Verschillende talen en ecosystemen: Traditionele simulatie-engines (zoals LAMMPS, GROMACS, Quantum ESPRESSO) zijn vaak geschreven in Fortran, C of C++, terwijl moderne ML-frameworks (PyTorch, JAX, scikit-learn) voornamelijk in Python draaien.
• Interoperabiliteitsbarrières: Het koppelen van ML-modellen aan simulatie-engines vereist vaak maatwerk interfaces. Dit leidt tot een "O(M × N)" probleem waarbij elke combinatie van model en engine een unieke interface nodig heeft, wat onderhoud kostbaar en tijdrovend maakt.
• Data-uitwisseling: Bestaande dataformaten ondersteunen vaak niet de specifieke behoeften van atomaire ML, zoals het opslaan van gradiënten (krachten, virial) naast de data zelf, of het efficiënt omgaan met sparsiteit (lege waarden) en metadata over atoomtypes en systemen.
• Modelverspreiding: ML-modellen bestaan niet alleen uit gewichten, maar ook uit code. Het delen en uitvoeren van deze modellen in verschillende omgevingen is complex zonder een gestandaardiseerd formaat.

Methodologie

De auteurs introduceren twee kernbibliotheken, metatensor en metatomic, die als fundamentele lagen dienen om deze kloof te overbruggen.

1. metatensor: Data-opslag en -manipulatie

• Doel: Een zelfbeschrijvend, sparselijk opgeslagen array-formaat voor atomaire data.
• Kernobjecten:
  • Labels: Metadata die multidimensionale indexen definiëren (bijv. atoomindex, systeem, eigenschappen).
  • TensorBlock: Bevat de daadwerkelijke data (dichte float-array) en de bijbehorende metadata (samples, components, properties). Het ondersteunt ook gradiënten (afgeleiden) als geneste TensorBlocks.
  • TensorMap: Een sleutel-waarde map die meerdere TensorBlocks groepeert. Dit maakt een "block-sparse" opslag mogelijk, ideaal voor data die per atoomtype of symmetrie verschilt.
• Implementatie: Geschreven in Rust met een C-API, wat bindings mogelijk maakt voor Python, C++, Rust en TorchScript.
• Serialisatie: Gebruikt het npz-formaat (NumPy) voor taal-onafhankelijke opslag van data en metadata.
• Extensies:
  • metatensor-operations: Biedt wiskundige bewerkingen (optellen, vermenigvuldigen, reduceren) die automatisch metadata en gradiënten bijhouden.
  • metatensor-learn: Biedt hoog-niveau abstrcties voor het bouwen en trainen van ML-modellen (vergelijkbaar met PyTorch nn), met ondersteuning voor symmetriebehoudende (equivariante) netwerken.

2. metatomic: Model-uitwisseling en Interface

• Doel: Een gestandaardiseerde interface tussen ML-modellen en simulatie-engines.
• Werking:
  • Definieert een universeel protocol waarbij een model zijn capaciteiten (output: energie, krachten, etc.) en vereisten (input: atoomposities, buurlijsten) declareert.
  • Vermindert de integratiecomplexiteit van O(M × N) naar O(M + N). Een model dat de interface ondersteunt, werkt met elke engine die deze ook ondersteunt.
  • Modellen worden verpakt in een AtomisticModel object dat de code (vaak gecompileerd via TorchScript), gewichten en metadata bevat.
• Performance: Tests tonen aan dat de overhead van de interface verwaarloosbaar is (~2 µs/atoom) vergeleken met de modeluitvoeringstijd.

Kernbijdragen en Ecosysteem

De auteurs presenteren een volledig modulair ecosysteem gebouwd rond deze bibliotheken:

1. metatrain: Een command-line tool voor het trainen en evalueren van ML-modellen. Het ondersteunt diverse architecturen (van GAP tot Graph Neural Networks) en exporteert modellen direct naar het metatomic-formaat.
2. featomic: Een bibliotheek voor het berekenen van atomaire beschrijvers (descriptors) zoals SOAP en ACE. Geschreven in Rust, biedt het hoge prestaties en lage geheugengebruik, vooral bij het berekenen van gradiënten.
3. torch-spex: Een PyTorch-gebaseerde bibliotheek voor het berekenen van sferische expansies, ondersteunend voor GPU-versnelling en automatische differentiatie.
4. torch-pme: Een bibliotheek voor efficiënte berekening van lange-afstand interacties (Coulomb, dispersion) met ondersteuning voor TorchScript.
5. vesin: Een snelle, compacte bibliotheek voor het genereren van buurlijsten (neighbor lists) in C++.
6. sphericart: Een bibliotheek voor snelle evaluatie van sferische harmonischen.

Resultaten en Gevallenstudies

Het artikel demonstreert de effectiviteit van het ecosysteem aan de hand van drie modellen en diverse integraties:

• PET-MAD: Een universeel interatomair potentieel getraind met metatrain op een diverse dataset (MAD). Het toont aan dat modellen kunnen worden getraind, geëvalueerd en geëxporteerd naar simulatie-engines zoals LAMMPS en ASE. Prestatiemetingen tonen een significante versnelling op GPUs (LAMMPS met KOKKOS) ten opzichte van CPU-only implementaties.
• ShiftML: Een model voor het voorspellen van NMR-chemische verschuivingen. Het toont de interoperabiliteit door het gebruik van metatomic om een model te koppelen aan i-PI voor Path Integral Molecular Dynamics (PIMD) simulaties, waarbij kwantumnucleaire effecten worden meegenomen.
• FlashMD: Een methode die ML gebruikt om direct MD-trajecten te voorspellen (bypassend de berekening van krachten via gradiënten), wat een versnelling van 1-2 orde van grootte mogelijk maakt.
• Integraties: Het ecosysteem is succesvol geïntegreerd met LAMMPS, i-PI, eOn (voor saddle point zoektochten), TorchSim, ASE, PLUMED (voor collectieve variabelen) en chemiscope (voor data-exploratie).

Significantie

De paper markeert een paradigmaverschuiving in atomaire simulaties:

1. Interoperabiliteit: Het lost het probleem van softwarefragmentatie op door een gemeenschappelijke taal (data en modelinterface) te bieden tussen traditionele HPC-code en moderne ML-frameworks.
2. Herbruikbaarheid: Onderzoekers hoeven niet langer specifieke interfaces te schrijven voor elke nieuwe combinatie van model en engine.
3. FAIR-principes: De software is open-source, goed gedocumenteerd, en gebruikt open standaarden, wat de reproduceerbaarheid en toegankelijkheid van wetenschappelijk werk verhoogt.
4. Toekomstperspectief: De auteurs plannen uitbreidingen naar andere talen (Julia, Fortran) en het decoupleren van metatomic van TorchScript om ook pure Python- en C++-modellen te ondersteunen.

Kortom, metatensor en metatomic vormen de ruggengraat voor een nieuw, samenhangend ecosysteem dat de adoptie van machine learning in de computationele chemie en materiaalkunde versnelt en vereenvoudigt.