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

Massimiliano Lupo Pasini; Jong Youl Choi; Kshitij Mehta; Richard Messerly; Rylie Weaver; Linda Ungerboeck; Isaac Lyngaas; Benajmin Stump

cond-mat.mtrl-sciApr 15, 2026

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

Authors: Massimiliano Lupo Pasini, Jong Youl Choi, Kshitij Mehta, Richard Messerly, Rylie Weaver, Linda Ungerboeck, Isaac Lyngaas, Benajmin Stump,

Organizations: Computational Sciences and Engineering Division, Oak Ridge National Laboratory · Computer Science and Mathematics Division, Oak Ridge National Laboratory · National Center for Computational Sciences Division, Oak Ridge National Laboratory · Bredesen Center Data Science and Engineering, University of Tennessee, Knoxville ·

Abstract

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.

Figures & tables

Attribute	Value
Category of achievement	scalability / time-to-solution / peak performance
Type of method used	n/a
Results reported on the basis of	whole application including I/O
Precision reported	mixed precision
System scale	results measured on full-scale system
Measurement mechanism	timers (GPTL); other (Omnistat telemetry)

TABLE I: Performance Attributes Required by ACM Gordon Bell Nomination

Dataset	Type of compounds	Approximation theory	Samples	Files
QM7-X [ 8 ]	Organic molecules	DFT (PBE0+MBD)	8,390,474	8
QCML-DFT [ 9 ]	Organic molecules	DFT (PBE0)	33,496,167	10
Transition1x [ 10 ]	Organic molecules	DFT ( $\omega$ b97x/6-31G(d))	9,644,740	1
ANI-1x [ 11 ]	Organic molecules	DFT ( $\omega$ b97x/6-31G(d))	4,956,005	1
$\nabla^{2}$ -DFT [ 12 ]	Organic molecules	DFT ( $\omega$ b97x-D/def2-SVP)	10,392,954	3
Materials Project Trajectories (MPTrj) [ 13 ]	3D bulk alloys	DFT (PBE0)	1,580,395	1

TABLE II: Multi-source, multi-fidelity first-principles datasets used in this work.

Fig. 1: Multi-task learning training for MLIPs on multi-source, multi-fidelity data, with shared message-passing representations and dataset-specific supervision heads. Message passing layers learn fundamental features that describe basic physics of interatomic interactions common to all datasets. Output heads customize features to the physics captured by the approximation theory used to generate a specific dataset

Hyperparameter	Range
Convolution layers	2–6
Message-passing hidden dimension	EGNN/SchNet: 100–3000; DimeNet: 10–100; MACE/PaiNN/PNAEq: 100–1000
Prediction-head layers	2–4
Prediction-head hidden dimension	300–2000
Force weight	$10$ – $100$
Learning rate	$10^{-5}$ – $3\times 10^{-3}$

TABLE III: Hyperparameters tuned in all single-model HPO campaigns. Convolution depth and hidden dimensions were sampled as integers, force weight and learning rate were sampled as a continuous real-valued variable.

MPNN	Additional tuned hyperparameters
EGNN	None beyond the shared search space.
SchNet	Interaction filters: 6–300; Gaussian basis functions: 3–100.
DimeNet	Basis embedding size: 8–200; envelope exponent: 4–6; interaction embedding size: 16–500; output embedding size: 24–4000; layers after skip: 0–2; layers before skip: 0–2; radial basis count: 3–12; spherical basis count: 2–8.
MACE	Radial basis type: {Bessel, Gaussian, Chebyshev}; distance transform: {Agnesi, Soft}; radial basis count: 3–12; max $\ell$ : 2–3; node max $\ell$ : 1–3; correlation order: 2–3.
PaiNN	Radial basis count: 3–12.
PNAEq	Radial basis count: 3–12.

TABLE IV: Additional HPO dimensions exposed for each MPNN family in the local SC26 HPO driver.

Fig. 2: Validation-loss trajectories for all trials from the six Frontier HPO campaigns, with one campaign per MPNN backbone. It shows 181 successful DeepHyper trials out of 574 explored configurations across MPNN backbones.

Parameter	Value
Convolution layers	2
Hidden dimension	337
Interaction filters	126
Radial basis functions	4 (Bessel)
Cutoff radius	5 Å
Max neighbors	20

TABLE V: Lead PaiNN model configuration selected from the HPO campaign and used for all downstream analysis.

Fig. 3: Strong and weak scaling results, along with time decomposition, are presented for Perlmutter, Aurora, and Frontier. Speedup (y-axis) is normalized to the 32-node performance on each system, with the number of GPUs shown on the x-axis. The decomposed timing corresponds to the weak scaling results, detailing time spent on data loading, forward pass, backward pass, and optimization.

Dataset	Size	Domain	Target(s)
QM9 [ 52 ]	${\sim}$ 130k	Small organic molecules (up to 9 heavy atoms; C, H, O, N, F)	Atomization energy
MD17 [ 53 ]	${\sim}$ 134k	Uracil ab-initio MD trajectory	Total energy and forces
Wiggle150 [ 54 ]	150	Strained conformers of adenosine, benzylpenicillin, efavirenz	Relative conformer energy
MS25 - MgO-2 $\times$ 2 [ 55 ]	10k	MgO(100) oxide surface	Per-atom energy
MS25 - MgO-4 $\times$ 4 [ 55 ]	10k	MgO(100) oxide surface	Per-atom energy
MS25 -Reaction [ 55 ]	10k	CH 4 dissociation on Pt(111)	Per-atom energy

TABLE VI: Downstream fine-tuning task suite. All tasks predict graph-level (per-structure) properties. Size indicates the number of structures available for fine-tuning.

Dataset	Split	Epochs	Batch Size	Learning Rate
QM9	70/15/15	100	32	$10^{-3}$
MD17	70/15/15	100	32	$10^{-4}$
Wiggle150	70/15/15	500	32	$10^{-4}$ a
MS25	70/15/15	10	2	$10^{-4}$
OQMD	70/15/15	100	32	$10^{-4}$
ABX 3	70/15/15	100	32	$10^{-4}$

TABLE VII: Training configuration for each benchmark. All runs use AdamW and FP64 unless otherwise noted.

Method	MAE (eV/atom)
Training from scratch	0.047
Frozen backbone	0.050
Unfrozen backbone	0.004

TABLE VIII: QM9 validation MAE (eV/atom) after 100 epochs for atomization energy at 0 K (

U_{0}

, property index 12).

	Energy MAE	Forces MAE
Method	kcal/mol	kcal/(mol Å)
Training from scratch	3.421	8.438
Frozen backbone	4.273	13.663
Unfrozen backbone	1.481	6.543

TABLE IX: MD17 uracil MLIP benchmark: best validation over 100 epochs.

Method	MAE (kcal/mol)	RMSE (kcal/mol)
Training from scratch	19.914	29.203
Frozen backbone	24.715	34.212
Unfrozen backbone	18.732	23.953

TABLE X: Validation performance on Wiggle150 relative conformer energies (kcal/mol).

System	Scratch	Frozen Backbone	Unfrozen Backbone
MgO-2 $\times$ 2	0.030	0.030	0.017
MgO-4 $\times$ 4	0.042	0.046	0.003
HEA	0.032	0.032	0.031
Reaction	0.047	0.046	0.019
Zr-O	0.293	0.293	0.152

TABLE XI: MS25 Validation MAE (eV/atom) after 10 epochs. Energy-only supervision; no force labels used.

Method	OQMD	ABX 3
FP64 training from scratch	0.596	0.444
BF16 unfrozen backbone	0.684	0.352
FP32 unfrozen backbone	0.245	0.341
FP64 unfrozen backbone	0.227	0.310

TABLE XII: Validation MAE (eV/atom) after 100 epochs for OQMD and ABX 3 .

Fig. 4: Validation-loss trajectories for all fine-tuning trials on the OQMD dataset. A smoothing window of 5 epochs was used.

Method	Avg. MAE	Std. MAE	Avg. RMSE	Max. Error
Scratch	47.763	15.215	116.756	1,561.365
Frozen	50.714	15.835	116.951	1,525.047
Unfrozen	48.346	14.418	112.694	1,536.159

TABLE XIII: MB-jdft2d (Exfoliation Energy in meV/Atom) Five-Fold Cross Validation Statistics after 100 epochs of fine-tuning.

	ROCAUC		F1	Accuracy
Method	Avg.	Std.	Avg.	Avg.
Training from scratch	0.911	0.002	0.809	0.835
Frozen backbone	0.924	0.003	0.823	0.847
Unfrozen backbone	0.924	0.003	0.825	0.848

TABLE XIV: MB-is_metal Five-Fold Cross Validation Statistics after 10 epochs of fine-tuning.

Fig. 5: Fused Gradient Inference Pipeline. The encoder executes once and its activations are cached. All 16 decoder heads run forward-only, their energies are combined into a weighted sum, and a single backward pass computes forces.

	FP64		FP32
Optimization	MI250	MI350	MI250	MI350
Baseline	23	106	42	190
$+$ Encoder reuse	35	169	66	313
$+$ Branch skip	165	743	308	1250
$+$ Fused gradient	306	1428	602	2134
$+$ torch.compile	738	3012	1405	6050

TABLE XV: Per-GPU Throughput (structures/s).

Graph Foundation Models

Graph Foundation Models are pretrained graph-learning systems designed to transfer across datasets, tasks, and domains with different structures and features.

See all 96 papers

MMAI Gym for Science: Training Liquid Foundation Models for Drug Discovery
Mar 3, 2026Maksim Kuznetsov, Zulfat Miftahutdinov, Rim Shayakhmetov +17Molecular Optimization Molecular Property Prediction
Insilico Medicine · Liquid AI
Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials
Apr 17, 2026Yuanchang Zhou, Hongyu Wang, Yiming Du +12Machine Learning Interatomic Potentials Communication-Efficient Distributed Training
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · Independent Researcher · Sun Yat-Sen University
SCALE: Simulation-Calibrated Amortized Learning for Energy Materials (A hybrid architecture connecting deterministic modeling, real-world data, and transformer-scale inference for accelerated energy-materials discovery)
Sep 3, 2026Kuan Huang, Bo BaiMaterials Discovery Materials Property Prediction
Alpha Ladder Particle Pte. Ltd., 168 Robinson Road, #26-02, Singapore 068912

Mar 3, 2026J/K move · Enter open · S save

MMAI Gym for Science: Training Liquid Foundation Models for Drug Discovery

Maksim Kuznetsov, Zulfat Miftahutdinov, Rim Shayakhmetov, Mikolaj Mizera+16

Insilico Medicine · Liquid AI

Molecular Optimization Molecular Property Prediction Molecular Foundation Models Retrosynthesis Molecular Representation Learning Drug Discovery

Open paper AskPDF arXiv

General-purpose large language models (LLMs) that rely on in-context learning do not reliably deliver the scientific understanding and performance required for drug discovery tasks. Simply increasing model size or introducing reasoning tokens does not yield significant performance gains. To address this gap, we introduce the MMAI Gym for Science, a one-stop shop molecular data formats and modalities as well as task-specific reasoning, training, and benchmarking recipes designed to teach foundation models the 'language of molecules' in order to solve practical drug discovery problems. We use MMAI Gym to train an efficient Liquid Foundation Model (LFM) for these applications, demonstrating that smaller, purpose-trained foundation models can outperform substantially larger general-purpose or specialist models on molecular benchmarks. Across essential drug discovery tasks - including molecular optimization, ADMET property prediction, retrosynthesis, drug-target activity prediction, and functional group reasoning - the resulting model achieves near specialist-level performance and, in the majority of settings, surpasses larger models, while remaining more efficient and broadly applicable in the domain.

Abstract

Graph Foundation Models

Explore similar work

MMAI Gym for Science: Training Liquid Foundation Models for Drug Discovery

Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials

SCALE: Simulation-Calibrated Amortized Learning for Energy Materials (A hybrid architecture connecting deterministic modeling, real-world data, and transformer-scale inference for accelerated energy-materials discovery)