面向训练十亿参数图神经网络以进行原子模拟

近期在图神经网络（GNNs）用于原子模拟建模方面取得的进展有可能彻底改变催化剂的发现，这是实现应对气候变化所需的能源突破的关键步骤之一。然而，对于这一任务最为有效的GNN模型需要大量的内存资源，因为它们模拟了图中的高阶相互作用，例如原子三元组或四元组之间的相互作用，这使得这些模型的扩展变得非常困难。本文中，我们提出了一种称为图并行化的方法，该方法可以将输入图分布在多个GPU上，从而能够训练拥有数亿甚至数十亿参数的大型GNN模型。我们通过大幅增加最近提出的DimeNet++和GemNet模型的参数数量来对我们的方法进行实证评估。在大规模的Open Catalyst 2020（OC20）数据集上，这些图并行化的模型在S2EF任务的力均方误差（MAE）指标上相对提高了15%，在IS2RS任务的AFbT指标上相对提高了21%，从而确立了新的最先进结果。