Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data

• URL: http://arxiv.org/abs/2506.21788v1
• Date: Thu, 26 Jun 2025 22:04:05 GMT
• Title: Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data
• Authors: Massimiliano Lupo Pasini, Jong Youl Choi, Pei Zhang, Kshitij Mehta, Rylie Weaver, Ashwin M. Aji, Karl W. Schulz, Jorda Polo, Prasanna Balaprakash,
• Abstract summary: Graph foundation models using graph neural networks promise sustainable, efficient atomistic modeling.<n>To tackle challenges of processing multi-source, multi-fidelity data during pre-training, recent studies employ multi-task learning.<n>We propose a multi-task parallelism method that distributes each head across computing resources with GPU acceleration.
• Score: 4.3387776186428
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: 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.

Related papers

• NeuroTrails: Training with Dynamic Sparse Heads as the Key to Effective Ensembling [35.837527844931266]
  We introduce $textbfNeuroTrails$, a sparse multi-head architecture with dynamically evolving topology.<n>NeuroTrails displays efficacy with convolutional and transformer-based architectures on computer vision and language tasks.
  arXiv  Detail & Related papers  (2025-05-23T13:53:21Z)
• OVERLORD: Ultimate Scaling of DataLoader for Multi-Source Large Foundation Model Training [16.91538022228882]
  Modern frameworks for training large foundation models (LFMs) employ dataloaders in a data-parallel manner.<n>We present Omniload, an industrial-grade distributed data loading architecture for LFMs.
  arXiv  Detail & Related papers  (2025-04-14T03:31:22Z)
• A Theoretical Framework for Data Efficient Multi-Source Transfer Learning Based on Cramér-Rao Bound [16.49737340580437]
  We propose a theoretical framework that answers the question: what is the optimal quantity of source samples needed from each source task to jointly train the target model?<n>Specifically, we introduce a generalization error measure that aligns with cross-entropy loss, and minimize it based on the Cram'er-Rao Bound to determine the optimal transfer quantity for each source task.<n>We develop an architecture-agnostic and data-efficient algorithm OTQMS to implement our theoretical results for training deep multi-source transfer learning models.
  arXiv  Detail & Related papers  (2025-02-06T17:32:49Z)
• MAmmoTH-VL: Eliciting Multimodal Reasoning with Instruction Tuning at Scale [66.73529246309033]
  multimodal large language models (MLLMs) have shown significant potential in a broad range of multimodal tasks.<n>Existing instruction-tuning datasets only provide phrase-level answers without any intermediate rationales.<n>We introduce a scalable and cost-effective method to construct a large-scale multimodal instruction-tuning dataset with rich intermediate rationales.
  arXiv  Detail & Related papers  (2024-12-06T18:14:24Z)
• NVLM: Open Frontier-Class Multimodal LLMs [64.00053046838225]
  We introduce NVLM 1.0, a family of frontier-class multimodal large language models (LLMs) that achieve state-of-the-art results on vision-language tasks. We propose a novel architecture that enhances both training efficiency and multimodal reasoning capabilities. We develop production-grade multimodality for the NVLM-1.0 models, enabling them to excel in vision-language tasks.
  arXiv  Detail & Related papers  (2024-09-17T17:59:06Z)
• Training Deep Surrogate Models with Large Scale Online Learning [48.7576911714538]
  Deep learning algorithms have emerged as a viable alternative for obtaining fast solutions for PDEs. Models are usually trained on synthetic data generated by solvers, stored on disk and read back for training. It proposes an open source online training framework for deep surrogate models.
  arXiv  Detail & Related papers  (2023-06-28T12:02:27Z)
• Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning [101.66860222415512]
  Multi-Task Diffusion Model (textscMTDiff) is a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis. For generative planning, we find textscMTDiff outperforms state-of-the-art algorithms across 50 tasks on Meta-World and 8 maps on Maze2D.
  arXiv  Detail & Related papers  (2023-05-29T05:20:38Z)
• Decentralized Training of Foundation Models in Heterogeneous Environments [77.47261769795992]
  Training foundation models, such as GPT-3 and PaLM, can be extremely expensive. We present the first study of training large foundation models with model parallelism in a decentralized regime over a heterogeneous network.
  arXiv  Detail & Related papers  (2022-06-02T20:19:51Z)
• Consistency Training of Multi-exit Architectures for Sensor Data [0.07614628596146598]
  We present a novel and architecture-agnostic approach for robust training of multi-exit architectures termed consistent exit training. We leverage weak supervision to align model output with consistency training and jointly optimize dual-losses in a multi-task learning fashion over the exits in a network.
  arXiv  Detail & Related papers  (2021-09-27T17:11:25Z)
• Multi-modal AsynDGAN: Learn From Distributed Medical Image Data without Sharing Private Information [55.866673486753115]
  We propose an extendable and elastic learning framework to preserve privacy and security. The proposed framework is named distributed Asynchronized Discriminator Generative Adrial Networks (AsynDGAN)
  arXiv  Detail & Related papers  (2020-12-15T20:41:24Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.