Graphcore拟未帮助美国太平洋西北国家实验室利用GNN加速3D分子建模

公众号新闻

2023-02-27 07:02

Graphcore拟未和美国太平洋西北国家实验室（PNNL）^[1]之间的合作大幅加快了为计算化学应用训练计算机模型的速度。

合作结果表明，在使用拟未IPU作为AI加速器时，人们可以通过模型预训练和微调大幅缩减分子图神经网络（GNN）的训练时间。

该成果通过《减少下游时间：使用异构AI加速器预训练分子图神经网络》^[2]在线发表，并于近期在NeurIPS机器学习和物理科学研讨会^[3]上发布。

分子动力学轨迹，比较了图神经网络的预训练结果和微调结果。对比显示，微调后的模型能更好地再现地面实况方法，同时评估速度明显加快。

合作促进发现

PNNL是化学、数据分析、地球和生命科学领域的领先科学发现中心，也是可持续能源和国家安全的技术创新中心。

该合作将PNNL在化学领域中的巨大优势与拟未先进的AI系统融合起来，加速分子图神经网络训练。

截至目前，该合作的工作重点是将SchNet GNN架构^[4]应用于一个大型数据集——HydroNet^[5]数据集，该数据集由500万个水簇的几何结构组成。

这个图神经网络是为探索分子的几何结构与功能关系而设计的，其计算成本仅是传统计算化学方法的一小部分。

在这项研究中，使用HydroNet数据集是一个自然而然的选择：它不仅是迄今为止报告的最大水簇能量学数据集，而且其开发也是由PNNL的研究人员带头进行的。

用IPU加速计算化学

SchNet模型通过学习从分子结构到量子化学特性的映射来工作。PNNL和拟未的研究人员正在训练图神经网络预测水簇的结合能量。

这种能量可以简单地理解为对整个原子网络被化学键以怎样的强度固定在一起的一种量化衡量标准。

HydroNet数据集包括使用高度精确的量子化学方法计算的结合能量^[6]。

水簇内的化学键包括与附近原子的强共价键以及通过氢键的长程相互作用的混合物。

众所周知，这些相互作用很难进行性能建模。SchNet GNN很好地处理了这些问题，它使用了一个空间分界线来定义成对的相互作用，这种相互作用随着被建模的原子数量的增加而线性增长，并利用多个消息传递步骤来有效地在整个结构中传播成对的相互作用，以更好地纳入键合网络的多体性质。

将SchNet用于HydroNet数据集的初步调研只限于使用总数据集的10%^[7]。

即使限制了数据集的大小，使用4个NVIDIA V100 GPU，每epoch时间4.5分钟，仍需2.7天的训练时间。

下图显示了使用拟未在同样的HydroNet数据集小子集（10%）上加速训练SchNet所带来的每epoch时间的提升。

使用4个V100 NVIDIA GPU^[2]相对于基线实现的速度提升（越高越好）。在50万个水簇的小子集上，使用64个IPU可以实现超过60倍的加速。^[4]

结果表明，相比之前使用传统硬件加速器的基线实现，拟未Bow Pod₁₆实现了40倍以上的加速。

从这些结果来看，按照公开的Paperspace gradient机器价格^[8]，在4个V100 GPU上2.7天的训练时间将花费596美元。相比之下，在Bow Pod₁₆上，同样的训练工作量只需要1.4小时，花费37美元。

这些比较只是针对HydroNet数据集的小子集。在扩展到更大的数据集上训练，或是在不同的超参数上进行扫描时，这些成本和模型开发时间上的节约会进一步地扩大。

使用Paperspace云服务在HydroNet数据集上训练SchNet GNN的计算成本比较

与基线实现相比，在一系列Bow IPU系统上收集到的每epoch时间（越短越好）。

更重要的是，通过使用拟未Pod系统，我们与PNNL的合作已经证明了整个数据集上的高效训练。在《基于图神经网络的量子化学预测模型的极限加速》一文中，我们更详细地探讨了拟未IPU支持高效训练的特点和能力^[9]。

对于使用机器学习进一步了解原子尺度上复杂的多体相互作用，以及这些相互作用与宏观观测指标之间的关系来说，这项工作是前进的一小步，但这一小步却是非常重要的一步。在一系列应用中（从太阳能面板和电池到药物研发等众多领域），对化学键的理解发挥着至关重要作用。这些方法的使用可以帮助人们在这些领域中取得突破。

[1]https://www.pnnl.gov/

[2]Bilbrey, J. A. et al. Reducing Down(stream)time: Pretraining Molecular GNNs using Heterogeneous AI Accelerators. Preprint at https://arxiv.org/abs/2211.04598 (2022).

[3]https://ml4physicalsciences.github.io/2022/

[4]Schütt, K. T., Sauceda, H. E., Kindermans, P. J., Tkatchenko, A. & Müller, K. R. SchNet - A deep learning architecture for molecules and materials. J. Chem. Phys. 148, 1–11 (2018).

[5]Choudhury, S. et al. HydroNet: Benchmark Tasks for Preserving Intermolecular Interactions and Structural Motifs in Predictive and Generative Models for Molecular Data. Preprint at https://arxiv.org/abs/2012.00131 (2020).

[6]Rakshit, A., Bandyopadhyay, P., Heindel, J. P. & Xantheas, S. S. Atlas of putative minima and low-lying energy networks of water clusters n = 3-25. J. Chem. Phys. 151, (2019).

[7]Bilbrey, J. A. et al. A look inside the black box: Using graph-theoretical descriptors to interpret a Continuous-Filter Convolutional Neural Network (CF-CNN) trained on the global and local minimum energy structures of neutral water clusters. J. Chem. Phys. 153, (2020).

[8]https://docs.paperspace.com/gradient/machines/

[9]Helal, H. et al. Extreme Acceleration of Graph Neural Network-based Prediction Models for Quantum Chemistry. Preprint at https://arxiv.org/abs/2211.13853 (2022).