Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866915361292550144 |
|---|---|
| author | Pasini, Massimiliano Lupo Choi, Jong Youl Zhang, Pei Mehta, Kshitij Weaver, Rylie Aji, Ashwin M. Schulz, Karl W. Polo, Jorda Balaprakash, Prasanna |
| author_facet | Pasini, Massimiliano Lupo Choi, Jong Youl Zhang, Pei Mehta, Kshitij Weaver, Rylie Aji, Ashwin M. Schulz, Karl W. Polo, Jorda Balaprakash, Prasanna |
| contents | Graph foundation models using graph neural networks promise sustainable, efficient atomistic modeling. To tackle challenges of processing multi-source, multi-fidelity data during pre-training, recent studies employ multi-task learning, in which shared message passing layers initially process input atomistic structures regardless of source, then route them to multiple decoding heads that predict data-specific outputs. This approach stabilizes pre-training and enhances a model's transferability to unexplored chemical regions. Preliminary results on approximately four million structures are encouraging, yet questions remain about generalizability to larger, more diverse datasets and scalability on supercomputers. We propose a multi-task parallelism method that distributes each head across computing resources with GPU acceleration. Implemented in the open-source HydraGNN architecture, our method was trained on over 24 million structures from five datasets and tested on the Perlmutter, Aurora, and Frontier supercomputers, demonstrating efficient scaling on all three highly heterogeneous super-computing architectures. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_21788 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data Pasini, Massimiliano Lupo Choi, Jong Youl Zhang, Pei Mehta, Kshitij Weaver, Rylie Aji, Ashwin M. Schulz, Karl W. Polo, Jorda Balaprakash, Prasanna Machine Learning Materials Science Artificial Intelligence Atomic and Molecular Clusters 68T07, 68T09 I.2; I.2.5; I.2.11 Graph foundation models using graph neural networks promise sustainable, efficient atomistic modeling. To tackle challenges of processing multi-source, multi-fidelity data during pre-training, recent studies employ multi-task learning, in which shared message passing layers initially process input atomistic structures regardless of source, then route them to multiple decoding heads that predict data-specific outputs. This approach stabilizes pre-training and enhances a model's transferability to unexplored chemical regions. Preliminary results on approximately four million structures are encouraging, yet questions remain about generalizability to larger, more diverse datasets and scalability on supercomputers. We propose a multi-task parallelism method that distributes each head across computing resources with GPU acceleration. Implemented in the open-source HydraGNN architecture, our method was trained on over 24 million structures from five datasets and tested on the Perlmutter, Aurora, and Frontier supercomputers, demonstrating efficient scaling on all three highly heterogeneous super-computing architectures. |
| title | Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data |
| topic | Machine Learning Materials Science Artificial Intelligence Atomic and Molecular Clusters 68T07, 68T09 I.2; I.2.5; I.2.11 |
| url | https://arxiv.org/abs/2506.21788 |