mlsanityData to 7 Oct 2026

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

arXiv 2604.15380 ↗First posted 15 Apr 2026cond-mat.mtrl-sci cs.AIcited by 0

Not on the map. Why?

Abstract (arXiv)

We present an exascale workflow for materials discovery using atomistic graph foundation models built on HydraGNN. We jointly train on 16 open first-principles datasets (544+ million structures covering 85+ elements) using a multi-task architecture with per-dataset heads and a scalable ADIOS2/DDStore data pipeline. On Frontier, we execute six large-scale DeepHyper hyperparameter optimization campaigns in FP64 and promote the top-performing message-passing models to sustained 2,048-node training, yielding a PaiNN-based lead model. The resulting model enables billion-scale screening, evaluating 1.1 billion atomistic structures in 50 seconds, compressing a workload that would require years of first-principles computation, and supports data-scarce fine-tuning across diverse downstream tasks. We quantify precision-performance tradeoffs (BF16/FP32/FP64), demonstrate transfer across twelve chemically diverse downstream tasks, and establish seamless strong- and weak-scaling across Frontier, Aurora, and Perlmutter. This work allows fast and reliable exploration of vast chemical design spaces that are otherwise inaccessible to first-principles methods.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 23 of its 55 references that are in the map. First posted on arXiv 15 Apr 2026.

Builds on

23 of 55 references.

References · 23

  1. 2017Mixed Precision Trainingcited by 1,135
  2. 2021E(n) Equivariant Graph Neural Networkscited by 951
  3. 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 641
  4. 2017SchNet: A continuous-filter convolutional neural network for modeling quantum interactionscited by 569
  5. 2020Directional Message Passing for Molecular Graphscited by 556
  6. 2020Principal Neighbourhood Aggregation for Graph Netscited by 477
  7. 2022MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fieldscited by 467
  8. 2021Equivariant message passing for the prediction of tensorial properties and molecular spectracited by 438
  9. 2021GemNet: Universal Directional Graph Neural Networks for Moleculescited by 342
  10. 2022Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphscited by 258
  11. 2022Learning Local Equivariant Representations for Large-Scale Atomistic Dynamicscited by 254
  12. 2020The Open Catalyst 2020 (OC20) Dataset and Community Challengescited by 246
11 more
  1. 2023EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representationscited by 201
  2. 2025UMA: A Family of Universal Models for Atomscited by 99
  3. 2017DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamicscited by 79
  4. 2022GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasetscited by 61
  5. 2024Open Materials 2024 (OMat24) Inorganic Materials Dataset and Modelscited by 50
  6. 2023From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Predictioncited by 39
  7. 2022Transition1x -- a Dataset for Building Generalizable Reactive Machine Learning Potentialscited by 37
  8. 2023Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasetscited by 37
  9. 2023The Open DAC 2023 Dataset and Challenges for Sorbent Discovery in Direct Air Capturecited by 18
  10. 2023CHGNet: Pretrained universal neural network potential for charge-informed atomistic modelingcited by 13
  11. 2024Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNNcited by 8

Built on by

Not cited yet.