arXiv:2504.10700v1 公告类型: cross 摘要: 利用图神经网络（GNNs）处理三维分子图结构的化学基础模型（CFMs）已成为计算化学家和材料科学家不可或缺的工具。这些模型有助于人们理解物质和发现新的分子和材料。与在大型同质图上操作的GNNs不同，CFMs使用的GNNs处理的是大量尺寸各异的几何图，这需要与为大型同质GNNs开发的不同优化策略。本文对CFM训练的两个关键阶段——数据分布和模型训练进行了优化，针对的是目前最先进的CFM——MACE。我们通过将其表述为多目标装载问题来应对数据分布中的负载均衡挑战。我们提出了一种迭代算法，提供了高效、快速且实用的解决方案，确保了有效的数据分布。对于训练阶段，我们确定了MACE中的核心计算内核是张量对称收缩，并优化了这个内核以提高整体性能。我们结合了均衡数据分布和内核优化的方法，显著提高了MACE的训练过程。实验结果表明，使用260万个样本数据集和740个GPU时，每轮训练的执行时间减少了10倍，从12分钟缩短至2分钟。