Technische Samenvatting: ElemeNet

Probleemstelling

Hoewel deep learning de voorspelling van moleculaire eigenschappen heeft gerevolutioneerd, blijven state-of-the-art (SOTA) modellen grotendeels beperkt tot gespecialiseerde codebases die een gebrek aan generaliseerbaarheid vertonen over diverse chemische ruimtes. Bestaande algemene tools, zoals Chemprop, zijn primair beperkt tot 2D-representaties (SMILES-strings) en organische chemie, en bieden geen natuurlijke ondersteuning voor stereochemie, complexe bindingen in organometaalverbindingen, zware elementen of biologische macromoleculen. Bovendien missen deze tools vaak ondersteuning voor 3D-equivariante architecturen, expliciete conditionering op lading en spin-toestanden, en robuuste onzekerheidskwantificering (UQ). Er is een kritieke behoefte aan een verenigd softwareframework dat ondersteuning biedt voor het volledige periodiek systeem (elementen 1–100), schaalt naar miljoenen moleculen en gekalibreerde onzekerheidsschattingen levert zonder dat daar diepgaande expertise in machine learning voor nodig is.

Methodologie

De auteurs introduceren ElemeNet, een modulaire, algemene softwarepackage voor moleculaire machine learning, ontworpen om deze hiaten te overbruggen.

1. Representaties en Input

• Elementaire reikwijdte: Ondersteunt atoomnummers 1–100 (Waterstof tot Fermium).
• Inputformaten: Accepteert diverse 2D- en 3D-formaten (cif, cml, mol, mol2, sdf, SMILES, xyz).
• Grafenconstructie:
  • 2D: Afgeleid van topologische connectiviteit; coördinaten worden alleen gebruikt voor de interpretatie van bindingen en daarna weggegooid.
  • 3D: Geconstrueerd uit atomaire coördinaten met randen gedefinieerd door een instelbare afkapstraal, waarbij paarsgewijze afstanden worden gebruikt om E(3)-equivariantie op te leggen.
• Feature Engineering:
  • Node Features: Een gebalanceerde initiële representatie bestaande uit continue getabelleerde atomaire eigenschappen (massa, radii, elektronegativiteit, etc.) en categorische kenmerken (groep, periode, valentie-elektronen) via one-hot encoding.
  • Edge Features: Codeert het type binding, de geschatte bindingslengte en de inclusie van metalen.
  • Graph Features: Codeert expliciet moleculaire lading en spin-multipliciteit via een dedicated neuraal netwerk, waardoor modellen eigenschappen kunnen leren als functies van de elektronische toestand.
  • Flexibiliteit: Ondersteunt door de gebruiker opgegeven features (bijv. partiële ladingen, Wiberg-indices) en impliciete/expliciete waterstofafhandeling.

2. Modelarchitecturen

ElemeNet maakt gebruik van een modulaire encoder-readout architectuur:

• Encoders:
  • 2D GNNs: Ondersteunt directed message-passing (D-MPNN), grafconvolutie, graph isomorphism en graph attention networks.
  • 3D Equivariante GNNs (EGNN): Gebruikt aparte message-passing lagen voor invariante scalaire features en 3D-coördinaten, waarbij interatomaire afstanden worden geprojecteerd via Gaussische basisfuncties om E(3)-equivariantie te waarborgen zonder sferische harmonieken.
• Readouts:
  • MLP: Standaard multilayer perceptron.
  • Transformer: Multi-head self-attention mechanismen voor dynamische weging van grafencomponenten.
• Hyperparameter Optimalisatie: Maakt gebruik van Optuna voor Bayesiaanse optimalisatie over een enorme ontwerpplek, inclusiële netwerkvormen (constant, toenemend, afnemend, hourglass, piramide), activatiefuncties en convolutie-operators.

3. Multiscale Leertaken

Het framework ondersteunt voorspellingen op meerdere schalen binnen één enkele graaf:

• Graph-niveau: Moleculaire eigenschappen (bijv. totale energie, toxiciteit).
• Node-niveau: Atoomgecentreerde eigenschappen (bijv. partiële ladingen).
• Edge-niveau: Bindingseigenschappen (bijv. bindingslengtes).
• Moiety-niveau (Nieuw): Een nieuwe capaciteit om eigenschappen te voorspellen die gelokaliseerd zijn aan een specifieke subgraph (k-hop buurt) gecentreerd rond een node, waarbij de rest van het molecuul tijdens de loss-berekening wordt gemaskeerd. Dit richt zich op functionele groepen of actieve sites.

4. Onzekerheidskwantificering (UQ)

ElemeNet implementeert twee niveaus van UQ:

• Heuristische Maten: Latente ruimteafstand (Euclidische en cosinus afstand) en informatie-entropie voor classificatie. Deze zijn goedkoop, maar niet formeel gekalibreerd aan de foutgrootte.
• Ensemble-gebaseerde Maten:
  • Regressie: Gebruikt shallow ensembles waarbij de uiteindelijke readout-laag $N$ keer wordt gerepliceerd terwijl de encoder-trunk gedeeld wordt. Het model minimaliseert de Negative Log-Likelihood (NLL) van een Gaussische distributie geparametriseerd door het ensemble-gemiddelde en de variantie.
  • Classificatie: Breidt shallow ensembles uit met een Mean-of-Probabilities (MoP) schema. Logits worden gemiddeld over de heads na het toepassen van de activatiefunctie (sigmoid/softmax), en de standaarddeviatie van het ensemble in de waarschijnlijkheidsruimte dient als de onzekerheidsmetriek.

Belangrijkste Resultaten

De auteurs hebben ElemeNet getest tegenover literatuur-baselines over vier diverse datasets:

1. QM9 (Organische Chemie):

   • 3D EGNN-modellen presteerden significant beter dan 2D GNN's, met name voor dipoolmomenten (6x reductie in MAE), wat de noodzaak van 3D-informatie voor geometrisch afhankelijke eigenschappen bevestigt.
   • ElemeNet-modellen (2D en 3D) evenaarden of overtroffen de prestaties van Chemprop v1/v2.
   • Shallow ensembling verbeterde de onzekerheidskalibratie (Spearman rangcorrelatie met MAE) zonder de voorspellende nauwkeurigheid significant te compromitteren.

2. tmQMg (Anorganische/Overgangsmetaalchemie):

   • EGNN-modellen presteerden beter dan 2D GNN's op het gebied van HOMO-LUMO gaten, dipolen en polariseerbaarheid.
   • Het beste ElemeNet-model presteerde beter dan het gemiddelde van bestaande 3D-architecturen uit de literatuur, ondanks het gebruik van standaard hyperparameters voor de transformer readout.
   • Een multitask EGNN-model bereikte competitieve prestaties over alle drie de targets simultaan.

3. pydentate (Coördinatiechemie):

   • ElemeNet bereikte een vergelijkbare nauwkeurigheid als gepubliceerde 2D-modellen voor het voorspellen van het coördinerende atoom en het coördinatiegetal.
   • Transformer readouts verbeterden de gebalanceerde nauwkeurigheid voor de voorspelling van het coördinerende atoom (96,1% vs. 92,4% voor MLP).
   • Stelde nieuwe benchmarks vast voor edge-level (bindingslengtes) en subgraph-level (metaal-ligand bindingslengtes) voorspellingen, die voorheen niet gerapporteerd waren.

4. GEMS (Biologische Chemie, 2,7M Moleculen):

   • Toonde schaalbaarheid naar datasets met miljoenen vermeldingen.
   • Modellen getraind met standaard hyperparameters bereikten competitieve MAE's voor elektronische energie en dipoolgrootte.
   • Expliciete lading- en spin-encoding bleek cruciaal, aangezien het foutief labelen van geladen soorten als neutraal de voorspellingsfout aanzienlijk verhoogde.

Betekenis en Claims

Het artikel claimt dat ElemeNet dient als een praktisch, verenigd platform voor het trainen van performante deep learning modellen over het gehele periodiek systeem. De primaire betekenis ligt in:

• Democratisering: Het verlagen van de drempel voor niet-computationele onderzoekers via een beknopte command-line interface en geautomatiseerde hyperparameter optimalisatie.
• Generaliseerbaarheid: Verder gaan dan organische chemie door ondersteuning te bieden voor organometaal-, anorganische en biologische systemen met complexe bindingen en elektronische toestanden.
• Multiscale Capaciteit: Introductie van moiety-level voorspellingen en ondersteuning voor node, edge, graph en subgraph targets binnen één enkel framework.
• Robuustheid: Het bieden van goed gekalibreerde onzekerheidsschattingen via shallow ensembles, een cruciale voorwaarde voor active learning en betrouwbare inzetbaarheid.
• Prestaties: Het behalen van SOTA of competitieve prestaties over diverse chemische domeinen met minimale gebruikersinspanning en bescheiden computationele budgetten.

De auteurs concluderen dat ElemeNet onderzoekers in staat stelt om moderne deep learning methoden te benutten voor moleculaire ontdekking en optimalisatie binnen de chemische en fysische wetenschappen.