[论文解读] GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets

Johannes Gasteiger, Muhammed Shuaibi+5 more|arXiv (Cornell University)|Apr 6, 2022

Machine Learning in Materials Science被引用 54

一句话总结

GemNet-OC 在 OC20 上实现了最先进的结果，同时训练速度大约快10倍，并且论文分析数据集规模、多样性和领域偏移如何影响 GNN 设计，提出 OC-2M 作为高效开发的代表子集。

ABSTRACT

Recent years have seen the advent of molecular simulation datasets that are orders of magnitude larger and more diverse. These new datasets differ substantially in four aspects of complexity: 1. Chemical diversity (number of different elements), 2. system size (number of atoms per sample), 3. dataset size (number of data samples), and 4. domain shift (similarity of the training and test set). Despite these large differences, benchmarks on small and narrow datasets remain the predominant method of demonstrating progress in graph neural networks (GNNs) for molecular simulation, likely due to cheaper training compute requirements. This raises the question -- does GNN progress on small and narrow datasets translate to these more complex datasets? This work investigates this question by first developing the GemNet-OC model based on the large Open Catalyst 2020 (OC20) dataset. GemNet-OC outperforms the previous state-of-the-art on OC20 by 16% while reducing training time by a factor of 10. We then compare the impact of 18 model components and hyperparameter choices on performance in multiple datasets. We find that the resulting model would be drastically different depending on the dataset used for making model choices. To isolate the source of this discrepancy we study six subsets of the OC20 dataset that individually test each of the above-mentioned four dataset aspects. We find that results on the OC-2M subset correlate well with the full OC20 dataset while being substantially cheaper to train on. Our findings challenge the common practice of developing GNNs solely on small datasets, but highlight ways of achieving fast development cycles and generalizable results via moderately-sized, representative datasets such as OC-2M and efficient models such as GemNet-OC. Our code and pretrained model weights are open-sourced.

研究动机与目标

• 研究 GNN 在小数据集上的改进是否会迁移到大规模、多样化的分子数据

提出的方法

• 基于 OC20 构建的 GemNet-OC 基线，具备多层次交互层级以及边/原子嵌入
• 用固定的最近邻数量替代固定距离截断来进行图构建
• 简化并优化基函数以降低径向/角向项的计算成本
• 引入包含原子-原子、边-原子、边-边交互的分层交互方案，并提供长程原子级路径
• 输出并连接来自交互块的嵌入，以提高最终的能量/力预测
• 提供开源代码和预训练权重，便于在 OC20 和 OC-2M 上实现可重复开发

实验结果

研究问题

• RQ1 四个数据集属性（化学多样性、体系大小、数据集规模和领域偏移）如何影响 GNN 设计决策？
• RQ2 在 OC20 上训练得到的模型能否在降低训练成本的情况下达到最先进的性能？
• RQ3 OC-2M 是否是一个与 OC20 趋势相关、便于更快开发的可靠子集？
• RQ4 模型组件的效应是否会随着小数据集与大数据集、以及 OC20 的子集的不同而变化？
• RQ5 哪些架构与训练策略最能实现对大型分子数据集的快速、可扩展学习？

主要发现

• GemNet-OC 在 OC20 任务上实现了最先进的结果，并且比先前的大模型训练速度约快10倍
• GemNet-OC 在 OC-2M 和 OC20 的训练数据上优于以前的模型，同时使用明显更少的训练数据
• 在小数据集与大、不同数据集上，模型组件的表现存在显著差异
• OC-2M 与完整 OC20 的相关性良好，使快速且具有代表性的模型开发成为可能
• 带有长程原子级路径的交互层级和优化的基函数实现了对大系统的高效处理