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

📄 arXiv: 2506.21788v1 📥 PDF

作者: Massimiliano Lupo Pasini, Jong Youl Choi, Pei Zhang, Kshitij Mehta, Rylie Weaver, Ashwin M. Aji, Karl W. Schulz, Jorda Polo, Prasanna Balaprakash

分类: cs.LG, cond-mat.mtrl-sci, cs.AI, physics.atm-clus

发布日期: 2025-06-26

备注: 15 pages, 4 figures, 2 tables

💡 一句话要点

提出多任务并行方法以增强图基础模型的鲁棒性

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 图基础模型 多任务学习 图神经网络 原子建模 GPU加速 超级计算机 数据处理 迁移学习

📋 核心要点

  1. 现有方法在处理多源、多保真数据时存在稳定性不足和迁移能力有限的问题。
  2. 提出的多任务并行方法通过GPU加速将解码头分布到计算资源上，提高了预训练的效率和稳定性。
  3. 在超过2400万个结构的训练中，该方法在多种超级计算机上展示了良好的扩展性和性能提升。

📝 摘要（中文）

图基础模型利用图神经网络在原子建模中展现出可持续和高效的潜力。为了解决在预训练过程中处理多源、多保真数据的挑战，近期研究采用了多任务学习的方法，通过共享的消息传递层处理输入的原子结构，然后将其路由到多个解码头以预测特定数据的输出。这种方法稳定了预训练过程，并增强了模型在未探索化学区域的迁移能力。尽管在约四百万个结构上的初步结果令人鼓舞，但对于更大、更具多样性数据集的泛化能力和在超级计算机上的可扩展性仍存在疑问。本文提出了一种多任务并行方法，将每个解码头分布到计算资源上，并利用GPU加速。该方法在开源的HydraGNN架构中实现，训练数据超过2400万个结构，并在Perlmutter、Aurora和Frontier超级计算机上进行了测试，展示了在这三种高度异构的超级计算架构上的高效扩展性。

🔬 方法详解

问题定义：本文旨在解决在多源、多保真原子建模数据的预训练过程中，现有方法在稳定性和迁移能力方面的不足。

核心思路：提出的多任务并行方法通过将每个解码头分布到不同的计算资源上，利用GPU加速来提高模型的训练效率和稳定性。这样的设计旨在应对大规模数据集的处理需求。

技术框架：整体架构包括共享的消息传递层和多个解码头，消息传递层负责处理输入的原子结构，而解码头则针对特定数据输出进行预测。该方法在开源的HydraGNN架构中实现，支持大规模并行计算。

关键创新：最重要的技术创新在于多任务并行方法的提出，使得每个解码头可以独立并行处理，从而显著提高了训练效率和模型的可扩展性。这与传统的集中式训练方法形成了鲜明对比。

关键设计：在模型设计中，采用了适应性损失函数和优化的网络结构，以确保在处理大规模数据时的稳定性和高效性。具体的参数设置和网络结构细节在实验部分进行了详细描述。

📊 实验亮点

实验结果显示，提出的方法在2400万个结构的训练中，成功在Perlmutter、Aurora和Frontier超级计算机上实现了高效的扩展性，显著提升了模型的训练速度和稳定性，具体性能数据尚未公开。

🎯 应用场景

该研究的潜在应用领域包括化学分子建模、材料科学和药物发现等。通过提高图基础模型在多源数据上的鲁棒性和迁移能力，能够加速新材料的发现和优化过程，具有重要的实际价值和长远影响。

📄 摘要（原文）

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.