论文标题
量子化学基于图神经网络的预测模型的极端加速度
Extreme Acceleration of Graph Neural Network-based Prediction Models for Quantum Chemistry
论文作者
论文摘要
分子特性计算是化学物理学的基岩。高保真\ textIt {ab intio}用于计算分子特性的建模技术可能非常昂贵,并激励机器学习模型的开发,从而使相同的预测更有效。大分子数据库上的训练图神经网络引入了独特的计算挑战,例如需要处理数百万个具有可变大小的小图和支持通信模式,这些图与在社交网络等大图上的学习不同。本文展示了一种新型的硬件软件共同设计方法,可扩展图形神经网络的训练以进行分子属性预测。我们引入了一种算法,将分子图的批次合并到固定尺寸包中,以消除与替代填充技术相关的冗余计算和内存,并通过最大程度地减少通信来改善吞吐量。我们通过在GraphCore Intelligence处理单元(IPU)上提供了公认的分子财产预测模型(IPU)来证明我们的共设计方法的有效性。我们以不同程度的图计数,尺寸和稀疏度在多个分子图数据库上评估训练性能。我们证明,这种共同设计方法可以将这种分子性质预测模型的训练时间从几天减少到不到两个小时,从而为AI驱动的科学发现开辟了新的可能性。
Molecular property calculations are the bedrock of chemical physics. High-fidelity \textit{ab initio} modeling techniques for computing the molecular properties can be prohibitively expensive, and motivate the development of machine-learning models that make the same predictions more efficiently. Training graph neural networks over large molecular databases introduces unique computational challenges such as the need to process millions of small graphs with variable size and support communication patterns that are distinct from learning over large graphs such as social networks. This paper demonstrates a novel hardware-software co-design approach to scale up the training of graph neural networks for molecular property prediction. We introduce an algorithm to coalesce the batches of molecular graphs into fixed size packs to eliminate redundant computation and memory associated with alternative padding techniques and improve throughput via minimizing communication. We demonstrate the effectiveness of our co-design approach by providing an implementation of a well-established molecular property prediction model on the Graphcore Intelligence Processing Units (IPU). We evaluate the training performance on multiple molecular graph databases with varying degrees of graph counts, sizes and sparsity. We demonstrate that such a co-design approach can reduce the training time of such molecular property prediction models from days to less than two hours, opening new possibilities for AI-driven scientific discovery.