量子化学中基于图神经网络预测模型的极致加速
机器学习
2022-11-28 v1 硬件体系结构
化学物理
摘要
分子性质计算是化学物理学的基石。用于计算分子性质的高保真 \textit{ab initio} 建模技术可能极其昂贵，这促使了能以更高效率作出相同预测的机器学习模型的发展。在大型分子数据库上训练图神经网络带来了独特的计算挑战，例如需要处理数百万个大小可变的小图，并支持不同于社交网络等大型图上学习的通信模式。本文展示了一种新颖的软硬件协同设计方法，以扩展分子性质预测图神经网络的训练规模。我们引入一种算法，将分子图批次合并为固定大小的包，以消除替代填充技术带来的冗余计算与内存，并通过最小化通信来提高吞吐量。我们通过在 Graphcore 智能处理单元（IPU）上实现一种成熟的分子性质预测模型，证明了协同设计方法的有效性。我们在具有不同图数量、大小与稀疏度的多个分子图数据库上评估了训练性能。我们表明，这种协同设计方法可将此类分子性质预测模型的训练时间从数天缩短至不到两小时，为 AI 驱动的科学发现开辟了新可能。
引用
@article{arxiv.2211.13853,
title = {Extreme Acceleration of Graph Neural Network-based Prediction Models for Quantum Chemistry},
author = {Hatem Helal and Jesun Firoz and Jenna Bilbrey and Mario Michael Krell and Tom Murray and Ang Li and Sotiris Xantheas and Sutanay Choudhury},
journal= {arXiv preprint arXiv:2211.13853},
year = {2022}
}