[论文解读] Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

Anuroop Sriram, Abhishek Das+3 more|arXiv (Cornell University)|Mar 18, 2022

Machine Learning in Materials Science被引用 9

一句话总结

本文提出图并行（Graph Parallelism），一种通过将输入图分布到多个GPU上，实现百亿参数图神经网络（GNN）在原子模拟中训练的方法，从而支持对更高阶相互作用（如原子三元组）的建模。该方法在Open Catalyst 2020（OC20）数据集上取得当前最优结果，相比先前模型，力的平均绝对误差（force MAE）降低15%，AFbT指标提升21%。

ABSTRACT

Recent progress in Graph Neural Networks (GNNs) for modeling atomic simulations has the potential to revolutionize catalyst discovery, which is a key step in making progress towards the energy breakthroughs needed to combat climate change. However, the GNNs that have proven most effective for this task are memory intensive as they model higher-order interactions in the graphs such as those between triplets or quadruplets of atoms, making it challenging to scale these models. In this paper, we introduce Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling us to train very large GNNs with hundreds of millions or billions of parameters. We empirically evaluate our method by scaling up the number of parameters of the recently proposed DimeNet++ and GemNet models by over an order of magnitude. On the large-scale Open Catalyst 2020 (OC20) dataset, these graph-parallelized models lead to relative improvements of 1) 15% on the force MAE metric for the S2EF task and 2) 21% on the AFbT metric for the IS2RS task, establishing new state-of-the-art results.

研究动机与目标

• 为解决在原子模拟中训练具有百亿参数的大型GNN的挑战，特别是针对建模更高阶相互作用（如原子三元组和四元组）的模型。
• 克服现有GNN训练方法在内存和计算上的局限性，这些方法不适用于包含数百万个中等规模图的数据集。
• 实现对大规模Open Catalyst 2020（OC20）数据集上最先进GNN模型（如DimeNet++和GemNet-T）的可扩展训练。
• 证明通过图并行扩展模型规模，可在预测原子受力和弛豫结构方面带来显著性能提升。
• 提供一种可推广的框架，适用于广泛的消息传递GNN架构，包括等变和能量守恒模型，且在扩展图网络（EGN）框架内有效。

提出的方法

• 提出图并行（Graph Parallelism），一种将输入图在多个GPU之间划分的技术，以支持无法在单个设备上运行的大型GNN训练。
• 将图网络（GN）框架扩展为扩展图网络（EGN）框架，以显式建模原子三元组和四元组等更高阶相互作用。
• 将更高阶相互作用（如键角）的计算拆分到多个GPU上，实现内存和计算的高效分布。
• 采用类似数据并行的策略，每个GPU处理输入图的一个子集，同时保持对更高阶项的完整图连通性。
• 与现有的模型并行技术（如GPipe）集成，实现对超大GNN的广度和深度并行扩展。
• 采用混合训练策略，结合图划分与模型并行，实现GNN中参数数量和相互作用复杂度的双重扩展。

实验结果

研究问题

• RQ1 尽管高阶相互作用带来高昂的内存和计算需求，是否仍能有效训练百亿参数的GNN用于原子模拟？
• RQ2 在包含数百万个小型图的分子数据集上，图并行与现有分布式GNN训练方法相比，在可扩展性和性能方面表现如何？
• RQ3 通过图并行训练的更大GNN在关键原子模拟基准（如力MAE和AFbT）上的预测精度提升程度如何？
• RQ4 图并行是否可推广至不同GNN架构，包括建模能量守恒和旋转等变性的模型？
• RQ5 训练此类大规模模型的环境影响如何？可通过战略性地选择计算资源来减轻吗？

主要发现

• 图并行使DimeNet++和GemNet-T模型的参数量超过1亿，相比之前模型提升一个数量级以上。
• 图并行化的GemNet-XL模型在S2EF任务上的力MAE相比之前SOTA模型降低15%。
• 同一模型在IS2RS任务的AFbT指标上实现21%的相对提升，成为OC20数据集上的新SOTA。
• 该方法具有良好的通用性，适用于广泛的消息传递GNN架构，包括等变和能量守恒模型，且在EGN框架内有效。
• 在配备Tesla V100 32GB的云GPU上训练GemNet-XL模型，估算的CO2排放量为3490–8052 kg CO2 eq.，相当于从洛杉矶飞往纽约的16次往返航班。
• 作者建议选择具有强碳抵消承诺的计算资源以减少环境影响，指出其训练过程已实现100%碳抵消。